trees.avl, implement assoc-size

extra/trees/avl/avl-tests.factor

@@ -119,3 +119,7 @@ IN: trees.avl.tests
 { f } [ test-tree 9 over delete-at 9 of ] unit-test
 { "replaced seven" } [ test-tree 9 over delete-at 7 of ] unit-test
 { "nine" } [ test-tree 7 over delete-at 4 over delete-at 9 of ] unit-test
+
+! test assoc-size
+{ 3 } [ test-tree assoc-size ] unit-test
+{ 2 } [ test-tree 9 over delete-at assoc-size ] unit-test

extra/trees/avl/avl.factor

@@ -63,24 +63,25 @@ TUPLE: avl-node < node balance ;
 
 DEFER: avl-set
 
-: avl-insert ( value key node -- node taller? )
+: avl-insert ( value key node -- node taller? created? )
     2dup key>> before? left right ? [
-        [ node-link avl-set ] keep swap
-        [ [ set-node-link ] keep ] dip
-        [ current-side get increase-balance balance-insert ]
-        [ f ] if
+        [ node-link avl-set ] keep -rot
+        [ [ set-node-link ] keep ] 2dip swap
+        [ [ current-side get increase-balance balance-insert ] dip ]
+        [ f swap ] if
     ] with-side ;
 
-: (avl-set) ( value key node -- node taller? )
+: (avl-set) ( value key node -- node taller? created? )
     2dup key>> = [
-        -rot pick key<< >>value f
+        -rot pick key<< >>value f f
     ] [ avl-insert ] if ;
 
-: avl-set ( value key node -- node taller? )
-    [ (avl-set) ] [ swap <avl-node> t ] if* ;
+: avl-set ( value key node -- node taller? created? )
+    [ (avl-set) ] [ swap <avl-node> t t ] if* ;
 
 M: avl set-at ( value key node -- )
-    [ avl-set drop ] change-root drop ;
+    [ avl-set nip swap ] change-root
+    swap [ dup inc-count ] when drop ;
 
 : delete-select-rotate ( node -- node shorter? )
     dup node+link balance>> zero? [
@@ -147,7 +148,8 @@ M: avl-node avl-delete ( key node -- node shorter? deleted? )
     ] if-zero ;
 
 M: avl delete-at ( key node -- )
-    [ avl-delete 2drop ] change-root drop ;
+    [ avl-delete nip swap ] change-root
+    swap [ dup dec-count ] when drop ;
 
 M: avl new-assoc 2drop <avl> ;