Cleaning up trees code a little bit

extra/trees/avl/avl-tests.factor

@@ -1,5 +1,5 @@
 USING: kernel tools.test trees trees.avl math random sequences
-assocs accessors ;
+assocs accessors trees.avl.private trees.private ;
 IN: trees.avl.tests
 
 [ "key1" 0 "key2" 0 ] [

extra/trees/avl/avl.factor

@@ -2,7 +2,8 @@
 ! See http://factorcode.org/license.txt for BSD license.
 USING: combinators kernel generic math math.functions
 math.parser namespaces io sequences trees shuffle
-assocs parser accessors math.order prettyprint.custom ;
+assocs parser accessors math.order prettyprint.custom
+trees.private ;
 IN: trees.avl
 
 TUPLE: avl < tree ;
@@ -10,6 +11,8 @@ TUPLE: avl < tree ;
 : <avl> ( -- tree )
     avl new-tree ;
 
+<PRIVATE
+
 TUPLE: avl-node < node balance ;
 
 : <avl-node> ( key value -- node )
@@ -20,11 +23,14 @@ TUPLE: avl-node < node balance ;
     swap [ + ] change-balance drop ;
 
 : rotate ( node -- node )
-    dup node+link dup node-link pick set-node+link
+    dup node+link
-    tuck set-node-link ;    
+    dup node-link
+    pick set-node+link
+    [ set-node-link ] keep ;    
 
 : single-rotate ( node -- node )
-    0 over (>>balance) 0 over node+link 
+    0 >>balance
+    0 over node+link 
     (>>balance) rotate ;
 
 : pick-balances ( a node -- balance balance )
@@ -61,7 +67,7 @@ DEFER: avl-set
 : avl-insert ( value key node -- node taller? )
     2dup key>> before? left right ? [
         [ node-link avl-set ] keep swap
-        [ tuck set-node-link ] dip
+        [ [ set-node-link ] keep ] dip
         [ dup current-side get increase-balance balance-insert ]
         [ f ] if
     ] with-side ;
@@ -146,6 +152,8 @@ M: avl delete-at ( key node -- )
 
 M: avl new-assoc 2drop <avl> ;
 
+PRIVATE>
+
 : >avl ( assoc -- avl )
     T{ avl f f 0 } assoc-clone-like ;
 

extra/trees/splay/splay.factor

@@ -1,7 +1,8 @@
 ! Copyright (c) 2005 Mackenzie Straight.
 ! See http://factorcode.org/license.txt for BSD license.
 USING: arrays kernel math namespaces sequences assocs parser
-trees generic math.order accessors prettyprint.custom shuffle ;
+trees generic math.order accessors prettyprint.custom
+trees.private combinators ;
 IN: trees.splay
 
 TUPLE: splay < tree ;
@@ -9,6 +10,8 @@ TUPLE: splay < tree ;
 : <splay> ( -- tree )
     \ splay new-tree ;
 
+<PRIVATE
+
 : rotate-right ( node -- node )
     dup left>>
     [ right>> swap (>>left) ] 2keep
@@ -27,32 +30,35 @@ TUPLE: splay < tree ;
     swap [ rot [ (>>right) ] 2keep
     drop dup right>> swapd ] dip swap ;
 
-: cmp ( key node -- obj node -1/0/1 )
+: cmp ( key node -- obj node <=> )
-    2dup key>> key-side ;
+    2dup key>> <=> ;
 
-: lcmp ( key node -- obj node -1/0/1 ) 
+: lcmp ( key node -- obj node <=> ) 
-    2dup left>> key>> key-side ;
+    2dup left>> key>> <=> ;
 
-: rcmp ( key node -- obj node -1/0/1 ) 
+: rcmp ( key node -- obj node <=> ) 
-    2dup right>> key>> key-side ;
+    2dup right>> key>> <=> ;
 
 DEFER: (splay)
 
 : splay-left ( left right key node -- left right key node )
     dup left>> [
-        lcmp 0 < [ rotate-right ] when
+        lcmp +lt+ = [ rotate-right ] when
         dup left>> [ link-right (splay) ] when
     ] when ;
 
 : splay-right ( left right key node -- left right key node )
     dup right>> [
-        rcmp 0 > [ rotate-left ] when
+        rcmp +gt+ = [ rotate-left ] when
         dup right>> [ link-left (splay) ] when
     ] when ;
 
 : (splay) ( left right key node -- left right key node )
-    cmp dup 0 <
+    cmp {
-    [ drop splay-left ] [ 0 > [ splay-right ] when ] if ;
+        { +lt+ [ splay-left ] }
+        { +gt+ [ splay-right ] }
+        { +eq+ [ ] }
+    } case ;
 
 : assemble ( head left right node -- root )
     [ right>> swap (>>left) ] keep
@@ -64,18 +70,18 @@ DEFER: (splay)
     [ T{ node } clone dup dup ] 2dip
     (splay) nip assemble ;
 
-: splay ( key tree -- )
+: do-splay ( key tree -- )
     [ root>> splay-at ] keep (>>root) ;
 
 : splay-split ( key tree -- node node )
-    2dup splay root>> cmp 0 < [
+    2dup do-splay root>> cmp +lt+ = [
         nip dup left>> swap f over (>>left)
     ] [
         nip dup right>> swap f over (>>right) swap
     ] if ;
 
 : get-splay ( key tree -- node ? )
-    2dup splay root>> cmp 0 = [
+    2dup do-splay root>> cmp +eq+ = [
         nip t
     ] [
         2drop f f
@@ -95,7 +101,7 @@ DEFER: (splay)
     ] if* ;
 
 : remove-splay ( key tree -- )
-    tuck get-splay nip [
+    [ get-splay nip ] keep [
         dup dec-count
         dup right>> swap left>> splay-join
         swap (>>root)
@@ -128,6 +134,8 @@ M: splay delete-at ( key tree -- )
 M: splay new-assoc
     2drop <splay> ;
 
+PRIVATE>
+
 : >splay ( assoc -- tree )
     T{ splay f f 0 } assoc-clone-like ;
 

extra/trees/trees.factor

@@ -2,22 +2,27 @@
 ! See http://factorcode.org/license.txt for BSD license.
 USING: kernel generic math sequences arrays io namespaces
 prettyprint.private kernel.private assocs random combinators
-parser math.order accessors deques make prettyprint.custom 
+parser math.order accessors deques make prettyprint.custom ;
-shuffle ;
 IN: trees
 
 TUPLE: tree root count ;
 
+<PRIVATE
+
 : new-tree ( class -- tree )
     new
         f >>root
         0 >>count ; inline
 
+PRIVATE>
+
 : <tree> ( -- tree )
     tree new-tree ;
 
 INSTANCE: tree assoc
 
+<PRIVATE
+
 TUPLE: node key value left right ;
 
 : new-node ( key value class -- node )
@@ -61,7 +66,7 @@ CONSTANT: right 1
 : set-node+link ( child node -- ) t set-node-link@ ;
 
 : with-side ( side quot -- )
-    [ swap current-side set call ] with-scope ; inline
+    [ current-side ] dip with-variable ; inline
 
 : with-other-side ( quot -- )
     current-side get neg swap with-side ; inline
@@ -137,9 +142,9 @@ DEFER: delete-node
 
 : (prune-extremity) ( parent node -- new-extremity )
     dup node-link [
-        rot drop (prune-extremity)
+        [ nip ] dip (prune-extremity)
     ] [
-        tuck delete-node swap set-node-link
+        [ delete-node ] [ set-node-link ] bi
     ] if* ;
 
 : prune-extremity ( node -- new-extremity )
@@ -183,9 +188,15 @@ DEFER: delete-node
     2dup key>> key-side dup 0 eq? [
         drop nip delete-node
     ] [
-        [ tuck node-link delete-bst-node over set-node-link ] with-side
+        [
+            [ node-link delete-bst-node ]
+            [ set-node-link ]
+            [ ] tri
+        ] with-side
     ] if ;
 
+PRIVATE>
+
 M: tree delete-at
     [ delete-bst-node ] change-root drop ;