mdw@git.distorted.org.uk Git - lisp/blob - aa-tree.lisp

   1 ;;; -*-lisp-*-
   2 ;;;
   3 ;;; Andersson tree implementation
   4 ;;;
   5 ;;; (c) 2006 Straylight/Edgeware
   6 ;;;
   7
   8 ;;;----- Licensing notice ---------------------------------------------------
   9 ;;;
  10 ;;; This program is free software; you can redistribute it and/or modify
  11 ;;; it under the terms of the GNU General Public License as published by
  12 ;;; the Free Software Foundation; either version 2 of the License, or
  13 ;;; (at your option) any later version.
  14 ;;;
  15 ;;; This program is distributed in the hope that it will be useful,
  16 ;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
  17 ;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  18 ;;; GNU General Public License for more details.
  19 ;;;
  20 ;;; You should have received a copy of the GNU General Public License
  21 ;;; along with this program; if not, write to the Free Software Foundation,
  22 ;;; Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
  23
  24 ;;;--------------------------------------------------------------------------
  25 ;;; Package.
  26
  27 (defpackage #:aa-tree
  28   (:use #:common-lisp #:mdw.base))
  29 (in-package #:aa-tree)
  30
  31 ;;;--------------------------------------------------------------------------
  32 ;;; The underlying implementation.
  33
  34 (deftype stack-pointer () '(integer 0 255))
  35
  36 (defstruct (tree-node
  37              (:conc-name node-)
  38              (:type vector)
  39              (:constructor make-tree-node
  40                            (key &optional data level left right)))
  41   "Structure representing a node in an Andersson tree."
  42   (left nil :type (or null tree-node))
  43   (right nil :type (or null tree-node))
  44   (level 0 :type stack-pointer)
  45   key
  46   data)
  47
  48 (deftype tree-node () 'simple-vector)
  49
  50 (export '(make-aa-tree aa-tree aa-tree-p aa-tree-key<))
  51 (defstruct (aa-tree
  52              (:predicate treep)
  53              (:constructor make-aa-tree
  54                            (key<-name
  55                             &aux
  56                             (key< (functionify key<-name))))
  57              (:conc-name tree-))
  58   "Structure representing an Andersson tree."
  59   (root nil :type (or null tree-node))
  60   (stack (make-array 32) :type simple-vector)
  61   (key< (slot-uninitialized) :read-only t :type (function (t t) t)))
  62
  63 (declaim (inline skew split))
  64
  65 (defun skew (node)
  66   "Implements the `skew' operation on a tree node, eliminating left-pointing
  67    internal pointers by applying right-rotation.  Returns the replacement
  68    node."
  69   (declare (type tree-node node))
  70   (let ((left (node-left node)))
  71     (when (and left (= (node-level node) (node-level left)))
  72       (shiftf (node-left node) (node-right left) node left))
  73     node))
  74
  75 (defun split (node)
  76   "Implements the `split' operation on a tree node, eliminating overly-large
  77    pseudo-nodes by applying left-rotation.  Returns the replacement node."
  78   (declare (type tree-node node))
  79   (let* ((right (node-right node))
  80          (rright (and right (node-right right))))
  81     (when (and rright (= (node-level node) (node-level rright)))
  82       (shiftf (node-right node) (node-left right) node right)
  83       (incf (node-level node)))
  84     node))
  85
  86 (defun get-tree-stack (tree)
  87   "Return the current stack for the TREE.  This is used to remember the path
  88    taken during a search in tree, so we can fix it up afterwards.  Keeping
  89    just one stack for the tree saves on consing; it's not safe to do
  90    simultaneous destructive operations on a tree anyway, so this is a
  91    reasonable thing to do.  This function ensures that the stack attached to
  92    the tree is actually large enough before returning it."
  93   (declare (type aa-tree tree))
  94   (let* ((root (tree-root tree))
  95          (want (* 4 (+ (if root (node-level root) 0) 2)))
  96          (stack (tree-stack tree))
  97          (size (array-dimension (tree-stack tree) 0)))
  98     (if (>= size want)
  99         stack
 100         (do ((need (ash size 1) (ash need 1)))
 101             ((>= need want) (setf (tree-stack tree) (make-array need)))))))
 102
 103 (export 'getaa)
 104 (defun getaa (tree key &optional default)
 105   "Look up the given KEY in an Andersson TREE; if the KEY was found, return
 106    the corresponding data and t, otherwise return DEFAULT and nil."
 107   (declare (type aa-tree tree))
 108   (let ((key< (tree-key< tree))
 109         (node (tree-root tree))
 110         (candidate nil)
 111         (candidate-key nil))
 112     (declare (type (function (t t) t) key<)
 113              (type (or null tree-node) node candidate))
 114     (flet ((key< (x y)
 115              (funcall key< x y)))
 116       (declare (inline key<))
 117       (loop (cond (node
 118                    (let ((node-key (node-key node)))
 119                      (if (key< key node-key)
 120                          (setf node (node-left node))
 121                          (setf candidate node
 122                                candidate-key node-key
 123                                node (node-right node)))))
 124                   ((and candidate (not (key< candidate-key key)))
 125                    (return (values (node-data candidate) t)))
 126                   (t
 127                    (return (values default nil))))))))
 128
 129 (defun tree-probe (tree key)
 130   "Do a search in an Andersson TREE for the KEY, returning three values.  The
 131    second and third are a stack of alternating nodes and direction bits, and
 132    a stack pointer (empty, ascending), which together describe the path from
 133    the tree root to the successor of the sought-for node.  The first is
 134    either the sought-for node itself, or nil if it wasn't there."
 135   (declare (type aa-tree tree))
 136   (let ((key< (tree-key< tree))
 137         (stack (get-tree-stack tree))
 138         (sp 0)
 139         (candidate nil)
 140         (candidate-key nil))
 141     (declare (type (function (t t) t) key<)
 142              (type simple-vector stack)
 143              (type stack-pointer sp)
 144              (type (or null tree-node) candidate))
 145     (flet ((pathpush (v i)
 146              (setf (svref stack sp) v
 147                    (svref stack (1+ sp)) i)
 148              (incf sp 2))
 149            (key< (x y)
 150              (funcall key< x y)))
 151       (declare (inline pathpush key<))
 152       (let ((node (tree-root tree)))
 153         (loop (when (null node)
 154                 (return))
 155               (let* ((node-key (node-key node))
 156                      (dir (cond ((key< key node-key) 0)
 157                                 (t (setf candidate node
 158                                          candidate-key node-key)
 159                                    1))))
 160                 (pathpush node dir)
 161                 (setf node (svref node dir)))))
 162       (values (if (and candidate (not (key< candidate-key key)))
 163                   candidate
 164                   nil)
 165               stack
 166               sp))))
 167
 168 (defun fixup-insert (tree stack sp node)
 169   "TREE is an Andersson tree, STACK and SP are the values from a failed call
 170    to tree-probe, and NODE is a newly-created node.  Insert the NODE into
 171    the tree, fix up its balance."
 172   (declare (type aa-tree tree)
 173            (type simple-vector stack)
 174            (type stack-pointer sp)
 175            (type tree-node node))
 176   (loop (when (zerop sp)
 177           (return))
 178         (decf sp 2)
 179         (let ((parent (svref stack sp))
 180               (dir (svref stack (1+ sp))))
 181           (setf (svref parent dir) node
 182                 node parent))
 183         (setf node (split (skew node))))
 184   (setf (tree-root tree) node))
 185
 186 (defun (setf getaa) (data tree key &optional ignore)
 187   "Inserts a new node with the given KEY into an Andersson TREE, if there
 188    wasn't one already.  Returns two values: the requested node, and either t
 189    if the node was inserted, or nil if it was already there."
 190   (declare (type aa-tree tree)
 191            (ignore ignore))
 192   (multiple-value-bind (node stack sp) (tree-probe tree key)
 193     (cond (node (setf (node-data node) data))
 194           (t (fixup-insert tree stack sp (make-tree-node key data)) data))))
 195
 196 (export 'updateaa)
 197 (defun updateaa (tree key func)
 198   "Search TREE for an item with the given KEY.  If it was found, call FUNC
 199    with arguments of the node's data and t, and store its result as the
 200    node's new data.  If it was absent, call FUNC with arguments nil and nil,
 201    and make a new node with the KEY and return value.  The FUNC can escape to
 202    prevent the node being created (though this is probably not useful)."
 203   (declare (type aa-tree tree))
 204   (multiple-value-bind (node stack sp) (tree-probe tree key)
 205     (cond (node (setf (node-data node) (funcall func (node-data node) t)))
 206           (t (let ((data (funcall func nil nil)))
 207                (fixup-insert tree stack sp (make-tree-node key data))
 208                data)))))
 209
 210 (export 'remaa)
 211 (defun remaa (tree key)
 212   "Deletes the node with the given KEY from an Andersson TREE.  Returns t if
 213    the node was found and deleted, or nil if it wasn't there to begin with."
 214   (declare (type aa-tree tree))
 215   (multiple-value-bind (candidate stack sp) (tree-probe tree key)
 216     (when candidate
 217       (decf sp 2)
 218       (let ((node (svref stack sp)))
 219
 220         ;; Unsplice the candidate node from the tree, leaving node as its
 221         ;; replacement.
 222         (if (eq candidate node)
 223             (setf node nil)
 224             (setf (node-key candidate) (node-key node)
 225                   (node-data candidate) (node-data node)
 226                   node (node-right node)))
 227
 228         ;; Now wander back up the tree, fixing it as we go.
 229         (loop (when (zerop sp)
 230                 (return))
 231               (decf sp 2)
 232               (let ((parent (svref stack sp))
 233                     (dir (svref stack (1+ sp))))
 234                 (setf (svref parent dir) node
 235                       node parent))
 236
 237               ;; If there's a level difference between this node and its
 238               ;; children, bring it (and, if it exists, its right
 239               ;; counterpart) down one level.
 240               (let ((level-1 (1- (node-level node)))
 241                     (left (node-left node))
 242                     (right (node-right node)))
 243                 (when (flet ((level (node)
 244                                (if node (node-level node) -1)))
 245                         (declare (inline level))
 246                         (or (< (level left) level-1)
 247                             (< (level right) level-1)))
 248                   (setf (node-level node) level-1)
 249                   (when (and right (> (node-level right) level-1))
 250                     (setf (node-level right) level-1))
 251
 252                   ;; Now we must fix up the balancing rules.  Apparently
 253                   ;; three skews and two splits suffice.
 254                   (setf node (skew node))
 255                   (let ((right (node-right node)))
 256                     (when right
 257                       (setf right (skew right)
 258                             (node-right node) right)
 259                       (let ((rright (node-right right)))
 260                         (when rright
 261                           (setf (node-right right) (skew rright))))))
 262                   (setf node (split node))
 263                   (let ((right (node-right node)))
 264                     (when right (setf (node-right node) (split right)))))))
 265
 266         ;; Store the new root.
 267         (setf (tree-root tree) node)))))
 268
 269 (export 'aa-tree-iterator)
 270 (defun aa-tree-iterator (tree)
 271   "Returns a tree iterator function for TREE.  The function returns three
 272    values.  For each node in the tree, it returns t, the key and the value;
 273    then, it returns nil three times."
 274   (let ((root (tree-root tree)))
 275     (if (null root)
 276         (lambda () (values nil nil nil))
 277         (let ((stack (make-array (* 2 (1+ (node-level root)))))
 278               (sp 0))
 279           (flet ((pushleft (node)
 280                    (do ((node node (node-left node)))
 281                        ((null node))
 282                      (setf (svref stack sp) node)
 283                      (incf sp))))
 284             (pushleft root)
 285             (lambda ()
 286               (cond ((zerop sp) (values nil nil nil))
 287                     (t (let ((node (svref stack (decf sp))))
 288                          (pushleft (node-right node))
 289                          (values t (node-key node) (node-data node)))))))))))
 290
 291 (export 'mapaa)
 292 (defun mapaa (func tree)
 293   "Apply FUNC to each key and value in the TREE."
 294   (labels ((walk (node)
 295              (when node
 296                (walk (node-left node))
 297                (funcall func (node-key node) (node-data node))
 298                (walk (node-right node)))))
 299     (walk (tree-root tree))
 300     nil))
 301
 302 (export 'doaa)
 303 (defmacro doaa ((key value tree &optional result) &body body)
 304   "Iterate over the items of TREE; for each one, bind KEY to its key and
 305    VALUE to the associated data, and evaluate BODY, which is an implicit
 306    tagbody.  Finally, return RESULT.  Either KEY or VALUE (or both!) may be
 307    nil to indicate `don't care'."
 308   (with-parsed-body (body decls) body
 309     (let ((ignores nil))
 310       (unless key (setf key (gensym "KEY")) (push key ignores))
 311       (unless value (setf value (gensym "VALUE")) (push value ignores))
 312       `(block nil
 313          (mapaa (lambda (,key ,value)
 314                   ,@decls
 315                   ,@(and ignores `((declare (ignore ,@ignores))))
 316                   (tagbody ,@body))
 317                 ,tree)
 318          ,result))))
 319
 320 ;;;--------------------------------------------------------------------------
 321 ;;; Testing.
 322
 323 #+debug
 324 (defun tree-print (tree &optional (stream *standard-output*))
 325   "Print a TREE to an output STREAM in a comprehesible way."
 326   (labels ((walk (depth node)
 327              (when node
 328                (walk (1+ depth) (node-left node))
 329                (format stream "~v@T~A: ~S => ~S~%"
 330                        (* depth 2)
 331                        (node-level node)
 332                        (node-key node)
 333                        (node-data node))
 334                (walk (1+ depth) (node-right node)))))
 335     (walk 0 (tree-root tree))))
 336
 337 (defun tree-build (key< &rest items)
 338   "Return a new tree sorted according to KEY<, containing the given ITEMS."
 339   (let ((tree (make-aa-tree key<)))
 340     (dolist (item items)
 341       (setf (getaa tree item) nil))
 342     tree))
 343
 344 #+debug
 345 (defun test-iterator (tree)
 346   (let ((iter (aa-tree-iterator tree)))
 347     (mapaa (lambda (key value)
 348              (multiple-value-bind (iwin ikey ivalue) (funcall iter)
 349                (assert (and iwin
 350                             (eql key ikey)
 351                             (eql value ivalue)))))
 352            tree)
 353     (assert (null (nth-value 0 (funcall iter))))))
 354
 355 #+debug
 356 (defun tree-check (tree)
 357   "Checks the invariants on a TREE."
 358   (let ((key< (tree-key< tree)))
 359     (labels ((check (node)
 360                (if (null node)
 361                    (values nil nil)
 362                    (let ((key (node-key node))
 363                          (level (node-level node))
 364                          (left (node-left node))
 365                          (right (node-right node)))
 366                      (multiple-value-bind (lmin lmax) (check left)
 367                        (multiple-value-bind (rmin rmax) (check right)
 368                          (assert (or (null lmax) (funcall key< lmax key)))
 369                          (assert (or (null rmin) (funcall key< key rmin)))
 370                          (assert (if (null left)
 371                                      (= level 0)
 372                                      (= (node-level left) (- level 1))))
 373                          (assert (if (null right)
 374                                      (= level 0)
 375                                      (let ((rright (node-right right)))
 376                                        (or (= (node-level right) (- level 1))
 377                                            (and (= (node-level right) level)
 378                                                 (or (null rright)
 379                                                     (= (node-level rright)
 380                                                        (- level 1))))))))
 381                          (values (or lmin key) (or rmax key))))))))
 382       (check (tree-root tree)))))
 383
 384 #+debug
 385 (defun test (&key (state (make-random-state))
 386              (count nil)
 387              (items nil)
 388              (verbose 1))
 389   (let ((in (make-array 0 :element-type 'string
 390                         :adjustable t :fill-pointer 0))
 391         (out (make-array 0 :element-type 'string
 392                          :adjustable t :fill-pointer 0))
 393         (tree (make-aa-tree #'string<)))
 394
 395     ;; Slurp in the word list
 396     (with-open-file (dict #p"/usr/share/dict/words")
 397       (loop for line = (read-line dict nil)
 398             while (and line (not (eql items 0)))
 399             do (vector-push-extend line out)
 400             when items do (decf items)))
 401
 402     (labels ((add (v w)
 403                (vector-push-extend w v))
 404              (rm (v i)
 405                (let ((n (1- (length v))))
 406                  (setf (aref v i) (aref v n))
 407                  (decf (fill-pointer v))))
 408              (insert ()
 409                (let* ((i (random (length out) state))
 410                       (w (aref out i)))
 411                  (setf (getaa tree w) nil)
 412                  (rm out i)
 413                  (add in w)
 414                  (when (>= verbose 2) (format t "insert ~A~%" w))))
 415              (remove ()
 416                (let* ((i (random (length in) state))
 417                       (w (aref in i)))
 418                  (remaa tree w)
 419                  (rm in i)
 420                  (add out w)
 421                  (when (>= verbose 2) (format t "remove ~A~%" w))))
 422              (check ()
 423                (when (>= verbose 2) (format t "check...~%"))
 424                (tree-check tree)
 425                (sort in #'string<)
 426                (loop with i = (aa-tree-iterator tree)
 427                      for w across in
 428                      for (win key value) = (multiple-value-list (funcall i))
 429                      do (assert (eq w (and win key)))
 430                      while w
 431                      finally (assert (null (nth-value 0 (funcall i)))))))
 432       (loop with prob = (if count (/ count 100) 1000)
 433             until (eql count 0)
 434             when count do (decf count)
 435             do (case (random prob state)
 436                  (0 (check) (when (= verbose 1) (write-char #\?)))
 437                  (t (if (< (random (+ (length in) (length out)) state)
 438                            (length out))
 439                         (progn (insert)
 440                                (when (= verbose 1) (write-char #\+)))
 441                         (progn (remove)
 442                                (when (= verbose 1) (write-char #\-))))))
 443             do (force-output)
 444             finally (check)))))
 445
 446 ;;;----- That's all, folks --------------------------------------------------