trees, add range operations (subtree>alist etc.)

extra/trees/trees-docs.factor

@@ -24,6 +24,67 @@ HELP: height
 }
 { $description "Returns the height of " { $snippet "tree" } "." } ;
 
+HELP: headtree>alist[)
+{ $values
+    { "to-key" "a key" } { "tree" tree }
+    { "alist" "an array of key/value pairs" }
+}
+{ $description "Returns an alist of the portion of this tree whose keys are strictly less than to-key." } ;
+
+HELP: headtree>alist[]
+{ $values
+    { "to-key" "a key" } { "tree" tree }
+    { "alist" "an array of key/value pairs" }
+}
+{ $description "Returns an alist of the portion of this tree whose keys are less than or equal to to-key." } ;
+
+HELP: subtree>alist()
+{ $values
+    { "from-key" "a key" } { "to-key" "a key" } { "tree" tree }
+    { "alist" "an array of key/value pairs" }
+}
+{ $description "Returns an alist of the portion of this map whose keys range from fromKey (exclusive) to toKey (exclusive)." } ;
+
+HELP: subtree>alist(]
+{ $values
+    { "from-key" "a key" } { "to-key" "a key" } { "tree" tree }
+    { "alist" "an array of key/value pairs" }
+}
+{ $description "Returns an alist of the portion of this map whose keys range from fromKey (exclusive) to toKey (inclusive)." } ;
+
+HELP: subtree>alist[)
+{ $values
+    { "from-key" "a key" } { "to-key" "a key" } { "tree" tree }
+    { "alist" "an array of key/value pairs" }
+}
+{ $description "Returns an alist of the portion of this map whose keys range from fromKey (inclusive) to toKey (exclusive)." } ;
+
+HELP: subtree>alist[]
+{ $values
+    { "from-key" "a key" } { "to-key" "a key" } { "tree" tree }
+    { "alist" "an array of key/value pairs" }
+}
+{ $description "Returns an alist of the portion of this map whose keys range from fromKey (inclusive) to toKey (inclusive)." } ;
+
+HELP: tailtree>alist(]
+{ $values
+    { "from-key" "a key" } { "tree" tree }
+    { "alist" "an array of key/value pairs" }
+}
+{ $description "Returns an alist of the portion of this tree whose keys are strictly greater than to-key." } ;
+
+HELP: tailtree>alist[]
+{ $values
+    { "from-key" "a key" } { "tree" tree }
+    { "alist" "an array of key/value pairs" }
+}
+{ $description "Returns an alist of the portion of this tree whose keys are greater than or equal to to-key." } ;
+
+{
+    headtree>alist[) headtree>alist[] tailtree>alist(] tailtree>alist[]
+    subtree>alist() subtree>alist(] subtree>alist[) subtree>alist[]
+} related-words
+
 ARTICLE: "trees" "Binary search trees"
 "This is a library for unbalanced binary search trees. 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."
 { $subsections
@@ -32,6 +93,12 @@ ARTICLE: "trees" "Binary search trees"
     >tree
     POSTPONE: TREE{
     height
-} ;
+}
+"Trees support range operations:"
+{ $subsections
+    headtree>alist[) headtree>alist[] tailtree>alist(] tailtree>alist[]
+    subtree>alist() subtree>alist(] subtree>alist[) subtree>alist[]
+}
+;
 
 ABOUT: "trees"

extra/trees/trees-tests.factor

@@ -1,5 +1,6 @@
-USING: accessors assocs kernel namespaces random tools.test
+USING: accessors arrays assocs combinators kernel math
-trees trees.private ;
+math.combinatorics math.ranges namespaces random sequences
+sequences.product tools.test trees trees.private ;
 IN: trees.tests
 
 : test-tree ( -- tree )
@@ -74,3 +75,59 @@ M: constant-random random-32* pattern>> ;
         1 over delete-at
     ] with-variable
 ] unit-test
+
+CONSTANT: test-tree2 TREE{
+        { 110 110 }
+        { 114 114 }
+        { 106 106 }
+        { 108 108 }
+        { 104 104 }
+        { 112 112 }
+        { 116 116 }
+        { 118 118 }
+        { 120 120 }
+        { 102 102 }
+        { 100 100 }
+    }
+
+: ?a,b? ( a b ? ? -- range )
+    2array {
+        { { t t } [ [a,b] ] }
+        { { t f } [ [a,b) ] }
+        { { f t } [ (a,b] ] }
+        { { f f } [ (a,b) ] }
+    } case ;
+
+! subtree>alist
+: test-tree2-subtree>alist ( a b ? ? -- subalist )
+    ?a,b? >array [ even? ] filter [ dup 2array ] map ;
+
+: subtree>alist ( from-key to-key tree start-inclusive? end-inclusive? -- alist )
+    2array {
+        { { t f } [ subtree>alist[) ] }
+        { { f t } [ subtree>alist(] ] }
+        { { t t } [ subtree>alist[] ] }
+        { { f f } [ subtree>alist() ] }
+    } case ;
+
+99 121 [a,b] 2 all-combinations
+{ t f } dup 2array <product-sequence> 2array
+[ first2 [ first2 ] bi@
+    {
+        [ test-tree2-subtree>alist 1array ]
+        [ [ [ test-tree2 ] 2dip subtree>alist ] 2curry 2curry unit-test ]
+    } 4cleave
+] product-each
+
+{ { } } [ 100 120 TREE{ } clone subtree>alist[] ] unit-test
+{ { } } [ 120 TREE{ } clone headtree>alist[] ] unit-test
+{ { } } [ 100 TREE{ } clone tailtree>alist[] ] unit-test
+
+{ { 100 102 104 106 108 110 112 114 } }
+[ 114 test-tree2 headtree>alist[] keys ] unit-test
+{ { 100 102 104 106 108 110 112 } }
+[ 114 test-tree2 headtree>alist[) keys ] unit-test
+{ { 106 108 110 112 114 116 118 120 } }
+[ 106 test-tree2 tailtree>alist[] keys ] unit-test
+{ { 108 110 112 114 116 118 120 } }
+[ 106 test-tree2 tailtree>alist(] keys ] unit-test

