diff --git a/extra/trees/avl/avl-tests.factor b/extra/trees/avl/avl-tests.factor index 9b9ec8ea5b..a9086e5bf1 100644 --- a/extra/trees/avl/avl-tests.factor +++ b/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 diff --git a/extra/trees/avl/avl.factor b/extra/trees/avl/avl.factor index dacc21346a..7897b6dbe1 100644 --- a/extra/trees/avl/avl.factor +++ b/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 t ] if* ; +: avl-set ( value key node -- node taller? created? ) + [ (avl-set) ] [ swap 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 ;