From: Doug Coleman <doug.coleman@gmail.com>
Date: Mon, 8 Aug 2022 03:57:06 +0000 (-0500)
Subject: assocs.extras: Add collect-by-multi and push-at-each
X-Git-Tag: 0.99~1168
X-Git-Url: https://gitweb.factorcode.org/gitweb.cgi?p=factor.git;a=commitdiff_plain;h=af61f2cee40c953562b131f0800c877c3cc8aa08

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 a4d7a08efd..e3211882c9 100644
--- 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 69a2c3cbce..3425b13388 100644
--- 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