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