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

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

;; [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)))]))

;; ---------------------------------------------------------------------------
;; 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? (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)))