diff --git a/extra/trees/trees-tests.factor b/extra/trees/trees-tests.factor index 82cf3c812a..2172f9add0 100644 --- 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 b64168a573..8b3b414033 100644 --- 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 ] if* ] keep - [ set-node-link ] keep + [ node-link [ node-set ] [ swap t ] if* ] keep + swap [ [ set-node-link ] keep ] dip ] with-side ] if ; M: tree set-at - [ [ node-set ] [ swap ] if* ] change-root drop ; + [ [ node-set ] [ swap 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 ;