Implement Assoc Protocol and pretty printing

trees/cb/cb.factor

@@ -14,8 +14,9 @@
 ! critical bit) in a length field (encoded as an integer and a mask) and two
 ! successor pointers. Arbitrary data objects make up its leaves.
 
-USING: accessors arrays assocs combinators.short-circuit fry kernel locals math
-namespaces sequences serialize trees trees.private ;
+USING: accessors arrays assocs combinators.short-circuit fry kernel locals make
+math namespaces parser prettyprint.custom sequences serialize trees
+trees.private vectors ;
 IN: trees.cb
 
 TUPLE: cb < tree ;
@@ -75,8 +76,8 @@ SYMBOL: new-side
 
 : with-key ( key quot -- )
     [
-        { current-key key-bytes } swap
-        dup key>bytes 2array zip
+        { current-key key-bytes new-side } swap
+        dup key>bytes left 3array zip
     ] dip with-variables ; inline
 
 ! Extract the critical byte
@@ -173,8 +174,6 @@ M:: cb-node cb-insert ( n p -- )
         [ n c cb-insert ] if
     ] with-side ;
 
-PRIVATE>
-
 M: cb set-at ( value key cb -- )
     [ swap [
         [ root>> cb-best-fit cb-update ]
@@ -184,8 +183,6 @@ M: cb set-at ( value key cb -- )
 
 ! ** Deletion
 
-<PRIVATE
-
 GENERIC: cb-delete ( node -- node deleted? )
 
 M: f cb-delete f ;
@@ -215,8 +212,62 @@ M: cb-node cb-delete ( node -- node deleted? )
         ] when drop
     ] with-side ;
 
-PRIVATE>
-
 M: cb delete-at ( key cb -- )
     [ swap [ cb-delete ] with-key swap ] change-root
     swap [ dup dec-count ] when drop ;
+
+M: cb new-assoc
+    2drop <cb> ;
+
+GENERIC: (cb-node>alist) ( node -- )
+
+M: node (cb-node>alist)
+    entry, ;
+
+: cb-node>entry ( node -- entry ) [ byte#>> ] [ bits>> ] bi 2array ;
+
+: cb-entry, ( node -- ) cb-node>entry , ;
+
+M: cb-node (cb-node>alist)
+    [ left>> (cb-node>alist) ]
+    [ cb-entry, ]
+    [ right>> (cb-node>alist) ]
+    tri ;
+
+M: cb >alist
+    [ root>> (cb-node>alist) ] { } make ;
+
+! Post-order traversal
+!
+! Assumes =f= in =left= and =right= slots of leaf nodes and ≠ =f= in split
+! nodes.
+: each-leaf-node ( node quot: ( ... entry -- ... ) -- ... )
+    [ [ dup left>> ] dip over [ each-leaf-node drop ] [ nip [ node>entry ] dip call ] if ]
+    [ [ right>> ] dip over [ each-leaf-node ] [ 2drop ] if ]
+    2bi ; inline recursive
+
+: >cb-alist ( tree -- alist )
+    dup assoc-size <vector> [
+        [ push ] curry [ root>> ] dip each-leaf-node
+    ] keep ;
+
+M: cb assoc-clone-like
+    [ dup cb? [ >cb-alist ] when ] dip call-next-method ;
+
+PRIVATE>
+
+: >cb ( assoc -- tree )
+    <cb> assoc-clone-like ;
+
+SYNTAX: CBTREE{
+        \ } [ >cb ] parse-literal ;
+
+<PRIVATE
+
+M: cb assoc-like drop dup cb? [ >cb ] unless ;
+
+M: cb pprint-delims drop \ CBTREE{ \ } ;
+M: cb >pprint-sequence >alist ;
+M: cb pprint-narrow? drop t ;
+
+PRIVATE>