Implement insert

trees/cb/cb.factor

@@ -14,8 +14,8 @@
 ! 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 fry kernel math namespaces sequences serialize
-trees trees.private ;
+USING: accessors arrays assocs combinators.short-circuit fry kernel locals math
+namespaces sequences serialize trees trees.private ;
 IN: trees.cb
 
 TUPLE: cb < tree ;
@@ -75,8 +75,8 @@ SYMBOL: new-side
 
 : with-key ( key quot -- )
     [
-        { current-key key-bytes new-side } swap
-        dup key>bytes left 3array zip
+        { current-key key-bytes } swap
+        dup key>bytes 2array zip
     ] dip with-variables ; inline
 
 ! Extract the critical byte
@@ -93,42 +93,9 @@ SYMBOL: new-side
 ! Tree insertion must be done by traversing the tree from the root, as it is
 ! ordered.
 
-! ! Insert a node at a leaf-node by first inserting a split node.
-! : cb-insert ( value leaf-node -- node created? )
-!     [ current-key get swap <node> ] dip
-!     dup key>> key>bytes key-bytes get swap
-!     bytes-diff swap <cb-node>
-!     ! ( new-leaf old-leaf side cb-node -- )
-!     [ swap
-!       [ [ set-node-link ] keep
-!         set-node+link
-!       ] with-side
-!     ] keep t ;
-
-! DEFER: cb-set
-
-! ! Insert a node into a non-empty tree.
-! : (cb-set) ( value parent node -- parent node created? )
-!     dup cb-node? [
-!         select-side [
-!             [ node-link (cb-set) ] keep
-!             swap [ [ set-node-link ] keep ] dip
-!         ] with-side
-!     ] [
-!         cb-insert
-!     ] if ;
-
-! ! Insert a node into the tree
-! : cb-set ( value node -- node created? )
-!     [ dup (cb-set) ] [ current-key get swap <node> t ] if* ;
-
 ! ** Walking the tree for the best fit
 
-! We need to remember the parent node as it will be the split node where our new
-! value starts to differ, should it not be (key-)identical with the current
-! node. Start with =f= as the parent.
-
-GENERIC: cb-best-fit ( parent/f node -- node child )
+GENERIC: cb-best-fit ( node -- node )
 
 M: f cb-best-fit ;
 
@@ -136,61 +103,83 @@ M: node cb-best-fit ;
 
 M: cb-node cb-best-fit
     select-side [
-        nip dup node-link cb-best-fit
+        node-link cb-best-fit
     ] with-side ;
 
-GENERIC: cb-update ( value parent/f node -- node created? )
+GENERIC: cb-update ( value node -- node created? )
 
 M: f cb-update
-    2drop current-key get swap <node> t ;
+    drop current-key get swap <node> t ;
 
+! Attach a new leaf node and record =new-side=. New leaf node is attached
+! opposite of =new-side=.
 : attach-node ( value side cb-node -- cb-node )
-    swap [
+    swap [ new-side set ] keep
+    [
         [ current-key get swap <node> ] dip
         [ set-node+link ] keep
     ] with-side ;
 
+! Update the tree by either updating a leaf node with a new key object and value
+! or create a new split node and attach a fresh leaf node with the new key and
+! value.
 M: node cb-update
     dup key>> current-key get = [
         current-key get >>key
-        nip swap >>value f
+        swap >>value f
     ] [
-        dup key>> key>bytes key-bytes get swap
+        key>> key>bytes key-bytes get swap
         bytes-diff swap <cb-node>
         attach-node t
     ] if ;
 
-GENERIC: cb-insert ( new-node cur-node -- node )
+! Break off the search when:
+! - the top node is no longer a split node
+! - the top split node is larger than the new split node (i.e. refers to a later
+!   byte or more significant bit in the current byte)
+: break? ( new-node node -- ? )
+    { [ nip cb-node? not ]
+      [ [ byte#>> ] bi@ < ]
+      [ { [ [ byte#>> ] bi@ = ]
+          [ [ bits>> ] bi@ < ]
+        } 2&& ]
+    } 2|| ;
 
-M: f cb-insert drop ;
+! Walk the tree and insert =new-node= at the pre-determined best place. We have
+! to keep track of the parent instead of the current node, as we might need to
+! relink from the parent to =new-node=.
+GENERIC: cb-insert ( new-node parent -- )
 
-M: node cb-insert
-    new-side get [
-        swap [ set-node-link ] keep
+M:: cb cb-insert ( n p -- )
+    p root>> :> c
+    n c break?
+    [ n p root<<
+      new-side get [
+          c n set-node-link
+      ] with-side
+    ]
+    [ n c cb-insert ] if ;
+
+M:: cb-node cb-insert ( n p -- )
+    p select-side [
+        drop
+        p node-link :> c
+        n c break?
+        [ n p set-node-link
+          new-side get [
+              c n set-node-link
+          ] with-side
+        ]
+        [ n c cb-insert ] if
     ] with-side ;
 
-M: cb-node cb-insert
-    2dup
-    [ [ [ byte#>> ] bi@ < ]
-      [ [ [ byte#>> ] bi@ = ]
-        [ [ bits>> ] bi@ <]
-      ] 2&&
-    ] 2|| [
-        ! break
-        select-side [ set-node-link ] with-side
-    ] [
-        ! step down
-    ] if
-
-
 PRIVATE>
 
 M: cb set-at ( value key cb -- )
-    [ swap
-      [ ! ( value root -- new created? root )
-          [ f swap cb-best-fit cb-update ] keep swap
-          [ cb-insert t ] [ 2nip f ] if
-      ] with-key swap ] change-root
+    [ swap [
+        [ root>> cb-best-fit cb-update ]
+        [ swap [ cb-insert t ] [ 2drop f ] if ] bi
+    ] with-key ] keep
     swap [ dup inc-count ] when drop ;
 
 ! ** Deletion