extra/trees/trees.factor

@@ -1,7 +1,7 @@
 ! Copyright (C) 2007 Alex Chapman
 ! See http://factorcode.org/license.txt for BSD license.
 USING: accessors arrays assocs combinators
-combinators.short-circuit kernel make math math.order namespaces
+combinators.short-circuit kernel locals make math math.order namespaces
 parser prettyprint.custom random ;
 IN: trees
 
@@ -122,10 +122,13 @@ M: tree set-at
 
 : valid-tree? ( tree -- ? ) root>> valid-node? ;
 
+: node-alist, ( node -- )
+    [ key>> ] [ value>> ] bi 2array , ;
+
 : (node>alist) ( node -- )
     [
         [ left>> (node>alist) ]
-        [ [ key>> ] [ value>> ] bi 2array , ]
+        [ node-alist, ]
         [ right>> (node>alist) ]
         tri
     ] when* ;
@@ -133,6 +136,92 @@ M: tree set-at
 M: tree >alist
     [ root>> (node>alist) ] { } make ;
 
+:: (node>subalist-right) ( to-key node end-comparator: ( key1 key2 -- ? ) -- )
+    node [
+        node key>> to-key end-comparator call :> node-left?
+
+        node left>> node-left? [ (node>alist) ] [
+            [ to-key ] dip end-comparator (node>subalist-right)
+        ] if
+
+        node-left? [
+            node [ node-alist, ] [
+                right>> [ to-key ] dip
+                end-comparator (node>subalist-right)
+            ] bi
+        ] when
+    ] when ; inline recursive
+
+:: (node>subalist-left) ( from-key node start-comparator: ( key1 key2 -- ? ) -- )
+    node [
+        node key>> from-key start-comparator call :> node-right?
+
+        node-right? [
+            node [
+                left>> [ from-key ] dip
+                start-comparator (node>subalist-left)
+            ] [ node-alist, ] bi
+        ] when
+
+        node right>> node-right? [ (node>alist) ] [
+            [ from-key ] dip start-comparator (node>subalist-left)
+        ] if
+    ] when ; inline recursive
+
+:: (node>subalist) ( from-key to-key node start-comparator: ( key1 key2 -- ? ) end-comparator: ( key1 key2 -- ? ) -- )
+    node [
+        node key>> from-key start-comparator call :> node-right?
+        node key>> to-key end-comparator call :> node-left?
+
+        node-right? [
+            from-key node left>> node-left?
+            [ start-comparator (node>subalist-left) ]
+            [
+                [ to-key ] dip start-comparator
+                end-comparator (node>subalist)
+            ] if
+        ] when
+
+        node-right? node-left? and [ node node-alist, ] when
+
+        node-left? [
+            to-key node right>> node-right?
+            [ end-comparator (node>subalist-right) ]
+            [
+                 [ from-key ] 2dip start-comparator
+                 end-comparator (node>subalist)
+            ] if
+        ] when
+    ] when ; inline recursive
+
+PRIVATE>
+
+: subtree>alist[) ( from-key to-key tree -- alist )
+    [ root>> [ after=? ] [ before? ] (node>subalist) ] { } make ;
+
+: subtree>alist(] ( from-key to-key tree -- alist )
+    [ root>> [ after? ] [ before=? ] (node>subalist) ] { } make ;
+
+: subtree>alist[] ( from-key to-key tree -- alist )
+    [ root>> [ after=? ] [ before=? ] (node>subalist) ] { } make ;
+
+: subtree>alist() ( from-key to-key tree -- alist )
+    [ root>> [ after? ] [ before? ] (node>subalist) ] { } make ;
+
+: headtree>alist[) ( to-key tree -- alist )
+    [ root>> [ before? ] (node>subalist-right) ] { } make ;
+
+: headtree>alist[] ( to-key tree -- alist )
+    [ root>> [ before=? ] (node>subalist-right) ] { } make ;
+
+: tailtree>alist[] ( from-key tree -- alist )
+    [ root>> [ after=? ] (node>subalist-left) ] { } make ;
+
+: tailtree>alist(] ( from-key tree -- alist )
+    [ root>> [ after? ] (node>subalist-left) ] { } make ;
+
+<PRIVATE
+
 M: tree clear-assoc
     0 >>count
     f >>root drop ;