;; -*- Mode: Irken -*-

(include "lib/core.scm")

(datatype bral
  (:leaf 'a)
  (:node int (bral 'a) (bral 'a)))

(datatype digit
  (:zero)
  (:one (bral 'a)))

(define (bral/empty) '())
(define (bral/empty? ts) (eq? ts '()))

(define bral->size
  (bral:leaf x)     -> 1
  (bral:node w _ _) -> w
  )

(define (bral/link t1 t2)
  (bral:node (+ (bral->size t1) (bral->size t2)) t1 t2))

(define bral/cons-tree
  x ()                    -> (list:cons (digit:one x) '())
  x ((digit:zero) . ts)   -> (list:cons (digit:one x) ts)
  x ((digit:one t0) . ts) -> (list:cons (digit:zero) (bral/cons-tree (bral/link x t0) ts))
  )

(define bral/uncons-tree
  ()                   -> (raise (:empty))
  ((digit:one t))      -> (:pair t '())
  ((digit:one t) . ts) -> (:pair t (list:cons (digit:zero) ts))
  ((digit:zero) . ts)
  -> (match (bral/uncons-tree ts) with
       (:pair (bral:node _ t1 t2) ts0) -> (:pair t1 (list:cons (digit:one t2) ts0))
       _ -> (raise (:impossible-pattern "bral/uncons-tree")))
  )

(define (bral/cons x ts)
  (bral/cons-tree (bral:leaf x) ts))

(define (bral/head ts)
  (match (bral/uncons-tree ts) with
    (:pair (bral:leaf x) _) -> x
    _ -> (raise (:impossible-pattern "head"))))

(define (bral/tail ts)
  (pair->second (bral/uncons-tree ts)))

(define bral/lookup-tree
  0 (bral:leaf x) -> x
  i (bral:leaf x) -> (raise (:subscript))
  i (bral:node w t1 t2)
  -> (if (< i (/ w 2))
         (bral/lookup-tree i t1)
         (bral/lookup-tree (- i (/ w 2)) t2)))

(define (bral/update-tree i y t)
  (match i t with
    0 (bral:leaf x) -> (bral:leaf y)
    i (bral:leaf x) -> (raise (:subscript))
    i (bral:node w t1 t2)
    -> (if (< i (/ w 2))
           (bral:node w (bral/update-tree i y t1) t2)
           (bral:node w t1 (bral/update-tree (- i (/ w 2)) y t2)))
    ))

(define bral/lookup
  i () -> (raise (:subscript))
  i ((digit:zero) . ts) -> (bral/lookup i ts)
  i ((digit:one t) . ts)
  -> (if (< i (bral->size t))
         (bral/lookup-tree i t)
         (bral/lookup (- i (bral->size t)) ts)))

(define (bral/update i y t)
  (match t with
    () -> (raise (:subscript))
    ((digit:zero) . ts) -> (list:cons (digit:zero) (bral/update i y ts))
    ((digit:one t) . ts)
    -> (if (< i (bral->size t))
           (list:cons (digit:one (bral/update-tree i y t)) ts)
           (list:cons (digit:one t) (bral/update (- i (bral->size t)) y ts)))
    ))

(define bral/iterate-tree
  p (bral:leaf x) -> (p x)
  p (bral:node _ l r)
  -> (begin
       (bral/iterate-tree p l)
       (bral/iterate-tree p r)))

(define bral/iterate
  p () -> #u
  p ((digit:zero) . tl) -> (bral/iterate p tl)
  p ((digit:one t) . tl) -> (begin (bral/iterate-tree p t) (bral/iterate p tl))
  )

(let ((t0 (bral/empty)))
  (for-range
      i 10 ;; i.e. (9 8 7 ... 0)
      (set! t0 (bral/cons i t0)))
  (printn t0)
  (printn (bral/lookup 5 t0))
  (printn (bral/lookup 9 t0))
  (printn (bral/lookup 1 t0))
  (set! t0 (bral/update 4 19 t0))
  (printn t0)
  (printn (bral/head t0))
  (printn (bral/tail t0))
  (printn (bral/cons 999 (bral/tail t0)))
  (bral/iterate
   (lambda (x) (print x) (print-string " "))
   t0)
  (newline)
  )