splay tree work

collects/data/splay-tree.rkt

@@ -3,9 +3,6 @@
          racket/dict
          racket/contract)
 
-;; FIXME: need special handling of +/- inf.0 ! (otherwise, other keys get killed)
-;; Idea: in traversal, just treat +/-inf.0 as 0 for key-adjustment.
-
 ;; ======== Raw splay tree ========
 
 (struct node (key value left right) #:mutable #:transparent)
@@ -113,22 +110,35 @@ In (values status nroot pside pnode):
              ((>) (SPisplay (find cmp tx k* (node-right x) 'right x add-v) 'right x))))]
         [add-v
          (let ([new-node (node k (car add-v) #f #f)])
-           ;; link unnecessary? will be done in isplay/finish?
+           ;; FIXME: link unnecessary? will be done in isplay/finish?
            (when p (set-node-side! p p-side new-node))
            (SPunit/add tx new-node))]
         [else (SPfail tx)]))
 
-;; isplay! : ... -> node
+(define (find-min tx x)
+  (define (find-min-loop x)
+    (cond [(and x (node-left x))
+           (SPisplay (find-min-loop (node-left x)) 'left x)]
+          [x (SPunit tx x)]
+          [else (SPfail tx)]))
+  (SPfinish (find-min-loop x)))
+
+(define (find-max tx x)
+  (define (find-max-loop x)
+    (cond [(and x (node-right x))
+           (SPisplay (find-max-loop (node-right x)) 'right x)]
+          [x (SPunit tx x)]
+          [else (SPfail tx)]))
+  (SPfinish (find-max-loop x)))
+
+;; isplay! : ... -> SP
 ;; incremental splay
 (define (isplay! tx ok? x p-side p gp-side gp)
-  ;; (printf "splay! ~s\n" (list x p-side p gp-side gp))
-  (printf "splay!\n")
   (cond [(eq? x #f)
          ;; Then p-side = #f, p = #f
          ;; Overwrite new root with gp
          (values tx ok? gp #f #f)]
         [p-side ;; we have two splay path segments; splay
-         ;; First, link x as p.p-side
          (set-node-side! p p-side x)
          (cond [(eq? p-side gp-side)
                 ;; zig-zig
@@ -146,17 +156,14 @@ In (values status nroot pside pnode):
          (values tx ok? x gp-side gp)]))
 
 (define (finish tx ok? x p-side p)
-  ;; (printf "run ~s\n" (list x p-side p))
-  (printf "finish!\n")
   (cond [(eq? x #f)
          ;; Then p-side = #f, p = #f
          (values ok? #f)]
-        [p-side ;; one splay segment left; perform zig
+        [p-side ;; one splay path segment left; perform zig
-         ;; First, link x as p.p-side
          (set-node-side! p p-side x)
          (rotate! tx p p-side)
          (values ok? x)]
-        [else ;; no splay segments left
+        [else ;; no splay path segments left
          (values ok? x)]))
 
 (define (set-node-side! n side v)
@@ -179,8 +186,7 @@ In (values status nroot pside pnode):
        (set-node-key! p (- 0 Kx))
        (set-node-key! x (+ Kp Kx))
        (when B
-         (set-node-key! B (+ (node-key B) Kx))))
+         (set-node-key! B (+ (node-key B) Kx))))]))
-     (sanity! tx 'right x)]))
 
 (define (left! tx p)
   (match p
@@ -191,18 +197,7 @@ In (values status nroot pside pnode):
        (set-node-key! p (- 0 Kx))
        (set-node-key! x (+ Kp Kx))
        (when B
-         (set-node-key! B (+ (node-key B) Kx))))
+         (set-node-key! B (+ (node-key B) Kx))))]))
-     (sanity! tx 'left x)]))
-
-(define (sanity! tx who x0)
-  (when tx 
-    (let loop ([x x0] [sign? void])
-      (when (node? x)
-        (unless (sign? (node-key x))
-          (printf "x0 = ~s\n" x0)
-          (error 'insane! "~s: insane sub-node ~s" who x))
-        (loop (node-left x) negative?)
-        (loop (node-right x) positive?)))))
 
 ;; --------
 
@@ -210,7 +205,7 @@ In (values status nroot pside pnode):
 (define (join-left tx left right)
   (cond [(and left right)
          (let-values ([(_ok? left*) (find-max tx left)])
-           ;; left* is node, must have empty right branch
+           ;; left* is node, left*.right = #f
            (set-node-right! left* right)
            (when tx
              (set-node-key! right (- (node-key right) (node-key left*))))
@@ -258,41 +253,25 @@ In (values status nroot pside pnode):
   (let-values ([(left right) (split/drop-root tx root)])
     (join-left tx left right)))
 
-(define (find-min tx x)
+(define (remove-range! cmp tx root from to contract?)
-  (define (find-min-loop x)
-    (cond [(and x (node-left x))
-           (SPisplay (find-min-loop (node-left x)) 'left x)]
-          [x (SPunit tx x)]
-          [else (SPfail tx)]))
-  (SPfinish (find-min-loop x)))
-
-(define (find-max tx x)
-  (define (find-max-loop x)
-    (cond [(and x (node-right x))
-           (SPisplay (find-max-loop (node-right x)) 'right x)]
-          [x (SPunit tx x)]
-          [else (SPfail tx)]))
-  (SPfinish (find-max-loop x)))
-
-(define (contract! cmp tx root from to)
   ;; tx = #t... why pass as param?
-  (let*-values ([(ok? from-node) (find/root cmp tx root from (list #f))]
+  (let*-values ([(ok? from-node) (find/root cmp tx from root (list #f))]
                 [(left-tree right-tree)
                  (if (eq? ok? 'added)
                      (split/drop-root tx from-node)
                      (split/root-to-right tx from-node))]
-                [(ok? to-node) (find/root cmp tx right-tree to (list #f))]
+                [(ok? to-node) (find/root cmp tx to right-tree (list #f))]
                 [(mid-tree right-tree)
                  (if (eq? ok? 'added)
                      (split/drop-root tx to-node)
                      (split/root-to-right tx to-node))])
-    (when tx ;; ie, #t
+    (when (and tx contract?)
       (when right-tree
         (set-node-key! right-tree (+ (node-key right-tree) (- from to)))))
     (join-left tx left-tree right-tree)))
 
 (define (expand! cmp tx root from to)
-  (let*-values ([(ok? from-node) (find/root cmp tx root from (list #f))]
+  (let*-values ([(ok? from-node) (find/root cmp tx from root (list #f))]
                 [(left-tree right-tree)
                  (if (eq? ok? 'added)
                      (split/drop-root tx from-node)
@@ -319,15 +298,6 @@ In (values status nroot pside pnode):
 
 ;; ======== Splay tree ========
 
-(define (make-splay-tree <? =?)
-  (splay-tree #f 0 (lambda (x y) (if (=? x y) '= (if (<? x y) '< '>))) #f))
-
-#|
-In an integer splay tree, keys can be stored relative to their parent nodes.
-|#
-(define (make-integer-splay-tree)
-  (integer-splay-tree #f 0 (lambda (x y) (if (= x y) '= (if (< x y) '< '>))) #t))
-
 (define not-given (gensym 'not-given))
 
 (define (splay-tree-ref s x [default not-given])
@@ -348,8 +318,8 @@ In an integer splay tree, keys can be stored relative to their parent nodes.
     [(splay-tree root size cmp tx)
      (let-values ([(ok? root) (find/root cmp tx x root (list v))])
        (set-splay-tree-root! s root)
-       (when (eq? ok? 'added) (set-splay-tree-size! s (add1 size)))
+       (when (and (eq? ok? 'added) size)
-       ;; (printf "root = ~s\n" root)
+         (set-splay-tree-size! s (add1 size)))
        (unless (eq? (node-value root) v)
          (set-node-value! root v)))]))
 
@@ -359,10 +329,37 @@ In an integer splay tree, keys can be stored relative to their parent nodes.
      (let-values ([(ok? root) (find/root cmp tx x root #f)])
        (when ok? ;; => root is node
          (set-splay-tree-root! s (delete-root tx root))
-         (set-splay-tree-size! s (sub1 size))))]))
+         (when size (set-splay-tree-size! s (sub1 size)))))]))
 
 (define (splay-tree-count s)
-  (splay-tree-size s))
+  (let ([size (splay-tree-size s)])
+    (if size
+        size
+        (let ([size (let loop ([x (splay-tree-root s)] [n 0])
+                      (if x
+                          (loop (node-left x) (loop (node-right x) (add1 n)))
+                          n))])
+          (set-splay-tree-size! s size)
+          size))))
+
+(define (splay-tree-remove-range! s from to)
+  (match s
+    [(splay-tree root size cmp tx)
+     (set-splay-tree-root! s (remove-range! cmp tx root from to #f))
+     (set-splay-tree-size! s #f)]))
+
+(define (splay-tree-contract! s from to)
+  (match s
+    [(splay-tree root size cmp tx)
+     (set-splay-tree-root! s (remove-range! cmp tx root from to #t))
+     (set-splay-tree-size! s #f)]))
+
+(define (splay-tree-expand! s from to)
+  (match s
+    [(splay-tree root size cmp tx)
+     (set-splay-tree-root! s (expand! cmp tx root from to))]))
+
+;; ========
 
 #|
 Iteration in splay-trees is problematic.
@@ -375,8 +372,9 @@ Options
  1) position = parent chain (very likely to get out of sync)
  2) position = key (re-lookup each time)
  3) snapshot as alist (more allocation than necessary, sometimes much more)
+ 4) position = node (doesn't work with position-relative keys)
 
-(1) is no good. (3) is not very iterator-like.
+(1,4) are no good. (3) is not very iterator-like.
 
 (2) seems to be the best compromise.
 |#
@@ -393,10 +391,7 @@ Options
 (define (splay-tree-iterate-next s pos)
   (match pos
     [(splay-tree-iter key)
-     (let ([next (splay-tree-least-key/>? s key not-given)])
+     (splay-tree-iterate-least/>? s key)]))
-       (if (eq? next not-given)
-           #f
-           (splay-tree-iter next)))]))
 
 (define (splay-tree-iterate-key s pos)
   (match pos
@@ -404,9 +399,9 @@ Options
 
 (define (splay-tree-iterate-value s pos)
   (match pos
-    [(splay-tree-iter key)
+    [(splay-tree-iter key) (splay-tree-ref s key #f)]))
-     (splay-tree-ref s key #f)]))
 
+;; ========
 
 (struct splay-tree ([root #:mutable] [size #:mutable] cmp tx)
         #:transparent
@@ -426,7 +421,7 @@ Options
                                 splay-tree-iter?
                                 #f #f #f)))
 
-(struct integer-splay-tree splay-tree ()
+(struct splay-tree* splay-tree (key-c value-c)
         #:transparent
         #:property prop:dict/contract
         (list (vector-immutable splay-tree-ref
@@ -439,14 +434,49 @@ Options
                                 splay-tree-iterate-next
                                 splay-tree-iterate-key
                                 splay-tree-iterate-value)
-              (vector-immutable exact-integer?
+              (vector-immutable any/c
                                 any/c
                                 splay-tree-iter?
-                                #f #f #f)))
+                                (lambda (s) (splay-tree*-key-c s))
+                                (lambda (s) (splay-tree*-value-c s))
+                                #f)))
+
+(define-syntax-rule (mkcmp <? =?)
+  (lambda (x y) (cond [(=? x y) '=] [(<? x y) '<] [else '>])))
+
+(define (make-splay-tree <? =?
+                         #:key-contract [key-contract any/c]
+                         #:value-contract [value-contract any/c])
+  (cond [(and (eq? key-contract any/c) (eq? value-contract any/c))
+         (splay-tree #f 0 (mkcmp <? =?) #f)]
+        [else
+         (splay-tree* #f 0 (mkcmp <? =?) #f key-contract value-contract)]))
+
+#|
+In an integer splay tree, keys can be stored relative to their parent nodes.
+|#
+(define (make-integer-splay-tree #:adjust? [adjust? #f]
+                                 #:key-contract [key-contract any/c]
+                                 #:value-contract [value-contract any/c])
+  (splay-tree* #f 0 (mkcmp < =) (and adjust? #t)
+               (if (eq? key-contract any/c)
+                   exact-integer?
+                   (and/c exact-integer? key-contract))
+               value-contract))
+
+(define (splay-tree-with-adjust? s)
+  (splay-tree-tx s))
+
+(define (key-c s)
+  (if (splay-tree*? s) (splay-tree*-key-c s) any/c))
+(define (val-c s)
+  (if (splay-tree*? s) (splay-tree*-value-c s) any/c))
+
+;; ========
 
 ;; Order-based search
 
-(define (extreme who s key cmp-result has-X? find-X default)
+(define (extreme who s key cmp-result has-X? find-X)
   (match s
     [(splay-tree root size cmp tx)
      (let*-values ([(_ok? root) (find/root cmp tx key root #f)]
@@ -458,27 +488,16 @@ Options
                           [else
                            (values #f root)])])
        (set-splay-tree-root! s root)
-       (if ok?
+       (and ok? (splay-tree-iter root)))]))
-           (node-key root)
-           (cond [(eq? default not-given)
-                  (error who "no key found ~a than~a ~e"
-                         (if (memq '< cmp-result) "less" "greater")
-                         (if (memq '= cmp-result) " or equal to" "")
-                         key)]
-                 [(procedure? default) (default)]
-                 [else default])))]))
 
-(define (splay-tree-greatest-key/<=? s key [default not-given])
+(define (splay-tree-iterate-greatest/<=? s key)
-  (extreme 'splay-tree-greatest-key/<=? s key '(< =) has-prev? find-prev default))
+  (extreme 'splay-tree-iterate-greatest/<=? s key '(< =) has-prev? find-prev))
-
+(define (splay-tree-iterate-greatest/<? s key)
-(define (splay-tree-greatest-key/<? s key [default not-given])
+  (extreme 'splay-tree-iterate-greatest/<? s key '(<) has-prev? find-prev))
-  (extreme 'splay-tree-greatest-key/<? s key '(<) has-prev? find-prev default))
+(define (splay-tree-iterate-least/>=? s key)
-
+  (extreme 'splay-tree-iterate-least/>=? s key '(> =) has-next? find-next))
-(define (splay-tree-least-key/>=? s key [default not-given])
+(define (splay-tree-iterate-least/>? s key)
-  (extreme 'splay-tree-least-key/>=? s key '(> =) has-next? find-next default))
+  (extreme 'splay-tree-iterate-least/>? s key '(>) has-next? find-next))
-
-(define (splay-tree-least-key/>? s key [default not-given])
-  (extreme 'splay-tree-least-key/>? s key '(>) has-next? find-next default))
 
 ;; ========
 
@@ -499,20 +518,56 @@ Options
 ;; ========
 
 (provide/contract
- [make-splay-tree (-> (-> any/c any/c any/c) (-> any/c any/c any/c) splay-tree?)]
+ [make-splay-tree
- [make-integer-splay-tree (-> splay-tree?)]
+  (->* ((-> any/c any/c any/c) (-> any/c any/c any/c))
- [splay-tree? (-> any/c boolean?)]
+       (#:key-contract contract? #:value-contract contract?)
- [splay-tree-ref (->* (splay-tree? any/c) (any/c) any/c)]
+       splay-tree?)]
- [splay-tree-set! (-> splay-tree? any/c any/c void?)]
+ [make-integer-splay-tree
- [splay-tree-remove! (-> splay-tree? any/c void?)]
+  (->* ()
- [splay-tree-count (-> splay-tree? exact-nonnegative-integer?)]
+       (#:adjust? any/c #:key-contract contract? #:value-contract contract?)
- [splay-tree->list (-> splay-tree? (listof (cons/c any/c any/c)))]
+       splay-tree?)]
 
- [splay-tree-greatest-key/<=?
+ [splay-tree? (-> any/c boolean?)]
-  (->* (splay-tree? any/c) (any/c) any/c)]
+ [splay-tree-with-adjust? (-> splay-tree? boolean?)]
- [splay-tree-greatest-key/<?
+
-  (->* (splay-tree? any/c) (any/c) any/c)]
+ [splay-tree-ref
- [splay-tree-least-key/>=?
+  (->i ([s splay-tree?] [key (s) (key-c s)])
-  (->* (splay-tree? any/c) (any/c) any/c)]
+       ([default any/c])
- [splay-tree-least-key/>?
+       [_ (s default) (or/c (key-c s) (lambda (x) (eq? x default)))])]
-  (->* (splay-tree? any/c) (any/c) any/c)])
+ [splay-tree-set!
+  (->i ([s splay-tree?] [key (s) (key-c s)] [v (s) (val-c s)]) [_ void?])]
+ [splay-tree-remove!
+  (->i ([s splay-tree?] [key (s) (key-c s)]) [_ void?])]
+ [splay-tree-remove-range!
+  (->i ([s splay-tree?] [from (s) (key-c s)] [to (s) (key-c s)]) [_ void?])]
+ [splay-tree-count
+  (-> splay-tree? exact-nonnegative-integer?)]
+ [splay-tree->list
+  (->i ([s splay-tree?]) [_ (s) (listof (cons/c (key-c s) (val-c s)))])]
+
+ [splay-tree-contract!
+  (->i ([s (and/c splay-tree? splay-tree-with-adjust?)]
+        [from (s) (key-c s)] [to (s) (key-c s)])
+       [_ void?])]
+ [splay-tree-expand!
+  (->i ([s (and/c splay-tree? splay-tree-with-adjust?)]
+        [from (s) (key-c s)] [to (s) (key-c s)])
+       [_ void?])]
+
+ [splay-tree-iterate-first
+  (-> splay-tree? (or/c splay-tree-iter? #f))]
+ [splay-tree-iterate-next
+  (-> splay-tree? splay-tree-iter? (or/c splay-tree-iter? #f))]
+ [splay-tree-iterate-key
+  (-> splay-tree? splay-tree-iter? any/c)]
+ [splay-tree-iterate-value
+  (-> splay-tree? splay-tree-iter? any/c)]
+
+ [splay-tree-iterate-greatest/<=?
+  (->i ([s splay-tree?] [k (s) (key-c s)]) [_ (or/c splay-tree-iter? #f)])]
+ [splay-tree-iterate-greatest/<?
+  (->i ([s splay-tree?] [k (s) (key-c s)]) [_ (or/c splay-tree-iter? #f)])]
+ [splay-tree-iterate-least/>=?
+  (->i ([s splay-tree?] [k (s) (key-c s)]) [_ (or/c splay-tree-iter? #f)])]
+ [splay-tree-iterate-least/>?
+  (->i ([s splay-tree?] [k (s) (key-c s)]) [_ (or/c splay-tree-iter? #f)])])