;; -*- Mode: Irken -*-

(include "lib/core.scm")

;; http://eternallyconfuzzled.com/tuts/datastructures/jsw_tut_andersson.aspx
;; http://en.wikipedia.org/wiki/AA_tree

(define tree/nil
  {level = 0
   left  = tree/nil
   right = tree/nil
   data  = (magic #u)
   })

;; shame that doesn't just work.
(set! tree/nil.left tree/nil)
(set! tree/nil.right tree/nil)

(define (node/make level data left right)
  { level = level
    data  = data
    left  = left
    right = right })

;; Note: I can't make skew & split tail-recursive unless I
;;   rewrite them to work bottom-up rather than top-down.

(define (tree/skew d)
  (if (and (> d.level 0)
	   (= d.left.level d.level))
      (node/make
       d.level d.left.data d.left.left 
       (node/make
	d.level d.data d.left.right
	(tree/skew d.right)))
      d))

(define (tree/split b)
  (if (and (= b.right.right.level b.level)
	   (not (= b.level 0)))
      (node/make
       (+ 1 b.level) b.right.data
       (node/make b.level b.data b.left b.right.left)
       (tree/split b.right.right))
      b))
       
(define (tree/insert root data <)
  (let loop ((n root))
    (if (= n.level 0)
	(node/make 1 data tree/nil tree/nil)
	(tree/split
	 (tree/skew
	  (if (< data n.data)
	      (node/make n.level n.data (loop n.left) n.right)
	      (node/make n.level n.data n.left (loop n.right)))
	  )))))

;; XXX make this pure.

(define (tree/delete root data data-less? data-equal?)
  (let recur ((root root) (data data))
    (if (not (eq? root tree/nil))
	(if (data-equal? data root.data)
	    (if (and (not (eq? root.left tree/nil))
		     (not (eq? root.right tree/nil)))
		(let loop ((heir root.left))
		  (cond ((not (eq? heir.right tree/nil))
			 (loop heir.right))
			(else
			 (set! root.data heir.data)
			 (set! root.left (recur root.left root.data)))))
		(set! root (if (eq? root.left tree/nil) root.right root.left)))
	    (if (data-less? root.data data)
		(set! root.right (recur root.right data))
		(set! root.left (recur root.left data))
		)))
    (if (or (< root.left.level (- root.level 1))
	    (< root.right.level (- root.level 1)))
	(begin
	  (set! root.level (- root.level 1))
	  (if (> root.right.level root.level)
	      (set! root.right.level root.level))
	  (tree/skew (tree/split root)))
	root
	)))
  
(define indent
  0 -> #u
  n -> (begin (print-string "  ") (indent (- n 1))))

(define (tree/dump t)
  (if (> t.level 0)
      (begin
	(tree/dump t.left)
	(indent t.level)
	(print t.level)
	(print-string " ")
	(printn t.data)
	(tree/dump t.right))))

(define (t0)
  (let ((t tree/nil))
    (for-range
	i 20
	(set! t (tree/insert t i <)))
    (tree/dump t)
    (for-range
	i 20
	(set! t (tree/delete t i < =))
	(print-string "---------------\n")
	(tree/dump t))
    ))

(t0)