Implement set-at; cb-set for the empty tree

cb/cb.factor

@@ -7,14 +7,15 @@
 ! PATRICIA trees showing that fast insertion, deletion, exact searching and suffix
 ! searching is possible with this data structure.
 
-! The strength of this datastructure, according to its author, lies in its simple
-! design and its optimisation to be machine parsable using machine word-sized
-! operations where possible. Like PATRICIA trees, crit-bit trees are
-! prefix-compressed, with internal nodes storing next decision point (the critical
-! bit) in a length field (encoded as an integer and a mask) and two successor
-! pointers. Arbitrary data objects make up its leaves.
+! The strength of this data structure, according to its author, lies in its
+! simple design and its optimisation to be machine parsable using machine
+! word-sized operations where possible. Like PATRICIA trees, crit-bit trees are
+! prefix-compressed, with internal nodes storing next decision point (the
+! 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 kernel math sequences serialize trees trees.private ;
+USING: accessors arrays assocs kernel math namespaces sequences serialize trees
+trees.private ;
 IN: trees.cb;
 
 TUPLE: cb < tree ;
@@ -23,49 +24,66 @@ TUPLE: cb < tree ;
 
 <PRIVATE
 
-TUPLE: cb-node < node byte# bits ;
+TUPLE: cb-node < node { byte# integer } { bits integer } ;
 
 : <cb-node> ( byte# bits key value -- node )
     cb-node new-node
-    >>bits
-    >>byte# ; inline
+    swap >>bits
+    swap >>byte# ; inline
 
 ! 0 = left
 ! 1 = right
-: direction ( byte bits -- direction )
+: direction ( n n -- direction )
     bitor 1 + -8 shift ; inline
 
-: key-side ( byte bits -- side )
+: key-side ( bits byte -- side )
     direction 0 = -1 1 ? ;
 
 : get-byte-at ( byte# key -- byte/0 )
     object>bytes ?nth [ 0 ] unless* ;
 
-! ** Insertion
-! Explain...
-
 ! Keep the byte sequence of the current key in =key-bytes= and provide a working
 ! environment for it with =with-key=.
 SYMBOL: key-bytes
+SYMBOL: current-key
 
 : with-key ( key quot -- )
-    [ key-bytes ] dip with-variable ; inline
+    [
+        { current-key key-bytes }
+        [ dup object>bytes 2array ] dip zip
+    ] dip with-variables ; inline
+
+: byte-at ( byte# -- byte/0 )
+    key-bytes get ?nth [ 0 ] unless* ;
+
+! ** Insertion
+! Explain...
 
 DEFER: cb-set
 
-! TODO Rewrite with =with-key=.
 : cb-insert ( value key node -- node taller? created? )
-    2dup [ byte#>> ] [ bits>> ] bi
-    [ swap get-byte-at ] dip key-side [
-        node-link cb-set
-    ] with-side ;
+    f
+    swap [
+        dup [ bits>> ] [ byte#>> ] bi
+        byte-at key-side [
+            node-link cb-set
+        ] with-side
+    ] with-key ;
 
 ! Implement this and cb-insert using SYMBOL for current key bytes while
 ! traversing the tree.
 : (cb-set) ( value key node -- node taller? created? )
-    ;
+    drop
+    dup key>> current-key get = [
+        current-key get >>key swap >>value f f
+    ] [ current-key get swap cb-insert ] if ;
 
+! Insert a node into the tree
 : cb-set ( value key node -- node taller? created? )
-    [ (cb-set) ] [ swap [ [ 0 get-byte-at 0xfe ] keep ] dip <cb-node> t t ] if* ;
+    [ (cb-set) ] [ [ 0 byte-at 0xfe ] 2dip swap <cb-node> t t ] if* ;
 
 PRIVATE>
+
+M: cb set-at ( value key cb -- )
+    [ cb-set nip swap ] change-root
+    swap [ dup inc-count ] when drop ;