;; #!/usr/bin/csi
;; (C) 2008, 2010 Jörg F. Wittenberger.

;; Redistribution permitted under either GPL, LGPL or BSD style
;; license.

;; (use extras)

;;* Left Leaning Red Black Tree

;;** Code Generator

;; The code generating macro expander, written in procedural style.

(define-syntax (define-llrbtree-code x r c)

(define %define (r 'define))
(define %begin (r 'begin))
(define %lambda (r 'lambda))
(define %let (r 'let))
(define %let* (r 'let*))
(define %eq? (r 'eq?))
(define %not (r 'not))
(define %and (r 'and))
(define %or (r 'or))
(define %if (r 'if))
(define %cond (r 'cond))
(define %else (r 'else))
(define %set! (r 'set!))
(define %vector (r 'vector))
(define %vector-ref (r 'vector-ref))
(define %vector-set! (r 'vector-set!))

;; -------------------------------------------

(define (make-llrbtree-code
	 ;; The "features" is a list of symbols to control code
	 ;; expansion.  "pure" will use "update" and "update+",
	 ;; otherwise only "update!" will be used.  "ordered" will
	 ;; enforce total order among the element.  "debug" will
	 ;; expand a simple tree printer, and "leftmost" will include
	 ;; untested code to maintain a leftmost value of the tree.
	 features
	 ;; The "update*" expressions are lambda abstractions (sans
	 ;; the 'lambda' keyword) evaluated at compile time to produce
	 ;; the actual code to update a node.  These procedures take
	 ;; 1+ arguments.  A original node and a keyword: <value>
	 ;; ... list of desired updates.  Possible keywords are left:
	 ;; right: and color:

	 ;; "update" : If feature "pure" is set, "update" must expand
	 ;; to a newly allocated node, otherwise is should expand to a
	 ;; side effect full update of the original node.
	 update
	 init-root-node!		;; defined
	 t-lookup			;; defined
	 t-min				;; defined
	 t-fold				;; defined
	 t-for-each			;; defined
	 t-insert			;; defined
	 t-delete			;; defined
	 t-delete-min			;; defined
	 t-empty?			;; defined

	 ;; These procedures expand to code for comparision
	 ;; expressions.
	 t-k-eq?			;; key<>node-key "equal"
	 t-k-<?				;; key<>node-key "less then"
	 t-<?				;; node<>node "less then"

	 left set-left!
	 right set-right!
	 color set-color!
      ;;; This is an experiment too.  But since it adds non-constant
      ;;; complexity to the code, I recommend to pass #f here.  It's
      ;;; also not really tested.
	 set-leftmost!
	 )

  (define maintain-leftmost! (memq 'leftmost features))

  (define pure (memq 'pure features))

  (define ordered (memq 'ordered features))

  (define use-root-pointer (memq 'use-root-pointer features))

  (define root-node left)

  (define (with-n-node t node . steps)
    `(,%let ((n.n ,node))
       (,%let ((n.l (,left n.n))
	     (n.r (,right n.n))
	     (n.c (,color n.n)))
	 ,(let loop ((steps steps))
	    (if (null? steps)
		`(,%if ,(empty? t 'n.n) n.n
		     ,(update 'n.n left: 'n.l right: 'n.r color: 'n.c))
		`(,%begin
		   ,((car steps) t 'n.n 'n.l 'n.r 'n.c)
		   ,(loop (cdr steps))))))))

  (define empty?
    (if (or pure (not use-root-pointer))
	(lambda (t node) `(,%not ,node))
	(lambda (t node) `(,%eq? ,t ,node))))

  (define empty
    (if (or pure (not use-root-pointer))
	(lambda (t) #f)
	(lambda (t) t)))

  (define black
    (if (or pure (not use-root-pointer))
	(lambda (t) #t)
	(lambda (t) t)))

  (define (red)
    #f)

  (define (red? t)
    `(,%lambda (node)
       (,%if ,(empty? t 'node) #f (,%not (,color node)))))

  (define (ptred? t r sel)
    `(,%if ,(empty? t r) #f
	 (,(red? t) (,sel ,r))))

  (define (black? t)
    `(,%lambda (node)
       (,color node)))

  (define (color-black? t)
    (lambda (c) c))

  (define (color-flip-node! t n)
    `(,%if ,(empty? t n) ,n
	 ,(update n color: `(,%if (,(black? t) ,n) ,(red) ,(black t)))))

  (define (color-flip! t n.n n.l n.r n.c)
    `(,%if (,%not ,(empty? t n.n))
	 (,%begin
	   (,%set! ,n.l ,(color-flip-node! 't n.l))
	   (,%set! ,n.r ,(color-flip-node! 't n.r))
	   (,%set! ,n.c (,%if ,((color-black? t) n.c) ,(red) ,(black t))))))

  (define (rotate-left! t n.n n.l n.r n.c)
    `(,%begin
       (,%set! ,n.l ,(update n.n left: n.l right: `(,left ,n.r) color: (red)))
       (,%set! ,n.n ,n.r)
       (,%set! ,n.r (,right ,n.r))))

  (define (rotate-right! t n.n n.l n.r n.c)
    `(,%begin
       (,%set! ,n.r ,(update n.n left: `(,right ,n.l) right: n.r color: (red)))
       (,%set! ,n.n ,n.l)
       (,%set! ,n.l (,left ,n.l))))

  (define (fixup! t n.n n.l n.r n.c)
    `(,%begin
       (,%if (,(red? t) ,n.r)
	   ,(rotate-left! t n.n n.l n.r n.c))
       (,%if (,%and (,(red? t) ,n.l) ,(ptred? t n.l left))
	   ,(rotate-right! t n.n n.l n.r n.c))
       (,%if (,%and (,(red? t) ,n.l) (,(red? t) ,n.r))
	   ,(color-flip! t n.n n.l n.r n.c))))

  (define (move-red-right! t n.n n.l n.r n.c)
    `(,%begin
       ,(color-flip! t n.n n.l n.r n.c)
       (,%if ,(ptred? t n.l left)
	   (,%begin
	     ,(rotate-right! t n.n n.l n.r n.c)
	     ,(color-flip! t n.n n.l n.r n.c)))))

  (define (move-red-left! t n.n n.l n.r n.c)
    `(,%begin
       ,(color-flip! t n.n n.l n.r n.c)
       (,%if ,(ptred? t n.r left)
	   (,%begin
	     (,%set! ,n.r ,(with-n-node t n.r rotate-right!))
	     ,(rotate-left! t n.n n.l n.r n.c)
	     ,(color-flip! t n.n n.l n.r n.c)))))

  (define (define-delete-min t)
    `(,%define (delete-min set-leftmost! r n)
       (,%if ,(empty? t `(,left n))
	   (,%begin
	     (,%vector-set! r 0 n)
	     (,left n))
	   ,(with-n-node
	     t 'n
	     (lambda (t n.n n.l n.r n.c)
	       `(,%begin
		  (,%if (,%and (,%not (,(red? t) ,n.l))
			   (,%not ,(ptred? t n.l left)))
		      ,(move-red-left! t n.n n.l n.r n.c))
		  (,%set! ,n.l (delete-min set-leftmost! r ,n.l))
		  ,@(if maintain-leftmost!
			`((,%if (,%and set-leftmost! ,(empty? t n.l))
			      (set-leftmost! n)))
			'())
		  ,(fixup! t n.n n.l n.r n.c)))))))

  `(,%begin
     ,@(if init-root-node!
	   `((,%define (,init-root-node! t)
	       ,(update 't color: (black 't) left: (empty 't))))
	   '())
     ,@(if t-empty?
	   `((,%define (,t-empty? t)
	       ,(empty? 't `(,root-node t))))
	   '())
     ,@(if t-lookup
	   `((,%define (,t-lookup t k)
	       (,%let loop ((node (,root-node t)))
		 (,%cond
		  (,(empty? 't 'node) node)
		  (,(t-k-eq? 'k 'node) node)
		  (,(t-k-<? 'k 'node) (loop (,left node)))
		  (,%else (loop (,right node)))))))
	   '())
     ,@(if t-min
	   `((,%define (,t-min t)
	       (,%if ,(empty? 't `(,root-node t)) #f
		   (,%let loop ((node (,root-node t)))
		     (,%cond
		      (,(empty? 't `(,left node)) node)
		      (,%else (loop (,left node))))))))
	   '())
     ,@(if t-fold
	   `((,%define (,t-fold procedure init t)
	       (,%define (fold init node)
		 (,%if ,(empty? 't 'node)
		     init
		     (fold (procedure node (fold init (,right node))) (,left node))))
	       (fold init (,root-node t))))
	   '())
     ,@(if t-for-each
	   `((,%define (,t-for-each procedure t)
	       (,%let loop ((node (,root-node t)))
		 (,%or ,(empty? 't 'node)
		     (,%begin
		       (procedure node)
		       (loop (,left node))
		       (loop (,right node)))))))
	   '())
     ,@(if t-insert
	   `((,%define (,t-insert t ,@(if ordered '(k) '()) n . set-leftmost!)
	       ,@(if pure '()
		     (list (update 'n color: (red) left: (empty 't) right: (empty 't))))
	       (,%let ((nr
		      (,%let loop ((node (,root-node t))
				 (sl (,%and (pair? set-leftmost!) (car set-leftmost!))))
			(,%if ,(empty? 't 'node)
			    (,%if sl (,%begin (sl n) n) n)
			    ,(with-n-node
			      't 'node
			      (lambda (t n.n n.l n.r n.c)
				`(,%begin
				   ,(if ordered
					`(,%if ,(t-k-eq? 'k n.n)
					     (,%set! ,n.n
						   ,(update 'n left: n.l right: n.r
							    color: n.c))
					     (,%if ,(t-k-<? 'k n.n)
						 (,%set! ,n.l (loop ,n.l sl))
						 (,%set! ,n.r (loop ,n.r #f))))
					`(,%if ,(t-<? 'n n.n)
					     (,%set! ,n.l (loop ,n.l sl))
					     (,%set! ,n.r (loop ,n.r #f))))
				   (,%if (,(red? 't) ,n.r)
				       ,(rotate-left! t n.n n.l n.r n.c))
				   (,%if (,%and (,(red? 't) ,n.l)
					    (,(red? 't) (,left ,n.l)))
				       ,(rotate-right! t n.n n.l n.r n.c))
				   (,%if (,%and (,(red? 't) ,n.l) (,(red? 't) ,n.r))
				       ,(color-flip! t n.n n.l n.r n.c)))))))))
		 (,%if (,(red? 't) nr)
		     (,%set! nr ,(update 'nr color: (black 't))))
		 ,(update 't left: 'nr color: (black 't)))))
	   '())
     ,@(if t-delete-min
	   `((,%define (,t-delete-min t . cont)
	       ,(define-delete-min 't)
	       (,%define set-leftmost! (,%and (pair? cont) (pair? (cdr cont)) (cadr cont)))
	       (,%if ,(empty? 't `(,root-node t))
		   ,(if pure
			`(,%if (pair? cont) ((car cont) t #f) t)
			#f)
		   (,%let* ((min (,%vector #f))
			  (r (delete-min set-leftmost! min (,left t))))
		     ,@(if maintain-leftmost!
			   `((,%if (,%and set-leftmost! ,(empty? 't 'r))
				 (set-leftmost! r)))
			   '())
		     ,(if pure
			  `(,%let ((t ,(update 't left: 'r)))
			     (,%if (pair? cont) ((car cont) t (,%vector-ref min 0)) t))
			  `(,%begin
			     ,(update 't left: 'r color: (black 't))
			     (,%vector-ref min 0)))))))
	   '())
     ,@(if t-delete
	   `((,%define (,t-delete t n/k . rest)
	       ,(define-delete-min 't)
	       (,%define (delete! set-leftmost! h)
		 (,%if ,(if ordered
			  (t-k-<? 'n/k 'h)
			  `(,%and (,%not (,%eq? n/k h)) ,(t-<? 'n/k 'h)))
		     ,(with-n-node
		       't 'h
		       (lambda (t n.n n.l n.r n.c)
			 `(,%begin
			    (,%if (,%and (,%not (,(red? 't) ,n.l))
				     (,%not ,(ptred? 't n.l left)))
				,(move-red-left! t n.n n.l n.r n.c))
			    (,%set! ,n.l (,%if ,(empty? 't n.l) ,(empty 't)
					   (delete! set-leftmost! ,n.l)))
			    ,@(if maintain-leftmost!
				  `((,%if (,%and set-leftmost! ,(empty? 't n.l))
					(set-leftmost! ,n.l)))
				  '())
			    ,(fixup! t n.n n.l n.r n.c))))
		     ,(with-n-node
		       't 'h
		       (lambda (t n.n n.l n.r n.c)
			 `(,%begin
			    (,%if (,(red? 't) ,n.l)
				,(rotate-right! t n.n n.l n.r n.c))
			    (,%if (,%and ,(if ordered (t-k-eq? 'n/k n.n) `(,%eq? n/k ,n.n))
				     ,(empty? 't n.r))
				(,%set! ,n.n ,n.r)
				(,%begin
				  (,%if (,%and (,%not (,(red? 't) ,n.r))
					   (,%not ,(ptred? 't n.r left)))
				      ,(move-red-right! t n.n n.l n.r n.c))
				  (,%if ,(if ordered (t-k-eq? 'n/k n.n) `(,%eq? n/k ,n.n))
				      (,%let ((minv (,%vector #f)))
					(,%set! ,n.r (delete-min #f minv ,n.r))
					(,%set! ,n.n (,%vector-ref minv 0)))
				      (,%if (,%not ,(empty? 't n.r))
					  (,%set! ,n.r (delete! #f ,n.r))))
				  ,(fixup! t n.n n.l n.r n.c))))))))
	       (,%define set-leftmost! (,%and (pair? rest) (car rest)))
	       (,%if ,(empty? 't `(,root-node t)) t
		   (,%let ((r (delete! set-leftmost! (,root-node t))))
		     ,(update 't left: 'r color: (black 't))))))
	   '())
     ))


;; -------------------------------------------

  (apply
   (lambda
       (features
	update
	init-root-node!
	t-lookup
	t-min
	t-fold
	t-for-each
	t-insert
	t-delete
	t-delete-min
	t-empty?
	t-k-eq?
	t-k-<?
	t-<?
	left
	set-left!
	right
	set-right!
	color
	set-color!
	set-leftmost!)
     (make-llrbtree-code
      features
      (eval (cons 'lambda update))
      init-root-node!
      t-lookup
      t-min
      t-fold
      t-for-each
      t-insert
      t-delete
      t-delete-min
      t-empty?
      (eval (cons 'lambda t-k-eq?))
      (eval (cons 'lambda t-k-<?))
      (eval (cons 'lambda t-<?))
      left
      set-left!
      right
      set-right!
      color
      set-color!
      set-leftmost!))
      (cdr x)))