(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)
                                 (add-node-to-forest a-node
                                                     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)]))