+;;; -*-lisp-*-
+;;;
+;;; $Id$
+;;;
+;;; Andersson tree implementation
+;;;
+;;; (c) 2006 Straylight/Edgeware
+;;;
+
+;;;----- Licensing notice ---------------------------------------------------
+;;;
+;;; This program is free software; you can redistribute it and/or modify
+;;; it under the terms of the GNU General Public License as published by
+;;; the Free Software Foundation; either version 2 of the License, or
+;;; (at your option) any later version.
+;;;
+;;; This program is distributed in the hope that it will be useful,
+;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
+;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+;;; GNU General Public License for more details.
+;;;
+;;; You should have received a copy of the GNU General Public License
+;;; along with this program; if not, write to the Free Software Foundation,
+;;; Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
+
+;;;--------------------------------------------------------------------------
+;;; Package.
+
+(defpackage #:aa-tree
+ (:use #:common-lisp #:mdw.base)
+ (:export #:make-aa-tree #:aa-tree-p #:aa-tree-key<
+ #:getaa #:updateaa #:mapaa #:doaa #:aa-tree-iterator #:remaa))
+(in-package #:aa-tree)
+
+;;;--------------------------------------------------------------------------
+;;; The underlying implementation.
+
+(deftype stack-pointer () '(integer 0 255))
+
+(defstruct (tree-node
+ (:conc-name node-)
+ (:type vector)
+ (:constructor make-tree-node
+ (key &optional data level left right)))
+ "Structure representing a node in an Andersson tree."
+ (left nil :type (or null tree-node))
+ (right nil :type (or null tree-node))
+ (level 0 :type stack-pointer)
+ key
+ data)
+
+(deftype tree-node () 'simple-vector)
+
+(defstruct (aa-tree
+ (:predicate treep)
+ (:constructor make-aa-tree
+ (key<-name
+ &aux
+ (key< (functionify key<-name))))
+ (:conc-name tree-))
+ "Structure representing an Andersson tree."
+ (root nil :type (or null tree-node))
+ (stack (make-array 32) :type simple-vector)
+ (key< (slot-uninitialized) :read-only t :type (function (t t) t)))
+
+(declaim (inline skew split))
+
+(defun skew (node)
+ "Implements the `skew' operation on a tree node, eliminating left-pointing
+ internal pointers by applying right-rotation. Returns the replacement
+ node."
+ (declare (type tree-node node))
+ (let ((left (node-left node)))
+ (when (and left (= (node-level node) (node-level left)))
+ (shiftf (node-left node) (node-right left) node left))
+ node))
+
+(defun split (node)
+ "Implements the `split' operation on a tree node, eliminating overly-large
+ pseudo-nodes by applying left-rotation. Returns the replacement node."
+ (declare (type tree-node node))
+ (let* ((right (node-right node))
+ (rright (and right (node-right right))))
+ (when (and rright (= (node-level node) (node-level rright)))
+ (shiftf (node-right node) (node-left right) node right)
+ (incf (node-level node)))
+ node))
+
+(defun get-tree-stack (tree)
+ "Return the current stack for the TREE. This is used to remember the path
+ taken during a search in tree, so we can fix it up afterwards. Keeping
+ just one stack for the tree saves on consing; it's not safe to do
+ simultaneous destructive operations on a tree anyway, so this is a
+ reasonable thing to do. This function ensures that the stack attached to
+ the tree is actually large enough before returning it."
+ (declare (type aa-tree tree))
+ (let* ((root (tree-root tree))
+ (want (* 4 (+ (if root (node-level root) 0) 2)))
+ (stack (tree-stack tree))
+ (size (array-dimension (tree-stack tree) 0)))
+ (if (>= size want)
+ stack
+ (do ((need (ash size 1) (ash need 1)))
+ ((>= need want) (setf (tree-stack tree) (make-array need)))))))
+
+(defun getaa (tree key &optional default)
+ "Look up the given KEY in an Andersson TREE; if the KEY was found, return
+ the corresponding data and t, otherwise return DEFAULT and nil."
+ (declare (type aa-tree tree))
+ (let ((key< (tree-key< tree))
+ (node (tree-root tree))
+ (candidate nil)
+ (candidate-key nil))
+ (declare (type (function (t t) t) key<)
+ (type (or null tree-node) node candidate))
+ (flet ((key< (x y)
+ (funcall key< x y)))
+ (declare (inline key<))
+ (loop (cond (node
+ (let ((node-key (node-key node)))
+ (if (key< key node-key)
+ (setf node (node-left node))
+ (setf candidate node
+ candidate-key node-key
+ node (node-right node)))))
+ ((and candidate (not (key< candidate-key key)))
+ (return (values (node-data candidate) t)))
+ (t
+ (return (values default nil))))))))
+
+(defun tree-probe (tree key)
+ "Do a search in an Andersson TREE for the KEY, returning three values. The
+ second and third are a stack of alternating nodes and direction bits, and
+ a stack pointer (empty, ascending), which together describe the path from
+ the tree root to the successor of the sought-for node. The first is
+ either the sought-for node itself, or nil if it wasn't there."
+ (declare (type aa-tree tree))
+ (let ((key< (tree-key< tree))
+ (stack (get-tree-stack tree))
+ (sp 0)
+ (candidate nil)
+ (candidate-key nil))
+ (declare (type (function (t t) t) key<)
+ (type simple-vector stack)
+ (type stack-pointer sp)
+ (type (or null tree-node) candidate))
+ (flet ((pathpush (v i)
+ (setf (svref stack sp) v
+ (svref stack (1+ sp)) i)
+ (incf sp 2))
+ (key< (x y)
+ (funcall key< x y)))
+ (declare (inline pathpush key<))
+ (let ((node (tree-root tree)))
+ (loop (when (null node)
+ (return))
+ (let* ((node-key (node-key node))
+ (dir (cond ((key< key node-key) 0)
+ (t (setf candidate node
+ candidate-key node-key)
+ 1))))
+ (pathpush node dir)
+ (setf node (svref node dir)))))
+ (values (if (and candidate (not (key< candidate-key key)))
+ candidate
+ nil)
+ stack
+ sp))))
+
+(defun fixup-insert (tree stack sp node)
+ "TREE is an Andersson tree, STACK and SP are the values from a failed call
+ to tree-probe, and NODE is a newly-created node. Insert the NODE into
+ the tree, fix up its balance."
+ (declare (type aa-tree tree)
+ (type simple-vector stack)
+ (type stack-pointer sp)
+ (type tree-node node))
+ (loop (when (zerop sp)
+ (return))
+ (decf sp 2)
+ (let ((parent (svref stack sp))
+ (dir (svref stack (1+ sp))))
+ (setf (svref parent dir) node
+ node parent))
+ (setf node (split (skew node))))
+ (setf (tree-root tree) node))
+
+(defun (setf getaa) (data tree key &optional ignore)
+ "Inserts a new node with the given KEY into an Andersson TREE, if there
+ wasn't one already. Returns two values: the requested node, and either t
+ if the node was inserted, or nil if it was already there."
+ (declare (type aa-tree tree)
+ (ignore ignore))
+ (multiple-value-bind (node stack sp) (tree-probe tree key)
+ (cond (node (setf (node-data node) data))
+ (t (fixup-insert tree stack sp (make-tree-node key data)) data))))
+
+(defun updateaa (tree key func)
+ "Search TREE for an item with the given KEY. If it was found, call FUNC
+ with arguments of the node's data and t, and store its result as the
+ node's new data. If it was absent, call FUNC with arguments nil and nil,
+ and make a new node with the KEY and return value. The FUNC can escape to
+ prevent the node being created (though this is probably not useful)."
+ (declare (type aa-tree tree))
+ (multiple-value-bind (node stack sp) (tree-probe tree key)
+ (cond (node (setf (node-data node) (funcall func (node-data node) t)))
+ (t (let ((data (funcall func nil nil)))
+ (fixup-insert tree stack sp (make-tree-node key data))
+ data)))))
+
+(defun remaa (tree key)
+ "Deletes the node with the given KEY from an Andersson TREE. Returns t if
+ the node was found and deleted, or nil if it wasn't there to begin with."
+ (declare (type aa-tree tree))
+ (multiple-value-bind (candidate stack sp) (tree-probe tree key)
+ (when candidate
+ (decf sp 2)
+ (let ((node (svref stack sp)))
+
+ ;; Unsplice the candidate node from the tree, leaving node as its
+ ;; replacement.
+ (if (eq candidate node)
+ (setf node nil)
+ (setf (node-key candidate) (node-key node)
+ (node-data candidate) (node-data node)
+ node (node-right node)))
+
+ ;; Now wander back up the tree, fixing it as we go.
+ (loop (when (zerop sp)
+ (return))
+ (decf sp 2)
+ (let ((parent (svref stack sp))
+ (dir (svref stack (1+ sp))))
+ (setf (svref parent dir) node
+ node parent))
+
+ ;; If there's a level difference between this node and its
+ ;; children, bring it (and, if it exists, its right
+ ;; counterpart) down one level.
+ (let ((level-1 (1- (node-level node)))
+ (left (node-left node))
+ (right (node-right node)))
+ (when (flet ((level (node)
+ (if node (node-level node) -1)))
+ (declare (inline level))
+ (or (< (level left) level-1)
+ (< (level right) level-1)))
+ (setf (node-level node) level-1)
+ (when (and right (> (node-level right) level-1))
+ (setf (node-level right) level-1))
+
+ ;; Now we must fix up the balancing rules. Apparently
+ ;; three skews and two splits suffice.
+ (setf node (skew node))
+ (let ((right (node-right node)))
+ (when right
+ (setf right (skew right)
+ (node-right node) right)
+ (let ((rright (node-right right)))
+ (when rright
+ (setf (node-right right) (skew rright))))))
+ (setf node (split node))
+ (let ((right (node-right node)))
+ (when right (setf (node-right node) (split right)))))))
+
+ ;; Store the new root.
+ (setf (tree-root tree) node)))))
+
+(defun aa-tree-iterator (tree)
+ "Returns a tree iterator function for TREE. The function returns three
+ values. For each node in the tree, it returns t, the key and the value;
+ then, it returns nil three times."
+ (let ((root (tree-root tree)))
+ (if (null root)
+ (lambda () (values nil nil nil))
+ (let ((stack (make-array (* 2 (1+ (node-level root)))))
+ (sp 0))
+ (flet ((pushleft (node)
+ (do ((node node (node-left node)))
+ ((null node))
+ (setf (svref stack sp) node)
+ (incf sp))))
+ (pushleft root)
+ (lambda ()
+ (cond ((zerop sp) (values nil nil nil))
+ (t (let ((node (svref stack (decf sp))))
+ (pushleft (node-right node))
+ (values t (node-key node) (node-data node)))))))))))
+
+(defun mapaa (func tree)
+ "Apply FUNC to each key and value in the TREE."
+ (labels ((walk (node)
+ (when node
+ (walk (node-left node))
+ (funcall func (node-key node) (node-data node))
+ (walk (node-right node)))))
+ (walk (tree-root tree))
+ nil))
+
+(defmacro doaa ((key value tree &optional result) &body body)
+ "Iterate over the items of TREE; for each one, bind KEY to its key and
+ VALUE to the associated data, and evaluate BODY, which is an implicit
+ tagbody. Finally, return RESULT. Either KEY or VALUE (or both!) may be
+ nil to indicate `don't care'."
+ (with-parsed-body (decls body body)
+ (let ((ignores nil))
+ (unless key (setf key (gensym "KEY")) (push key ignores))
+ (unless value (setf value (gensym "VALUE")) (push value ignores))
+ `(block nil
+ (mapaa (lambda (,key ,value)
+ ,@decls
+ ,@(and ignores `((declare (ignore ,@ignores))))
+ (tagbody ,@body))
+ ,tree)
+ ,result))))
+
+;;;--------------------------------------------------------------------------
+;;; Testing.
+
+#+debug
+(defun tree-print (tree &optional (stream *standard-output*))
+ "Print a TREE to an output STREAM in a comprehesible way."
+ (labels ((walk (depth node)
+ (when node
+ (walk (1+ depth) (node-left node))
+ (format stream "~v@T~A: ~S => ~S~%"
+ (* depth 2)
+ (node-level node)
+ (node-key node)
+ (node-data node))
+ (walk (1+ depth) (node-right node)))))
+ (walk 0 (tree-root tree))))
+
+(defun tree-build (key< &rest items)
+ "Return a new tree sorted according to KEY<, containing the given ITEMS."
+ (let ((tree (make-aa-tree key<)))
+ (dolist (item items)
+ (setf (getaa tree item) nil))
+ tree))
+
+#+debug
+(defun test-iterator (tree)
+ (let ((iter (aa-tree-iterator tree)))
+ (mapaa (lambda (key value)
+ (multiple-value-bind (iwin ikey ivalue) (funcall iter)
+ (assert (and iwin
+ (eql key ikey)
+ (eql value ivalue)))))
+ tree)
+ (assert (null (nth-value 0 (funcall iter))))))
+
+#+debug
+(defun tree-check (tree)
+ "Checks the invariants on a TREE."
+ (let ((key< (tree-key< tree)))
+ (labels ((check (node)
+ (if (null node)
+ (values nil nil)
+ (let ((key (node-key node))
+ (level (node-level node))
+ (left (node-left node))
+ (right (node-right node)))
+ (multiple-value-bind (lmin lmax) (check left)
+ (multiple-value-bind (rmin rmax) (check right)
+ (assert (or (null lmax) (funcall key< lmax key)))
+ (assert (or (null rmin) (funcall key< key rmin)))
+ (assert (if (null left)
+ (= level 0)
+ (= (node-level left) (- level 1))))
+ (assert (if (null right)
+ (= level 0)
+ (let ((rright (node-right right)))
+ (or (= (node-level right) (- level 1))
+ (and (= (node-level right) level)
+ (or (null rright)
+ (= (node-level rright)
+ (- level 1))))))))
+ (values (or lmin key) (or rmax key))))))))
+ (check (tree-root tree)))))
+
+#+debug
+(defun test (&key (state (make-random-state))
+ (count nil)
+ (items nil)
+ (verbose 1))
+ (let ((in (make-array 0 :element-type 'string
+ :adjustable t :fill-pointer 0))
+ (out (make-array 0 :element-type 'string
+ :adjustable t :fill-pointer 0))
+ (tree (make-aa-tree #'string<)))
+
+ ;; Slurp in the word list
+ (with-open-file (dict #p"/usr/share/dict/words")
+ (loop for line = (read-line dict nil)
+ while (and line (not (eql items 0)))
+ do (vector-push-extend line out)
+ when items do (decf items)))
+
+ (labels ((add (v w)
+ (vector-push-extend w v))
+ (rm (v i)
+ (let ((n (1- (length v))))
+ (setf (aref v i) (aref v n))
+ (decf (fill-pointer v))))
+ (insert ()
+ (let* ((i (random (length out) state))
+ (w (aref out i)))
+ (setf (getaa tree w) nil)
+ (rm out i)
+ (add in w)
+ (when (>= verbose 2) (format t "insert ~A~%" w))))
+ (remove ()
+ (let* ((i (random (length in) state))
+ (w (aref in i)))
+ (remaa tree w)
+ (rm in i)
+ (add out w)
+ (when (>= verbose 2) (format t "remove ~A~%" w))))
+ (check ()
+ (when (>= verbose 2) (format t "check...~%"))
+ (tree-check tree)
+ (sort in #'string<)
+ (loop with i = (aa-tree-iterator tree)
+ for w across in
+ for (win key value) = (multiple-value-list (funcall i))
+ do (assert (eq w (and win key)))
+ while w
+ finally (assert (null (nth-value 0 (funcall i)))))))
+ (loop with prob = (if count (/ count 100) 1000)
+ until (eql count 0)
+ when count do (decf count)
+ do (case (random prob state)
+ (0 (check) (when (= verbose 1) (write-char #\?)))
+ (t (if (< (random (+ (length in) (length out)) state)
+ (length out))
+ (progn (insert)
+ (when (= verbose 1) (write-char #\+)))
+ (progn (remove)
+ (when (= verbose 1) (write-char #\-))))))
+ do (force-output)
+ finally (check)))))
+
+;;;----- That's all, folks --------------------------------------------------
~mdw / lisp / commitdiff