#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)
(define (tree->preorder-seq t)
(make-do-sequence (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))))
(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)))]))
(define (half n)
(arithmetic-shift n -1))
(define (build-tree i f) (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))))])))
(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)))