trees, implement assoc-size

diff --git a/extra/trees/trees-tests.factor b/extra/trees/trees-tests.factor
index 82cf3c812a5955e1bded281abe0bb52c51e347d3..2172f9add03c0aa11b2c155f61ee5071ee843bb5 100644 (file)
--- a/extra/trees/trees-tests.factor
+++ b/extra/trees/trees-tests.factor
@@ -51,3 +51,7 @@ IN: trees.tests
     { 7 "nine" }
     { 4 "four" }
 } height ] unit-test
+
+! test assoc-size
+{ 3 } [ test-tree assoc-size ] unit-test
+{ 2 } [ test-tree 9 over delete-at assoc-size ] unit-test

diff --git a/extra/trees/trees.factor b/extra/trees/trees.factor
index b64168a573b55d9409bed325806714c89bcce2a0..8b3b41403322fd575e56fd1aff1e290cec5f1caf 100644 (file)
--- a/extra/trees/trees.factor
+++ b/extra/trees/trees.factor
@@ -96,18 +96,19 @@ CONSTANT: right 1
 M: tree at*
     root>> node-at* ;
 
-: node-set ( value key node -- node )
+: node-set ( value key node -- node new? )
     2dup key>> key-side dup 0 eq? [
-        drop nip swap >>value
+        drop nip swap >>value f
     ] [
         [
-            [ node-link [ node-set ] [ swap <node> ] if* ] keep
-            [ set-node-link ] keep
+            [ node-link [ node-set ] [ swap <node> t ] if* ] keep
+            swap [ [ set-node-link ] keep ] dip
         ] with-side
     ] if ;
 
 M: tree set-at
-    [ [ node-set ] [ swap <node> ] if* ] change-root drop ;
+    [ [ node-set ] [ swap <node> t ] if* swap ] change-root
+    swap [ dup inc-count ] when drop ;
 
 : valid-node? ( node -- ? )
     [
@@ -183,21 +184,22 @@ DEFER: delete-node
         nip ! right but no left, or no children
     ] if* ;
 
-: delete-bst-node ( key node -- node )
+: delete-bst-node ( key node -- node deleted? )
     2dup key>> key-side dup 0 eq? [
-        drop nip delete-node
+        drop nip delete-node t
     ] [
         [
             [ node-link delete-bst-node ]
-            [ set-node-link ]
-            [ ] tri
+            [ swap [ set-node-link ] dip ]
+            [ swap ] tri
         ] with-side
     ] if ;
 
 PRIVATE>
 
 M: tree delete-at
-    [ delete-bst-node ] change-root drop ;
+    [ delete-bst-node swap ] change-root
+    swap [ dup dec-count ] when drop ;
 
 M: tree new-assoc
     2drop <tree> ;