private/tree.ss
#lang scheme
(require (planet schematics/schemeunit))
         
(provide (struct-out tree)
         (struct-out leaf)
         (struct-out node)
         tree-map
         build-tree)

(define-struct tree        (val)        #:prefab)
(define-struct (leaf tree) ()           #:prefab)
(define-struct (node tree) (left right) #:prefab)

;; Not currently used, but might come in handy at some point.
;; [Tree X] -> [Seq X]
(define (tree->preorder-seq t)
  (make-do-sequence  ;; position is a forest [Listof [Tree X]]
   (lambda ()
     (values
      (lambda (x) (tree-val (car x)))
      (lambda (p) 
        (cond [(leaf? (car p)) (cdr p)]
              [else
               (cons (node-left (car p))
                     (cons (node-right (car p))
                           (cdr p)))]))
      (list t) 
      cons?
      void
      void))))

;; [X -> Y] [Tree X] -> [Tree Y]
(define (tree-map f t)
  (cond [(leaf? t) (make-leaf (f (tree-val t)))]
        [(node? t) (make-node (f (tree-val t))
                              (tree-map f (node-left t))
                              (tree-map f (node-right t)))]))

;; Consumes n = 2^i-1 and produces 2^(i-1)-1.
;; Nat -> Nat
(define (half n)
  (arithmetic-shift n -1))

;; Nat [Nat -> X] -> [Tree X]
;; like build-list, but for complete binary trees
(define (build-tree i f) ;; i = 2^j-1
  (let rec ((i i) (o 0))
    (cond [(= 1 i) (make-leaf (f o))]
          [else
           (make-node o
                      (rec (half i) (+ o 1))
                      (rec (half i) (+ o 1 (half i))))])))


;; ---------------------------------------------------------------------------
;; Test suite

(define/provide-test-suite tree-tests
  
  (check-equal? (tree-map add1 (make-leaf 0))
                (make-leaf 1))
  (check-equal? (tree-map add1 (make-node 0 (make-leaf 1) (make-leaf 2)))
                (make-node 1 (make-leaf 2) (make-leaf 3)))
  
  (check-equal? (build-tree 1 (lambda (i) i)) (make-leaf 0))
  (check-equal? (build-tree 3 (lambda (i) i)) 
                (make-node 0
                           (make-leaf 1)
                           (make-leaf 2)))
  (check-equal? (build-tree 7 (lambda (i) i))
                (make-node 0
                           (make-node 1
                                      (make-leaf 2)
                                      (make-leaf 3))
                           (make-node 4
                                      (make-leaf 5)
                                      (make-leaf 6))))
  
  (check-equal? (for/list ([i (tree->preorder-seq (make-leaf 0))]) i)
                (list 0))
  
  (check-equal? 
   (for/list ([i (tree->preorder-seq
                  (make-node 0
                             (make-node 1 
                                        (make-leaf 2) 
                                        (make-leaf 3))
                             (make-node 4 
                                        (make-leaf 5) 
                                        (make-leaf 6))))])
             i)
   (list 0 1 2 3 4 5 6)))