;; tree-add1.ss

(define-datatype StackFrame StackFrame?
  [identity-k]
  [left-k (val number?)
          (right box?)]
  [right-k (val number?)
           (left box?)])

(define HEAP-SIZE 24)

(define stack (make-vector 20 (identity-k)))
(define stack-ptr 0)

(define heap (make-vector HEAP-SIZE 'unused))
(define heap-ptr 0)

(define (alloc/empty-tree)
  (if (< heap-ptr HEAP-SIZE)
      (begin0
        heap-ptr
        (vector-set! heap heap-ptr 'empty-tree)
        (set! heap-ptr (add1 heap-ptr)))
      (error 'alloc/empty-tree "out of memory!")))

(define (alloc/node val lchild rchild)
  (if (< (+ heap-ptr 3) HEAP-SIZE)
      (begin0
        heap-ptr
        (vector-set! heap heap-ptr 'node)
        (vector-set! heap (+ 1 heap-ptr) val)
        (vector-set! heap (+ 2 heap-ptr) lchild)
        (vector-set! heap (+ 3 heap-ptr) rchild)
        (set! heap-ptr (+ 4 heap-ptr)))
      (error 'alloc/node "out of memory!")))

(define (empty-tree? loc)
  (eq? (vector-ref heap loc) 'empty-tree))

(define (node? loc)
  (eq? (vector-ref heap loc) 'node))

(define (node-val loc)
  (vector-ref heap (add1 loc)))

(define (node-left loc)
  (vector-ref heap (+ 2 loc)))

(define (node-right loc)
  (vector-ref heap (+ 3 loc)))

(define empty-tree alloc/empty-tree)
(define node alloc/node)

(define (get-stack-roots)
  (let loop ([i (sub1 stack-ptr)] [roots empty])
    (if (> i 0)
      (cases StackFrame (vector-ref stack i)
        [identity-k () (loop (sub1 i) roots)]
        [left-k (val right) (loop (sub1 i) (cons right roots))]
        [right-k (val left) (loop (sub1 i) (cons left roots))])
      roots)))

(define (push frame)
  (vector-set! stack stack-ptr frame)
  (set! stack-ptr (add1 stack-ptr)))

(define (pop a-tree)
  (set! stack-ptr (sub1 stack-ptr))
  (cases StackFrame (vector-ref stack stack-ptr)
    [identity-k () a-tree]
    [left-k (val right)
            (push (right-k val (box a-tree)))
            (tree-add1/k (unbox right))]
    [right-k (val left)
             (pop
              (node (add1 val) (unbox left) a-tree))]))
            
(define (tree-add1 a-tree)
  (push (identity-k))
  (tree-add1/k a-tree))

(define (tree-add1/k a-tree)
  (cond
    [(empty-tree? a-tree)
     (pop (empty-tree))]
    [(node? a-tree)
     (push (left-k (node-val a-tree)
                   (box (node-right a-tree))))
     (tree-add1/k (node-left a-tree))]))

(define (location->tree loc)
  (cond
    [(empty-tree? loc) 'mt]
    [(node? loc) (list 'node (node-val loc)
                       (location->tree (node-left loc))
                       (location->tree (node-right loc)))]))