perf(powers): O(1) rbtree-size via a cached size field in a wrapper record
ober
ff079fa47100eb7c9719aaeb6e26a0d389be55ea
--- a/Makefile +++ b/Makefile @@ -20,3 +20,4 @@ build: test: build gxi tests/stm-engine-test.ss + gxi tests/powers-rbtree-test.ss --- a/powers.ss +++ b/powers.ss @@ -74,27 +74,59 @@ (def write-unlock! rwlock-write-unlock!) ;; ---- rbtree (Gerbil: put/get/remove; functional variants have no !) ---- -(def rbtree-insert rbtree-put) ; (rbtree-insert t k v) -> new tree -(def rbtree-lookup rbtree-get) ; (rbtree-lookup t k) -> value | #f -(def rbtree-delete rbtree-remove) +;; Gerbil's stdlib rbtree is fully opaque: it exposes no root/node access and no +;; size field, so the count cannot be read off the tree in O(1). The jerboa-named +;; adapters therefore thread a small wrapper record that caches the element count +;; and is updated on insert/delete, making rbtree-size O(1). Every adapter accepts +;; either a wrapper or a bare Gerbil rbtree (rbtree-unwrap), so a caller that +;; passes a bare rbtree still works — rbtree-size just falls back to an O(n) walk +;; for that value. Callers that build the tree through rbtree-insert/rbtree-delete +;; get the O(1) size automatically. +(defstruct powers-rbtree (tree size)) + +(def (rbtree-unwrap t) + (if (powers-rbtree? t) (powers-rbtree-tree t) t)) + +;; O(1) for a wrapper (cached); O(n) walk for a bare Gerbil rbtree. +(def (rbtree-count t) + (if (powers-rbtree? t) + (powers-rbtree-size t) + (let (n 0) (rbtree-for-each (lambda (k v) (set! n (+ n 1))) t) n))) + +(def (rbtree-insert t k v) ; (rbtree-insert t k v) -> new tree + (let* ((tree (rbtree-unwrap t)) + (new? (not (rbtree-contains? tree k))) + (size (+ (rbtree-count t) (if new? 1 0)))) + (make-powers-rbtree (rbtree-put tree k v) size))) + +(def (rbtree-delete t k) + (let* ((tree (rbtree-unwrap t)) + (old? (rbtree-contains? tree k)) + (size (- (rbtree-count t) (if old? 1 0)))) + (make-powers-rbtree (rbtree-remove tree k) size))) + +(def (rbtree-lookup t k) (rbtree-get (rbtree-unwrap t) k)) ; -> value | #f + ;; contains? walks the tree via rbtree-ref (O(log n)) instead of scanning ;; rbtree->list (O(n)); a sentinel distinguishes an absent key from a #f value. (def +rbtree-absent+ (list 'rbtree-absent)) (def (rbtree-contains? t k) - (not (eq? (rbtree-ref t k +rbtree-absent+) +rbtree-absent+))) -;; size/min/max avoid materializing rbtree->list: a single allocation-free -;; for-each pass (for-each visits keys in ascending order, so the first key seen -;; is min and the last is max). Gerbil's rbtree is opaque (no root/node access, -;; no size field), so these stay O(n) time but no longer allocate the full list. -(def (rbtree-size t) - (let (n 0) (rbtree-for-each (lambda (k v) (set! n (+ n 1))) t) n)) + (not (eq? (rbtree-ref (rbtree-unwrap t) k +rbtree-absent+) +rbtree-absent+))) + +(def (rbtree-size t) (rbtree-count t)) + +;; min/max stay O(n): for-each visits keys in ascending order, so the first key +;; seen is min and the last is max. With an opaque tree, keeping a cached min/max +;; correct across deletion of the current extreme would require a full walk on +;; that deletion anyway, so we walk on demand rather than risk a stale cache. (def (rbtree-min t) (let ((k* #f) (v* #f) (any? #f)) (rbtree-for-each (lambda (k v) (unless any? (set! k* k) (set! v* v) (set! any? #t))) - t) + (rbtree-unwrap t)) (and any? (cons k* v*)))) (def (rbtree-max t) (let ((k* #f) (v* #f) (any? #f)) - (rbtree-for-each (lambda (k v) (set! k* k) (set! v* v) (set! any? #t)) t) + (rbtree-for-each (lambda (k v) (set! k* k) (set! v* v) (set! any? #t)) + (rbtree-unwrap t)) (and any? (cons k* v*)))) new file mode 100644 --- /dev/null +++ b/tests/powers-rbtree-test.ss @@ -0,0 +1,63 @@ +;;; -*- Gerbil -*- + +(import :std/sugar + :std/misc/rbtree + :jerboa-compat/powers) + +(def failures 0) + +(defrule (check! name assertion) + (if assertion + (displayln "ok: " name) + (begin + (set! failures (+ failures 1)) + (displayln "FAIL: " name)))) + +(def (cmp a b) (if (< a b) -1 (if (> a b) 1 0))) + +;; Build through the jerboa-named adapters so the size cache is exercised. +(def t0 (make-rbtree cmp)) +(def t1 (rbtree-insert t0 5 'e)) +(def t2 (rbtree-insert t1 2 'b)) +(def t3 (rbtree-insert t2 8 'h)) +(def t4 (rbtree-insert t3 1 'a)) + +(check! "size is correct after adds" (= (rbtree-size t4) 4)) +(check! "re-inserting an existing key does not change size" + (= (rbtree-size (rbtree-insert t4 2 'B)) 4)) +(check! "lookup reads the latest value for a re-inserted key" + (eq? (rbtree-lookup (rbtree-insert t4 2 'B) 2) 'B)) +(check! "contains? finds a present key" (rbtree-contains? t4 8)) +(check! "contains? rejects an absent key" (not (rbtree-contains? t4 99))) +(check! "min after adds" (equal? (rbtree-min t4) (cons 1 'a))) +(check! "max after adds" (equal? (rbtree-max t4) (cons 8 'h))) + +;; Removing the current min/max must yield the next extreme. +(def t-no-min (rbtree-delete t4 1)) +(def t-no-max (rbtree-delete t4 8)) +(check! "size decrements on remove of a present key" + (= (rbtree-size t-no-min) 3)) +(check! "size unchanged on remove of an absent key" + (= (rbtree-size (rbtree-delete t4 99)) 4)) +(check! "min recomputed after removing the current min" + (equal? (rbtree-min t-no-min) (cons 2 'b))) +(check! "max recomputed after removing the current max" + (equal? (rbtree-max t-no-max) (cons 5 'e))) +(check! "lookup of a removed key is #f" + (eq? (rbtree-lookup t-no-min 1) #f)) + +;; Empty tree edge cases. +(check! "size of empty tree is 0" (= (rbtree-size t0) 0)) +(check! "min of empty tree is #f" (not (rbtree-min t0))) +(check! "max of empty tree is #f" (not (rbtree-max t0))) + +;; A bare Gerbil rbtree (not built through the adapters) still works via the +;; O(n) fallback path. +(def bare (rbtree-put (rbtree-put (make-rbtree cmp) 7 'g) 3 'c)) +(check! "size falls back to a walk for a bare rbtree" + (= (rbtree-size bare) 2)) +(check! "min works on a bare rbtree" (equal? (rbtree-min bare) (cons 3 'c))) +(check! "max works on a bare rbtree" (equal? (rbtree-max bare) (cons 7 'g))) + +(displayln "powers-rbtree-tests: ok") +(exit (if (> failures 0) 1 0))