assocs.extras: Add collect-by-multi and push-at-each

diff --git a/extra/assocs/extras/extras-tests.factor b/extra/assocs/extras/extras-tests.factor
index a4d7a08efd8edb24e0de631477f2667993bea9d8..e3211882c9c26eed8c96ab317f95c8d865011772 100644 (file)
--- a/extra/assocs/extras/extras-tests.factor
+++ b/extra/assocs/extras/extras-tests.factor
@@ -164,3 +164,22 @@ USING: arrays assocs.extras kernel math math.order sequences tools.test ;
     H{ { 3 30 } { 4 40 } } 3array
     [ min ] V{ } assoc-collapse-as
 ] unit-test
+
+{
+    H{ { 1 V{ 10 } } { 2 V{ 10 } } { 3 V{ 10 } } { 4 V{ 10 } } { 5 V{ 10 } } }
+} [
+    H{ } clone 10 { 1 2 3 4 5 } pick push-at-each
+] unit-test
+
+{
+    H{
+        { 1 V{ 10 20 30 40 50 60 } }
+        { 2 V{ 10 20 30 40 50 60 } }
+        { 3 V{ 10 20 30 40 50 60 } }
+        { 4 V{ 10 20 30 40 50 60 } }
+        { 5 V{ 10 20 30 40 50 60 } }
+    }
+} [
+    { 10 20 30 } [ drop { 1 2 3 4 5 } ] collect-by-multi
+    { 40 50 60 } [ drop { 1 2 3 4 5 } ] collect-by-multi!
+] unit-test
\ No newline at end of file

diff --git a/extra/assocs/extras/extras.factor b/extra/assocs/extras/extras.factor
index 69a2c3cbce77811544fadc58ee916446fd22f11a..3425b13388831313ccdf170065096469f92594a8 100644 (file)
--- a/extra/assocs/extras/extras.factor
+++ b/extra/assocs/extras/extras.factor
@@ -234,3 +234,14 @@ PRIVATE>
 : histogram-diff ( hashtable1 hashtable2 -- hashtable3 )
     [ neg swap pick at+ ] assoc-each
     [ 0 > ] filter-values ;
+
+: push-at-each ( value keys assoc -- )
+    '[ _ push-at ] with each ; inline
+
+: collect-by-multi! ( ... assoc seq quot: ( ... obj -- ... key ) -- ... assoc )
+    [ keep swap ] curry rot [
+        [ push-at-each ] curry compose each
+    ] keep ; inline
+
+: collect-by-multi ( ... seq quot: ( ... obj -- ... keys ) -- ... assoc )
+    [ H{ } clone ] 2dip collect-by-multi! ; inline