trees, add pop/slurp operations

extra/trees/trees-docs.factor

@@ -174,8 +174,39 @@ HELP: first-key
 
 { first-key first-entry last-key last-entry } related-words
 
+HELP: pop-tree-left
+{ $values
+    { "tree" tree }
+    { "pair/f" { $maybe pair } }
+}
+{ $description "Removes and returns a key-value mapping associated with the lowest key in this map, or " { $link f } " if the map is empty." } ;
+
+HELP: pop-tree-right
+{ $values
+    { "tree" tree }
+    { "pair/f" { $maybe pair } }
+}
+{ $description "Removes and returns a key-value mapping associated with the highest key in this map, or " { $link f } " if the map is empty." } ;
+
+{ pop-tree-left pop-tree-right } related-words
+
+HELP: slurp-tree-left
+{ $values
+    { "tree" tree } { "quot" { $quotation ( ... entry -- ... ) } }
+}
+{ $description "Removes entries from a tree from the left (lowest key) and processes them with the quotation until the tree is empty." } ;
+
+HELP: slurp-tree-right
+{ $values
+    { "tree" tree } { "quot" { $quotation ( ... entry -- ... ) } }
+}
+{ $description "Removes entries from a tree from the right (highest key) and processes them with the quotation until the tree is empty." } ;
+
+{ slurp-tree-left slurp-tree-right } related-words
+
 ARTICLE: "trees" "Binary search trees"
-"This is a library for unbalanced binary search " { $link tree } "s. It is not intended to be used directly in most situations but rather as a base class for new trees, because performance can degrade to linear time storage/retrieval by the number of keys. These binary search trees conform to the assoc protocol."
+"The " { $vocab-link "trees" } " vocabulary is a library for unbalanced binary search trees. A " { $link tree } " is not intended to be used directly in most situations but rather as a base class for new trees, because performance can degrade to linear time storage/retrieval by the number of keys. These binary search trees conform to the assoc protocol."
+$nl
 "Constructing trees:"
 { $subsections
     <tree>
@@ -198,6 +229,11 @@ ARTICLE: "trees" "Binary search trees"
     lower-key lower-entry higher-key higher-entry
     floor-key floor-entry ceiling-key ceiling-entry
 }
+"Pop/Slurp operations on trees:"
+{ $subsections
+    pop-tree-left pop-tree-right
+    slurp-tree-left slurp-tree-right
+}
 ;
 
 ABOUT: "trees"

extra/trees/trees-tests.factor

@@ -1,5 +1,5 @@
-USING: accessors arrays assocs combinators kernel math
-math.combinatorics math.ranges namespaces random sequences
+USING: accessors arrays assocs combinators fry kernel locals
+math math.combinatorics math.ranges namespaces random sequences
 sequences.product tools.test trees trees.private ;
 IN: trees.tests
 
@@ -191,3 +191,72 @@ CONSTANT: test-tree2 TREE{
 [ 106 test-tree2 tailtree>alist[] keys ] unit-test
 { { 108 110 112 114 116 118 120 } }
 [ 106 test-tree2 tailtree>alist(] keys ] unit-test
+
+
+{ { { 10 10 } TREE{ { 20 20 } { 30 30 } } } } [
+    TREE{ { 20 20 } { 10 10 } { 30 30 } } clone [
+        pop-tree-left
+    ] keep 2array
+] unit-test
+
+{ { { 30 30 } TREE{ { 20 20 } { 10 10 } } } } [
+    TREE{ { 20 20 } { 10 10 } { 30 30 } } clone [
+        pop-tree-right
+    ] keep 2array
+] unit-test
+
+{ { { 20 20 } TREE{ } } } [
+    TREE{ { 20 20 } } clone [
+        pop-tree-right
+    ] keep 2array
+] unit-test
+
+{ { { 20 20 } TREE{ } } } [
+    TREE{ { 20 20 } } clone [
+        pop-tree-left
+    ] keep 2array
+] unit-test
+
+{ f } [ TREE{ } pop-tree-left ] unit-test
+{ f } [ TREE{ } pop-tree-right ] unit-test
+
+: with-limited-calls ( n quot -- quot' )
+    [let
+        0 :> count!
+        '[ count _ >=
+            [ "too many calls" throw ]
+            [ count 1 + count! @ ] if
+         ]
+    ] ; inline
+
+
+{ V{ { 10 10 } { 15 10 } { 20 20 }
+     { 15 20 } { 30 30 } { 35 30 }
+} } [
+    TREE{ { 20 20 } { 10 10 } { 30 30 } } clone V{ } clone [
+        dupd 6 [ [
+                over first {
+                    { [ dup 20 mod zero? ] [ drop [ first2 swap 5 - ] dip set-at ] }
+                    { [ dup 10 mod zero? ] [ drop [ first2 swap 5 + ] dip set-at ] }
+                    [ 3drop ]
+                } cond
+            ] [ push ] bi-curry* bi
+        ] with-limited-calls 2curry slurp-tree-left
+    ] keep
+] unit-test
+
+{ V{
+    { 30 30 } { 25 30 } { 20 20 }
+    { 25 20 } { 10 10 } {  5 10 } }
+} [
+    TREE{ { 20 20 } { 10 10 } { 30 30 } } clone V{ } clone [
+        dupd 6 [ [
+                over first {
+                    { [ dup 20 mod zero? ] [ drop [ first2 swap 5 + ] dip set-at ] }
+                    { [ dup 10 mod zero? ] [ drop [ first2 swap 5 - ] dip set-at ] }
+                    [ 3drop ]
+                } cond
+            ] [ push ] bi-curry* bi
+        ] with-limited-calls 2curry slurp-tree-right
+    ] keep
+] unit-test

extra/trees/trees.factor

@@ -416,3 +416,37 @@ PRIVATE>
 
 : height ( tree -- n )
     root>> node-height ;
+
+<PRIVATE
+
+: (pop-tree-extremity) ( tree -- node/f )
+    dup root>> dup node-link
+    [ (prune-extremity) nip ]
+    [ [ delete-node swap root<< ] keep ] if* ;
+
+: pop-tree-extremity ( tree -- node/f )
+    [ (pop-tree-extremity) ] [ over [ dec-count ] [ drop ] if ] bi
+    node>entry ;
+
+: slurp-tree ( tree quot: ( ... entry -- ... ) -- ... )
+    [ drop [ count>> 0 = ] curry ]
+    [ [ [ pop-tree-extremity ] curry ] dip compose ] 2bi until ; inline
+
+: pop-tree ( tree -- entry )
+    dup root>> dup [
+        drop pop-tree-extremity
+    ] [ nip ] if ;
+
+PRIVATE>
+
+: pop-tree-left ( tree -- pair/f )
+    left [ pop-tree ] with-side ;
+
+: pop-tree-right ( tree -- pair/f )
+    right [ pop-tree ] with-side ;
+
+: slurp-tree-left ( tree quot: ( ... entry -- ... ) -- ... )
+    left [ slurp-tree ] with-side ; inline
+
+: slurp-tree-right ( tree quot: ( ... entry -- ... ) -- ... )
+    right [ slurp-tree ] with-side ; inline