From 46bc9c866f85354570baec691b6a858c0c46d7aa Mon Sep 17 00:00:00 2001
From: Doug Coleman <doug.coleman@gmail.com>
Date: Tue, 9 Aug 2022 22:03:02 -0500
Subject: [PATCH] assocs.extras: refactor collect-by words

---
 extra/assocs/extras/extras-tests.factor | 35 ++++++++++++++++++-------
 extra/assocs/extras/extras.factor       | 22 +++++++++-------
 2 files changed, 37 insertions(+), 20 deletions(-)

diff --git a/extra/assocs/extras/extras-tests.factor b/extra/assocs/extras/extras-tests.factor
index 25ded3f642..65aaf7798c 100644
--- a/extra/assocs/extras/extras-tests.factor
+++ b/extra/assocs/extras/extras-tests.factor
@@ -201,21 +201,36 @@ USING: arrays assocs.extras kernel math math.order sequences tools.test ;
     [ [ drop even? ] [ 2array ] 2bi ] collect-assoc-by
 ] unit-test
 
+
 {
-    H{ { t V{ 10 20 30 } } { f V{ 41 } } }
+    H{ { t V{ 10 21 } } { f V{ 30 41 } } }
 } [
-    { { 10 100 } { 20 200 } { 30 300 } { 41 401 } }
-    [ even? ] collect-key-by
-] unit-test
+    { { 10 100 } { 21 200 } { 30 301 } { 41 401 } }
+    [ nip even? ] collect-key-by
+ ] unit-test
 
 {
-    H{ { t V{ 100 200 300 } } { f V{ 401 } } }
+    H{ { t V{ 10 30 } } { f V{ 21 41 } } }
 } [
-    { { 10 100 } { 20 200 } { 30 300 } { 41 401 } }
-    [ even? ] collect-value-by
-] unit-test
+    { { 10 100 } { 21 200 } { 30 301 } { 41 401 } }
+    [ drop even? ] collect-key-by
+ ] unit-test
 
 
+{
+    H{ { t V{ 100 200 } } { f V{ 301 401 } } }
+} [
+    { { 10 100 } { 21 200 } { 30 301 } { 41 401 } }
+    [ nip even? ] collect-value-by
+ ] unit-test
+
+{
+    H{ { t V{ 100 301 } } { f V{ 200 401 } } }
+} [
+    { { 10 100 } { 21 200 } { 30 301 } { 41 401 } }
+    [ drop even? ] collect-value-by
+ ] unit-test
+
 {
     H{
         { 1 V{ 10 20 30 40 50 60 } }
@@ -244,7 +259,7 @@ USING: arrays assocs.extras kernel math math.order sequences tools.test ;
     }
 } [
     { { 10 100 } { 20 200 } { 30 300 } { 41 401 } }
-    [ dup 1 + 2array ] collect-key-by-multi
+    [ drop dup 1 + 2array ] collect-key-by-multi
 ] unit-test
 
 
@@ -261,7 +276,7 @@ USING: arrays assocs.extras kernel math math.order sequences tools.test ;
     }
 } [
     { { 10 100 } { 20 200 } { 30 300 } { 41 401 } }
-    [ dup 1 + 2array ] collect-value-by-multi
+    [ nip dup 1 + 2array ] collect-value-by-multi
 ] unit-test
 
 
diff --git a/extra/assocs/extras/extras.factor b/extra/assocs/extras/extras.factor
index 533e5ffa63..1314f8b616 100644
--- a/extra/assocs/extras/extras.factor
+++ b/extra/assocs/extras/extras.factor
@@ -245,16 +245,16 @@ PRIVATE>
 : collect-assoc-by ( ... input-assoc quot: ( ... key value -- ... key value ) -- ... assoc )
     [ H{ } clone ] 2dip collect-assoc-by! ; inline
 
-: collect-key-by! ( ... assoc input-assoc quot: ( ... key -- ... new-key ) -- ... assoc )
-    rot [ '[ drop _ keep swap _ push-at ] assoc-each ] keep ; inline
+: collect-key-by! ( ... assoc input-assoc quot: ( ... key value -- ... new-key ) -- ... assoc )
+    '[ _ keepd ] collect-assoc-by! ; inline
 
-: collect-key-by ( ... input-assoc quot: ( ... key -- ... new-key ) -- ... assoc )
+: collect-key-by ( ... input-assoc quot: ( ... key value -- ... new-key ) -- ... assoc )
     [ H{ } clone ] 2dip collect-key-by! ; inline
 
-: collect-value-by! ( ... assoc input-assoc quot: ( ... value -- ... new-key ) -- ... assoc )
-    rot [ '[ nip _ keep swap _ push-at ] assoc-each ] keep ; inline
+: collect-value-by! ( ... assoc input-assoc quot: ( ... key value -- ... new-key ) -- ... assoc )
+    '[ _ keep ] collect-assoc-by! ; inline
 
-: collect-value-by ( ... input-assoc quot: ( ... value -- ... new-key ) -- ... assoc )
+: collect-value-by ( ... input-assoc quot: ( ... key value -- ... new-key ) -- ... assoc )
     [ H{ } clone ] 2dip collect-value-by! ; inline
 
 
@@ -264,14 +264,16 @@ PRIVATE>
 : collect-assoc-by-multi ( ... assoc quot: ( ... key value -- ... new-keys value' ) -- ... assoc )
     [ H{ } clone ] 2dip collect-assoc-by-multi! ; inline
 
-: collect-key-by-multi! ( ... assoc input-assoc quot: ( ... key -- ... new-keys ) -- ... assoc )
-    rot [ '[ drop _ keep swap _ push-at-each ] assoc-each ] keep ; inline
+
+: collect-key-by-multi! ( ... assoc input-assoc quot: ( ... key value -- ... new-keys ) -- ... assoc )
+    '[ _ keepd ] collect-assoc-by-multi! ; inline
 
 : collect-key-by-multi ( ... assoc quot: ( ... key -- ... new-keys ) -- ... assoc )
     [ H{ } clone ] 2dip collect-key-by-multi! ; inline
 
-: collect-value-by-multi! ( ... assoc input-assoc quot: ( ... value -- ... new-keys ) -- ... assoc )
-    rot [ '[ nip _ keep swap _ push-at-each ] assoc-each ] keep ; inline
+
+: collect-value-by-multi! ( ... assoc input-assoc quot: ( ... key value -- ... new-keys ) -- ... assoc )
+    '[ _ keep ] collect-assoc-by-multi! ; inline
 
 : collect-value-by-multi ( ... assoc quot: ( ... value -- ... new-keys ) -- ... assoc )
     [ H{ } clone ] 2dip collect-value-by-multi! ; inline
-- 
2.34.1