trees, add navigation operations (lower-key etc.)

extra/trees/trees-docs.factor

@@ -1,4 +1,4 @@
-USING: assocs help.markup help.syntax math ;
+USING: arrays assocs help.markup help.syntax kernel math ;
 IN: trees
 
 HELP: TREE{
@@ -85,20 +85,119 @@ HELP: tailtree>alist[]
     subtree>alist() subtree>alist(] subtree>alist[) subtree>alist[]
 } related-words
 
+HELP: ceiling-entry
+{ $values
+    { "key" "a key" } { "tree" tree }
+    { "pair/f" { $maybe pair } }
+}
+{ $description "Returns a key-value mapping associated with the least key greater than or equal to the given key, or " { $link f } " if there is no such key." } ;
+
+HELP: ceiling-key
+{ $values
+    { "key" "a key" } { "tree" tree }
+    { "key/f" { $maybe "a key" } }
+}
+{ $description "Returns the least key greater than or equal to the given key, or " { $link f } " if there is no such key." } ;
+
+HELP: floor-entry
+{ $values
+    { "key" "a key" } { "tree" tree }
+    { "pair/f" { $maybe pair } }
+}
+{ $description "Returns a key-value mapping associated with the greatest key less than or equal to the given key, or " { $link f } " if there is no such key." } ;
+
+HELP: floor-key
+{ $values
+    { "key" "a key" } { "tree" tree }
+    { "key/f" { $maybe "a key" } }
+}
+{ $description "Returns the greatest key less than or equal to the given key, or " { $link f } " if there is no such key." } ;
+
+HELP: higher-entry
+{ $values
+    { "key" "a key" } { "tree" tree }
+    { "pair/f" { $maybe pair } }
+}
+{ $description "Returns a key-value mapping associated with the least key strictly greater than the given key, or " { $link f } " if there is no such key." } ;
+
+HELP: higher-key
+{ $values
+    { "key" "a key" } { "tree" tree }
+    { "key/f" { $maybe "a key" } }
+}
+{ $description "Returns the least key strictly greater than the given key, or " { $link f } " if there is no such key." } ;
+
+HELP: lower-entry
+{ $values
+    { "key" "a key" } { "tree" tree }
+    { "pair/f" { $maybe pair } }
+}
+{ $description "Returns a key-value mapping associated with the greatest key strictly less than the given key, or " { $link f } " if there is no such key." } ;
+
+HELP: lower-key
+{ $values
+    { "key" "a key" } { "tree" tree }
+    { "key/f" { $maybe "a key" } }
+}
+{ $description "Returns the greatest key strictly less than the given key, or " { $link f } " if there is no such key." } ;
+
+{ lower-key lower-entry higher-key higher-entry
+  floor-key floor-entry ceiling-key ceiling-entry } related-words
+
+HELP: last-entry
+{ $values
+    { "tree" tree }
+    { "pair/f" { $maybe pair } }
+}
+{ $description "Returns a key-value mapping associated with the last (highest) key in this tree, or " { $link f } " if the tree is empty." } ;
+
+HELP: last-key
+{ $values
+    { "tree" tree }
+    { "key/f" { $maybe "a key" } }
+}
+{ $description "Returns the last (highest) key in this tree, or " { $link f } " if the tree is empty." } ;
+
+HELP: first-entry
+{ $values
+    { "tree" tree }
+    { "pair/f" { $maybe pair } }
+}
+{ $description "Returns a key-value mapping associated with the first (lowest) key in this tree, or " { $link f } " if the tree is empty." } ;
+
+HELP: first-key
+{ $values
+    { "tree" tree }
+    { "key/f" { $maybe pair } }
+}
+{ $description "Returns the first (lowest) key in this tree, or " { $link f } " if the tree is empty." } ;
+
+{ first-key first-entry last-key last-entry } 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."
+"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."
+"Constructing trees:"
 { $subsections
-    tree
     <tree>
     >tree
     POSTPONE: TREE{
-    height
 }
-"Trees support range operations:"
+"Operations on trees: "
+{ $subsections
+    height
+    first-entry first-key
+    last-entry last-key
+}
+"Range operations on trees:"
 { $subsections
     headtree>alist[) headtree>alist[] tailtree>alist(] tailtree>alist[]
     subtree>alist() subtree>alist(] subtree>alist[) subtree>alist[]
 }
+"Navigation operations on trees:"
+{ $subsections
+    lower-key lower-entry higher-key higher-entry
+    floor-key floor-entry ceiling-key ceiling-entry
+}
 ;
 
 ABOUT: "trees"

extra/trees/trees-tests.factor

@@ -90,6 +90,66 @@ CONSTANT: test-tree2 TREE{
         { 100 100 }
     }
 
+: test-tree2-lower-key ( key -- key' )
+    dup 2 mod 2 swap - - ;
+: test-tree2-higher-key ( key -- key' )
+    dup 2 mod 2 swap - + ;
+: test-tree2-floor-key ( key -- key' )
+    dup 2 mod - ;
+: test-tree2-ceiling-key ( key -- key' )
+    dup 2 mod + ;
+
+{ f } [ 99 test-tree2 lower-node ]  unit-test
+{ f } [ 100 test-tree2 lower-node ]  unit-test
+100 121 (a,b] [
+    [ test-tree2-lower-key 1array ] keep [ test-tree2 lower-node key>> ] curry unit-test
+] each
+
+99 120 [a,b) [
+    [ test-tree2-higher-key 1array ] keep [ test-tree2 higher-node key>> ] curry unit-test
+] each
+{ f } [ 120 test-tree2 higher-node ]  unit-test
+{ f } [ 121 test-tree2 higher-node ]  unit-test
+
+{ f } [ 99 test-tree2 floor-node ]  unit-test
+100 121 [a,b] [
+    [ test-tree2-floor-key 1array ] keep [ test-tree2 floor-node key>> ] curry unit-test
+] each
+
+99 120 [a,b] [
+    [ test-tree2-ceiling-key 1array ] keep [ test-tree2 ceiling-node key>> ] curry unit-test
+] each
+{ f } [ 121 test-tree2 ceiling-node ]  unit-test
+
+{ 100 } [ test-tree2 first-node key>> ] unit-test
+{ 120 } [ test-tree2 last-node key>> ] unit-test
+
+{ f } [ 99 test-tree2 lower-entry ] unit-test
+{ f } [ 99 test-tree2 lower-key ] unit-test
+{ f } [ 121 test-tree2 higher-entry ] unit-test
+{ f } [ 121 test-tree2 higher-key ] unit-test
+{ f } [ 99 test-tree2 floor-entry ] unit-test
+{ f } [ 99 test-tree2 floor-key ] unit-test
+{ f } [ 121 test-tree2 ceiling-entry ] unit-test
+{ f } [ 121 test-tree2 ceiling-key ] unit-test
+{ { 108 108 } } [ 110 test-tree2 lower-entry ] unit-test
+{ 108  } [ 110 test-tree2 lower-key ] unit-test
+{ { 112 112 } } [ 110 test-tree2 higher-entry ] unit-test
+{ 112 } [ 110 test-tree2 higher-key ] unit-test
+{ { 110 110 } } [ 110 test-tree2 floor-entry ] unit-test
+{ 110 } [ 110 test-tree2 floor-key ] unit-test
+{ { 110 110 } } [ 110 test-tree2 ceiling-entry ] unit-test
+{ 110 } [ 110 test-tree2 ceiling-key ] unit-test
+
+{ f } [ TREE{ } clone first-key ] unit-test
+{ f } [ TREE{ } clone first-entry ] unit-test
+{ f } [ TREE{ } clone last-key ] unit-test
+{ f } [ TREE{ } clone last-entry ] unit-test
+{ { 100 100 } } [ test-tree2 first-entry ] unit-test
+{ 100 } [ test-tree2 first-key ] unit-test
+{ { 120 120 } } [ test-tree2 last-entry ] unit-test
+{ 120 } [ test-tree2 last-key ] unit-test
+
 : ?a,b? ( a b ? ? -- range )
     2array {
         { { t t } [ [a,b] ] }

extra/trees/trees.factor

@@ -1,8 +1,8 @@
 ! Copyright (C) 2007 Alex Chapman
 ! See http://factorcode.org/license.txt for BSD license.
 USING: accessors arrays assocs combinators
-combinators.short-circuit kernel locals make math math.order namespaces
+combinators.short-circuit kernel locals make math math.order
-parser prettyprint.custom random ;
+namespaces parser prettyprint.custom random sequences ;
 IN: trees
 
 TUPLE: tree root { count integer } ;
@@ -222,6 +222,91 @@ PRIVATE>
 
 <PRIVATE
 
+: (nodepath-at) ( key node -- )
+    [
+        dup ,
+        2dup key>> = [
+            2drop
+        ] [
+            choose-branch (nodepath-at)
+        ] if
+    ] [ drop ] if* ;
+
+: nodepath-at ( key tree -- path )
+    [ root>> (nodepath-at) ] { } make ;
+
+: right-extremity ( node -- node' )
+    [ dup right>> dup ] [ nip ] while drop ;
+
+: left-extremity ( node -- node' )
+    [ dup left>> dup ] [ nip ] while drop ;
+
+: lower-node-in-child? ( key node -- ? )
+    [ nip left>> ] [ key>> = ] 2bi and ;
+
+: higher-node-in-child? ( key node -- ? )
+    [ nip right>> ] [ key>> = ] 2bi and ;
+
+: lower-node ( key tree -- node )
+    dupd nodepath-at
+    [ drop f ] [
+        reverse 2dup first lower-node-in-child?
+        [ nip first left>> right-extremity ]
+        [ [ key>> after? ] with find nip ] if
+    ] if-empty ;
+
+: higher-node ( key tree -- node )
+    dupd nodepath-at
+    [ drop f ] [
+        reverse 2dup first higher-node-in-child?
+        [ nip first right>> left-extremity ]
+        [ [ key>> before? ] with find nip ] if
+    ] if-empty ;
+
+: floor-node ( key tree -- node )
+    dupd nodepath-at [ drop f ] [
+        reverse [ key>> after=? ] with find nip
+    ] if-empty ;
+
+: ceiling-node ( key tree -- node )
+    dupd nodepath-at [ drop f ] [
+        reverse [ key>> before=? ] with find nip
+    ] if-empty ;
+
+: first-node ( tree -- node ) root>> dup [ left-extremity ] when ;
+
+: last-node ( tree -- node ) root>> dup [ right-extremity ] when ;
+
+: node>entry ( node -- entry ) [ key>> ] [ value>> ] bi 2array ;
+
+PRIVATE>
+
+: lower-entry ( key tree -- pair/f ) lower-node dup [ node>entry ] when ;
+
+: higher-entry ( key tree -- pair/f ) higher-node dup [ node>entry ] when ;
+
+: floor-entry ( key tree -- pair/f ) floor-node dup [ node>entry ] when ;
+
+: ceiling-entry ( key tree -- pair/f ) ceiling-node dup [ node>entry ] when ;
+
+: first-entry ( tree -- pair/f ) first-node dup [ node>entry ] when ;
+
+: last-entry ( tree -- pair/f ) last-node dup [ node>entry ] when ;
+
+: lower-key ( key tree -- key/f ) lower-node dup [ key>> ] when ;
+
+: higher-key ( key tree -- key/f ) higher-node dup [ key>> ] when ;
+
+: floor-key ( key tree -- key/f ) floor-node dup [ key>> ] when ;
+
+: ceiling-key ( key tree -- key/f ) ceiling-node dup [ key>> ] when ;
+
+: first-key ( tree -- key/f ) first-node dup [ key>> ] when ;
+
+: last-key ( tree -- key/f ) last-node dup [ key>> ] when ;
+
+<PRIVATE
+
 M: tree clear-assoc
     0 >>count
     f >>root drop ;