(in-package :binary-trees)

(defun tree-successor (tree node)
  (cond
    ((null-node-p (right node) tree)
     (do ((node node candidate)
          (candidate (parent node) (parent candidate)))
         ((or (null-node-p candidate tree)
              (not (eq node (right candidate))))
          candidate)))
    (t (minimum-node tree (right node)))))

(defun tree-predecessor (tree node)
  (cond
    ((null-node-p (left node) tree)
     (do ((node node candidate)
          (candidate (parent node) (parent candidate)))
         ((or (null-node-p candidate tree)
              (not (eq node (left candidate))))
          candidate)))
    (t (maximum-node tree (left node)))))

(defun tree-for-each (func tree)
  (do ((node (minimum-node tree (root-node tree))
             (tree-successor tree node)))
      ((null-node-p node tree) (values))
    (funcall func (datum node))))

(defun reverse-tree-for-each (func tree)
  (do ((node (maximum-node tree (root-node tree))
             (tree-predecessor tree node)))
      ((null-node-p node tree) (values))
    (funcall func (datum node))))

#||
seemingly useful functions (?)
(defun preorder-tree-for-each (func tree)
  (labels ((traverse-node (node)
             (unless (null-node-p node tree)
               (funcall func (datum node))
               (traverse-node (left node))
               (traverse-node (right node)))))
    (traverse-node (root-node tree))))

(defun postorder-tree-for-each (func tree)
  (labels ((traverse-node (node)
             (unless (null-node-p node tree)
               (traverse-node (left node))
               (traverse-node (right node))
               (funcall func (datum node)))))
    (traverse-node (root-node tree))))

(defun reverse-preorder-tree-for-each (func tree)
  (labels ((traverse-node (node)
             (unless (null-node-p node tree)
               (funcall func (datum node))
               (traverse-node (right node))
               (traverse-node (left node)))))
    (traverse-node (root-node tree))))

(defun reverse-postorder-tree-for-each (func tree)
  (labels ((traverse-node (node)
             (unless (null-node-p node tree)
               (traverse-node (right node))
               (traverse-node (left node))
               (funcall func (datum node)))))
    (traverse-node (root-node tree))))
||#

(defmacro dotree ((obj-var tree-var &optional return-value) &body body)
  (let ((nodesym (gensym))
        (treesym (gensym)))
  `(let ((,treesym ,tree-var))
    (do ((,nodesym (minimum-node ,treesym (root-node ,treesym))
         (tree-successor ,treesym ,nodesym)))
    ((null-node-p ,nodesym ,treesym) ,return-value)
    (let ((,obj-var (datum ,nodesym)))
      ,@body)))))

(defmacro do-tree-range ((obj-var tree-var
                                  &key (type :key)
                                  (lower nil)
                                  (upper nil)) &body body)
  (macrolet ((invalid-type (type)
               `(error "Invalid :type supplied to DO-TREE-RANGE: ~A" ,type)))
    (let ((nodesym (gensym))
          (lower-exp (if lower
                         (cond
                           ((eq type :key)
                            `(upper-bound-node ,lower ,tree-var))
                           ((eq type :index)
                            `(select-node ,tree-var ,lower))
                           (t (invalid-type type)))
                         (cond
                           ((or (eq type :key) (eq type :index))
                            `(minimum-node ,tree-var (root-node ,tree-var)))
                           (t (invalid-type type)))))
          (upper-exp (if upper
                         (cond
                           ((eq type :key)
                            `(upper-bound-node ,upper ,tree-var))
                           ((eq type :index)
                            `(select-node ,tree-var ,upper))
                           (t (invalid-type type)))
                         (cond
                           ((or (eq type :key) (eq type :index))
                            `(sentinel-node ,tree-var))
                           (t (invalid-type type))))))
      `(do ((,nodesym ,lower-exp
             (tree-successor ,tree-var ,nodesym)))
        ((eq ,nodesym ,upper-exp) nil)
        (let ((,obj-var (datum ,nodesym)))
          ,@body)))))

;;; FIXME: FROM-END isn't necessarily very intuitive here.  find out
;;; how regular CL sequence functions treat it (especially with indices)
;;; and rewrite the macro to match.
(defmacro with-tree-iterator ((iter tree &key
                                    (from-end nil) (type :key) (start nil))
                              &body body)
  (let ((nodesym (gensym))
        (treesym (gensym)))
    `(let ((,treesym ,tree))
      (macrolet ((,iter ()
                   (values (datum ,nodesym)
                           (not (null-node-p ,nodesym ,treesym)))))
        (do ((,nodesym ,(cond
                         ((eq type :key)
                          (if start
                              `(lower-bound-node ,start ,treesym)
                              `(minimum-node ,treesym (root-node ,treesym))))
                         ((eq type :index)
                          (if start
                              `(select-node ,start ,treesym)
                              `(minimum-node ,treesym (root-node ,treesym)))))
                       (,(if from-end
                             'tree-predecessor
                             'tree-successor) ,treesym ,nodesym)))
            ()
          ,@body)))))

(defun reduce (tree function
                    &key (key #'identity)
                    (initial-value nil valuep)
                    (from-end nil))
  (let ((accum (if valuep
                   initial-value
                   (funcall function))))
    (flet ((reducer (object)
             (setf accum (funcall function accum (funcall key object)))))
      (declare (dynamic-extent #'reducer))
      (if from-end
          (reverse-tree-for-each #'reducer tree)
          (tree-for-each #'reducer tree))
      accum)))

(defun node-position (node tree from-end)
  (if (null-node-p node tree)
      nil                             ; same as CL:POSITION
      (do ((node node (parent node))
           (position (rank node)))
          ((eq node (root-node tree))
           (if from-end
               (- (size tree) (1- position))
               (1- position)))
        (unless (eq node (left (parent node)))
          (incf position (rank (parent node)))))))

(defun position (key tree &key from-end)
  (let ((node (find-node-with-key tree key)))
    (node-position node tree from-end)))