assocs.extras: go nuts on the crazy assoc words.

extra/assocs/extras/extras-tests.factor

@@ -1,5 +1,4 @@
-
-USING: assocs.extras kernel sequences tools.test ;
+USING: assocs.extras kernel math sequences tools.test ;
 
 { f } [ f { } deep-at ] unit-test
 { f } [ f { "foo" } deep-at ] unit-test
@@ -18,3 +17,41 @@ USING: assocs.extras kernel sequences tools.test ;
 
 { H{ } } [ H{ { 1 2 } } 2 over delete-value-at ] unit-test
 { H{ { 1 2 } } } [ H{ { 1 2 } } 3 over delete-value-at ] unit-test
+
+{
+    H{ { 1 3 } { 2 3 } }
+} [
+    {
+        { { 1 2 } 3 }
+    } expand-keys-set-at
+] unit-test
+
+{
+    H{ { 3 4 } }
+} [
+    {
+        { 3 { 1 2 } } { 3 4 }
+    } expand-values-set-at
+] unit-test
+
+{
+    H{ { 1 V{ 3 } } { 2 V{ 3 } } }
+} [
+    {
+        { { 1 2 } 3 }
+    } expand-keys-push-at
+] unit-test
+
+{
+    H{ { 3 V{ 1 2 4 } } }
+} [
+    {
+        { 3 { 1 2 } } { 3 4 }
+    } expand-values-push-at
+] unit-test
+
+{
+    H{ { 1 [ sq ] } { 2 [ sq ] } }
+} [
+    { { { 1 2 { 1 } { 2 } { 1 1 } } [ sq ] } } flatten-keys
+] unit-test

extra/assocs/extras/extras.factor

@@ -1,7 +1,7 @@
 ! Copyright (C) 2012 John Benediktsson, Doug Coleman
 ! See http://factorcode.org/license.txt for BSD license
-USING: arrays assocs assocs.private generalizations kernel math
-sequences ;
+USING: arrays assocs assocs.private fry generalizations kernel
+math sequences ;
 IN: assocs.extras
 
 : deep-at ( assoc seq -- value/f )
@@ -65,6 +65,12 @@ PRIVATE>
 : sequence>assoc! ( assoc seq map-quot: ( x -- ..y ) insert-quot: ( ..y assoc -- ) -- assoc )
     4 nrot (sequence>assoc) ; inline
 
+: assoc>object ( assoc map-quot insert-quot exemplar -- object )
+    clone [ swap curry compose assoc-each ] keep ; inline
+
+: assoc>object! ( assoc seq map-quot: ( x -- ..y ) insert-quot: ( ..y assoc -- ) -- object )
+    4 nrot assoc>object ; inline
+
 : sequence>assoc ( seq map-quot insert-quot exemplar -- assoc )
     clone (sequence>assoc) ; inline
 
@@ -77,3 +83,80 @@ PRIVATE>
 : sequence>hashtable ( seq map-quot insert-quot -- hashtable )
     H{ } sequence>assoc ; inline
 
+: expand-keys-set-at-as ( assoc exemplar -- hashtable' )
+    [
+        [ swap dup sequence? [ 1array ] unless ]
+        [ '[ _ set-at ] with each ]
+    ] dip assoc>object ;
+
+: expand-keys-set-at ( assoc -- hashtable' )
+    H{ } expand-keys-set-at-as ;
+
+: expand-keys-push-at-as ( assoc exemplar -- hashtable' )
+    [
+        [ swap dup sequence? [ 1array ] unless ]
+        [ '[ _ push-at ] with each ]
+    ] dip assoc>object ;
+
+: expand-keys-push-at ( assoc -- hashtable' )
+    H{ } expand-keys-push-at-as ; inline
+
+: expand-keys-push-as ( assoc exemplar -- hashtable' )
+    [
+        [ [ dup sequence? [ 1array ] unless ] dip ]
+        [ '[ 2array _ push ] curry each ]
+    ] dip assoc>object ;
+
+: expand-keys-push ( assoc -- hashtable' )
+    V{ } expand-keys-push-as ; inline
+
+: expand-values-set-at-as ( assoc exemplar -- hashtable' )
+    [
+        [ dup sequence? [ 1array ] unless swap ]
+        [ '[ _ set-at ] curry each ]
+    ] dip assoc>object ;
+
+: expand-values-set-at ( assoc -- hashtable' )
+    H{ } expand-values-set-at-as ; inline
+
+: expand-values-push-at-as ( assoc exemplar -- hashtable' )
+    [
+        [ dup sequence? [ 1array ] unless swap ]
+        [ '[ _ push-at ] curry each ]
+    ] dip assoc>object ;
+
+: expand-values-push-at ( assoc -- assoc )
+    H{ } expand-values-push-at-as ; inline
+
+: expand-values-push-as ( assoc exemplar -- assoc )
+    [
+        [ dup sequence? [ 1array ] unless ]
+        [ '[ 2array _ push ] with each ]
+    ] dip assoc>object ;
+
+: expand-values-push ( assoc -- sequence )
+    V{ } expand-values-push-as ; inline
+
+: assoc-any-key? ( ... assoc quot: ( ... key -- ... ? ) -- ... ? )
+    [ drop ] prepose assoc-find 2nip ; inline
+
+: assoc-any-value? ( ... assoc quot: ( ... key -- ... ? ) -- ... ? )
+    [ nip ] prepose assoc-find 2nip ; inline
+
+: assoc-all-key? ( ... assoc quot: ( ... key -- ... ? ) -- ... ? )
+    [ not ] compose assoc-any-key? not  ; inline
+
+: assoc-all-value? ( ... assoc quot: ( ... key -- ... ? ) -- ... ? )
+    [ not ] compose assoc-any-value? not  ; inline
+
+: any-multi-key? ( assoc -- ? )
+    [ sequence? ] assoc-any-key? ;
+
+: any-multi-value? ( assoc -- ? )
+    [ sequence? ] assoc-any-value? ;
+
+: flatten-keys ( assoc -- assoc' )
+    dup any-multi-key? [ expand-keys-set-at flatten-keys ] when ;
+
+: flatten-values ( assoc -- assoc' )
+    dup any-multi-value? [ expand-values-set-at flatten-values ] when ;