#lang scheme
(require (planet schematics/schemeunit))
         
(provide (struct-out tree)
         (struct-out leaf)
         (struct-out node)
         tree-map tree-map/n
         build-tree
         tr:make-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)))]
        [else      (make-node (f (tree-val t))
                              (tree-map f (node-left t))
                              (tree-map f (node-right t)))]))

;; [X Y Z ... -> R] [List [Tree X] [Tree Y] [Tree Z] ...] -> [Tree R]
(define (tree-map/n f ts)
  (let recr ((ts ts))
    (cond [(leaf? (car ts)) (make-leaf (apply f (map tree-val ts)))]
          [else (make-node (apply f (map tree-val ts))
                           (recr (map node-left ts))
                           (recr (map node-right ts)))])))

;; 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 (f o)
                      (rec (half i) (+ o 1))
                      (rec (half i) (+ o 1 (half i))))])))

;; Nat X -> [Tree X]
(define (tr:make-tree i x) ;; i = 2^j-1
  (let rec ((i i))
    (cond [(= 1 i) (make-leaf x)]
          [else
           (make-node x
                      (rec (half i))
                      (rec (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? (tree-map/n cons
                            (list (make-leaf 'a)
                                  (make-leaf 'z)))
                (make-leaf '(a . z)))
  
  (check-equal? (tree-map/n cons
                            (list (make-node 'a 
                                             (make-leaf 'b) 
                                             (make-leaf 'c))
                                  (make-node 'z 
                                             (make-leaf 'y) 
                                             (make-leaf 'x))))
                (make-node '(a . z)
                           (make-leaf '(b . y))
                           (make-leaf '(c . x))))
  
  
  
  (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? (build-tree 1 (lambda (i) 'x))
                (make-leaf 'x))
  (check-equal? (build-tree 3 (lambda (i) 'x))
                (make-node 'x
                           (make-leaf 'x)
                           (make-leaf 'x)))
  
  (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)))