fib-join-forest.ss
(module fib-join-forest mzscheme
(require (lib "etc.ss")
(lib "list.ss")
(lib "contract.ss")
(planet "contract-utils.ss" ("cobbe" "contract-utils.plt" 3 0)))

;; history:
;; dyoo: I ripped out this code from the rope.plt library; the new version of
;; rope.plt will delegate off to this package.

;; The code below roughly follows the node length-balancing strategy described
;; in the Ropes paper by Hans-Juergen Boehm, Russell R. Atkinson,
;; and Michael F. Plass.

;; join-forest: (listof X) (X X -> X) (X -> number) -> X
;; Joins all of the elements together, roughly keeping the nodes balanced.
(define (join-forest a-forest node-join-f node-weight-f)
(concatenate-forest (foldl (lambda (a-node a-forest)
a-forest
node-join-f
node-weight-f))
'()
a-forest)
node-join-f))

;; add-node-to-forest: X (listof X) (X X -> X) (X -> number) -> (listof X)
;; Adds a new node to our intermediate forest.  The invariant we maintain is
;; that the forest is in increasing weight order.
(define (add-node-to-forest a-node a-forest node-join-f node-weight-f)
(cond
[(empty? a-forest)
(list a-node)]
[(< (node-weight-f a-node)
(node-weight-f (first a-forest)))
(cons a-node a-forest)]
[else
(local
((define partial-forest
(merge-smaller-children a-forest
(node-weight-f a-node)
node-join-f
node-weight-f)))
(restore-forest-order (cons (node-join-f (first partial-forest)
a-node)
(rest partial-forest))
node-join-f
node-weight-f))]))

;; concatenate-first-two: (listof X) (X X -> X) -> (listof X)
;; Joins the first two nodes in the tree together.
(define (concatenate-first-two a-forest node-join-f)
(cons (node-join-f (second a-forest)
(first a-forest))
(rest (rest a-forest))))

;; merge-smaller-children: (listof X) number (X X -> X) (X -> number) -> (listof X)
;; Given a forest, merges the forest iteratively until the weight of the first element
;; in the forest is no larger than n.
(define (merge-smaller-children a-forest n node-join-f node-weight-f)
(cond
[(empty? (rest a-forest))
a-forest]
[(<= (node-weight-f (first a-forest)) n)
a-forest]
[else
(merge-smaller-children (concatenate-first-two a-forest node-join-f)
n
node-join-f
node-weight-f)]))

;; restore-forest-order: (listof X) (X X -> X) (X -> number) -> (listof X)
;; Ensures that the forest is ordered in increasing weight, with the precondition
;; that at worst the first two elements violate this property.
(define (restore-forest-order a-forest node-join-f node-weight-f)
(cond
[(empty? (rest a-forest))
a-forest]
[(>= (node-weight-f (first a-forest))
(node-weight-f (second a-forest)))
(restore-forest-order (concatenate-first-two a-forest node-join-f)
node-join-f
node-weight-f)]
[else
a-forest]))

;; concatenate-forest: (listof X) (X X -> X) -> X
;; Joins all of the forest elements together.
(define (concatenate-forest a-forest node-join-f)
(cond
[(empty? (rest a-forest))
(first a-forest)]
[else
(concatenate-forest (concatenate-first-two a-forest node-join-f)
node-join-f)]))

(provide/contract [join-forest
((nelistof/c any/c)
(any/c any/c . -> . any)
(any/c . -> . natural-number/c)
. -> . any)]))