lisp/emacs-lisp/radix-tree.el

   1 ;;; radix-tree.el --- A simple library of radix trees  -*- lexical-binding: t; -*-
   2
   3 ;; Copyright (C) 2016-2017 Free Software Foundation, Inc.
   4
   5 ;; Author: Stefan Monnier <monnier@iro.umontreal.ca>
   6 ;; Keywords:
   7
   8 ;; This file is part of GNU Emacs.
   9
  10 ;; GNU Emacs 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 3 of the License, or
  13 ;; (at your option) any later version.
  14
  15 ;; GNU Emacs 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 GNU Emacs.  If not, see <http://www.gnu.org/licenses/>.
  22
  23 ;;; Commentary:
  24
  25 ;; There are many different options for how to represent radix trees
  26 ;; in Elisp.  Here I chose a very simple one.  A radix-tree can be either:
  27 ;; - a node, of the form ((PREFIX . PTREE) . RTREE) where PREFIX is a string
  28 ;;   meaning that everything that starts with PREFIX is in PTREE,
  29 ;;   and everything else in RTREE.  It also has the property that
  30 ;;   everything that starts with the first letter of PREFIX but not with
  31 ;;   that whole PREFIX is not in RTREE (i.e. is not in the tree at all).
  32 ;; - anything else is taken as the value to associate with the empty string.
  33 ;; So every node is basically an (improper) alist where each mapping applies
  34 ;; to a different leading letter.
  35 ;;
  36 ;; The main downside of this representation is that the lookup operation
  37 ;; is slower because each level of the tree is an alist rather than some kind
  38 ;; of array, so every level's lookup is O(N) rather than O(1).  We could easily
  39 ;; solve this by using char-tables instead of alists, but that would make every
  40 ;; level take up a lot more memory, and it would make the resulting
  41 ;; data structure harder to read (by a human) when printed out.
  42
  43 ;;; Code:
  44
  45 (defun radix-tree--insert (tree key val i)
  46   (pcase tree
  47     (`((,prefix . ,ptree) . ,rtree)
  48      (let* ((ni (+ i (length prefix)))
  49             (cmp (compare-strings prefix nil nil key i ni)))
  50        (if (eq t cmp)
  51            (let ((nptree (radix-tree--insert ptree key val ni)))
  52              `((,prefix . ,nptree) . ,rtree))
  53          (let ((n (if (< cmp 0) (- -1 cmp) (- cmp 1))))
  54            (if (zerop n)
  55                (let ((nrtree (radix-tree--insert rtree key val i)))
  56                  `((,prefix . ,ptree) . ,nrtree))
  57              (let* ((nprefix (substring prefix 0 n))
  58                     (kprefix (substring key (+ i n)))
  59                     (pprefix (substring prefix n))
  60                     (ktree (if (equal kprefix "") val
  61                              `((,kprefix . ,val)))))
  62                `((,nprefix
  63                   . ((,pprefix . ,ptree) . ,ktree))
  64                  . ,rtree)))))))
  65     (_
  66      (if (= (length key) i) val
  67        (let ((prefix (substring key i)))
  68          `((,prefix . ,val) . ,tree))))))
  69
  70 (defun radix-tree--remove (tree key i)
  71   (pcase tree
  72     (`((,prefix . ,ptree) . ,rtree)
  73      (let* ((ni (+ i (length prefix)))
  74             (cmp (compare-strings prefix nil nil key i ni)))
  75        (if (eq t cmp)
  76            (pcase (radix-tree--remove ptree key ni)
  77              (`nil rtree)
  78              (`((,pprefix . ,pptree))
  79               `((,(concat prefix pprefix) . ,pptree) . ,rtree))
  80              (nptree `((,prefix . ,nptree) . ,rtree)))
  81          (let ((n (if (< cmp 0) (- -1 cmp) (- cmp 1))))
  82            (if (zerop n)
  83                (let ((nrtree (radix-tree--remove rtree key i)))
  84                  `((,prefix . ,ptree) . ,nrtree))
  85              tree)))))
  86     (_
  87      (if (= (length key) i) nil tree))))
  88
  89
  90 (defun radix-tree--lookup (tree string i)
  91   (pcase tree
  92     (`((,prefix . ,ptree) . ,rtree)
  93      (let* ((ni (+ i (length prefix)))
  94             (cmp (compare-strings prefix nil nil string i ni)))
  95        (if (eq t cmp)
  96            (radix-tree--lookup ptree string ni)
  97          (let ((n (if (< cmp 0) (- -1 cmp) (- cmp 1))))
  98            (if (zerop n)
  99                (radix-tree--lookup rtree string i)
 100              (+ i n))))))
 101     (val
 102      (if (and val (equal (length string) i))
 103          (if (integerp val) `(t . ,val) val)
 104        i))))
 105
 106 ;; (defun radix-tree--trim (tree string i)
 107 ;;   (if (= i (length string))
 108 ;;       tree
 109 ;;     (pcase tree
 110 ;;       (`((,prefix . ,ptree) . ,rtree)
 111 ;;        (let* ((ni (+ i (length prefix)))
 112 ;;               (cmp (compare-strings prefix nil nil string i ni))
 113 ;;               ;; FIXME: We could compute nrtree more efficiently
 114 ;;               ;; whenever cmp is not -1 or 1.
 115 ;;               (nrtree (radix-tree--trim rtree string i)))
 116 ;;          (if (eq t cmp)
 117 ;;              (pcase (radix-tree--trim ptree string ni)
 118 ;;                (`nil nrtree)
 119 ;;                (`((,pprefix . ,pptree))
 120 ;;                 `((,(concat prefix pprefix) . ,pptree) . ,nrtree))
 121 ;;                (nptree `((,prefix . ,nptree) . ,nrtree)))
 122 ;;            (let ((n (if (< cmp 0) (- -1 cmp) (- cmp 1))))
 123 ;;              (cond
 124 ;;               ((equal (+ n i) (length string))
 125 ;;                `((,prefix . ,ptree) . ,nrtree))
 126 ;;               (t nrtree))))))
 127 ;;       (val val))))
 128
 129 (defun radix-tree--prefixes (tree string i prefixes)
 130   (pcase tree
 131     (`((,prefix . ,ptree) . ,rtree)
 132      (let* ((ni (+ i (length prefix)))
 133             (cmp (compare-strings prefix nil nil string i ni))
 134             ;; FIXME: We could compute prefixes more efficiently
 135             ;; whenever cmp is not -1 or 1.
 136             (prefixes (radix-tree--prefixes rtree string i prefixes)))
 137        (if (eq t cmp)
 138            (radix-tree--prefixes ptree string ni prefixes)
 139          prefixes)))
 140     (val
 141      (if (null val)
 142          prefixes
 143        (cons (cons (substring string 0 i)
 144                    (if (eq (car-safe val) t) (cdr val) val))
 145              prefixes)))))
 146
 147 (defun radix-tree--subtree (tree string i)
 148   (if (equal (length string) i) tree
 149     (pcase tree
 150       (`((,prefix . ,ptree) . ,rtree)
 151        (let* ((ni (+ i (length prefix)))
 152               (cmp (compare-strings prefix nil nil string i ni)))
 153          (if (eq t cmp)
 154              (radix-tree--subtree ptree string ni)
 155            (let ((n (if (< cmp 0) (- -1 cmp) (- cmp 1))))
 156              (cond
 157               ((zerop n) (radix-tree--subtree rtree string i))
 158               ((equal (+ n i) (length string))
 159                (let ((nprefix (substring prefix n)))
 160                  `((,nprefix . ,ptree))))
 161               (t nil))))))
 162       (_ nil))))
 163
 164 ;;; Entry points
 165
 166 (defconst radix-tree-empty nil
 167   "The empty radix-tree.")
 168
 169 (defun radix-tree-insert (tree key val)
 170   "Insert a mapping from KEY to VAL in radix TREE."
 171   (when (consp val) (setq val `(t . ,val)))
 172   (if val (radix-tree--insert tree key val 0)
 173     (radix-tree--remove tree key 0)))
 174
 175 (defun radix-tree-lookup (tree key)
 176   "Return the value associated to KEY in radix TREE.
 177 If not found, return nil."
 178   (pcase (radix-tree--lookup tree key 0)
 179     (`(t . ,val) val)
 180     ((pred numberp) nil)
 181     (val val)))
 182
 183 (defun radix-tree-subtree (tree string)
 184   "Return the subtree of TREE rooted at the prefix STRING."
 185   (radix-tree--subtree tree string 0))
 186
 187 ;; (defun radix-tree-trim (tree string)
 188 ;;   "Return a TREE which only holds entries \"related\" to STRING.
 189 ;; \"Related\" is here defined as entries where there's a `string-prefix-p' relation
 190 ;; between STRING and the key."
 191 ;;   (radix-tree-trim tree string 0))
 192
 193 (defun radix-tree-prefixes (tree string)
 194   "Return an alist of all bindings in TREE for prefixes of STRING."
 195   (radix-tree--prefixes tree string 0 nil))
 196
 197 (eval-and-compile
 198   (pcase-defmacro radix-tree-leaf (vpat)
 199     ;; FIXME: We'd like to use a negative pattern (not consp), but pcase
 200     ;; doesn't support it.  Using `atom' works but generates sub-optimal code.
 201     `(or `(t . ,,vpat) (and (pred atom) ,vpat))))
 202
 203 (defun radix-tree-iter-subtrees (tree fun)
 204   "Apply FUN to every immediate subtree of radix TREE.
 205 FUN is called with two arguments: PREFIX and SUBTREE.
 206 You can test if SUBTREE is a leaf (and extract its value) with the
 207 pcase pattern (radix-tree-leaf PAT)."
 208   (while tree
 209     (pcase tree
 210       (`((,prefix . ,ptree) . ,rtree)
 211        (funcall fun prefix ptree)
 212        (setq tree rtree))
 213       (_ (funcall fun "" tree)
 214          (setq tree nil)))))
 215
 216 (defun radix-tree-iter-mappings (tree fun &optional prefix)
 217   "Apply FUN to every mapping in TREE.
 218 FUN is called with two arguments: KEY and VAL.
 219 PREFIX is only used internally."
 220   (radix-tree-iter-subtrees
 221    tree
 222    (lambda (p s)
 223      (let ((nprefix (concat prefix p)))
 224        (pcase s
 225          ((radix-tree-leaf v) (funcall fun nprefix v))
 226          (_ (radix-tree-iter-mappings s fun nprefix)))))))
 227
 228 ;; (defun radix-tree->alist (tree)
 229 ;;   (let ((al nil))
 230 ;;     (radix-tree-iter-mappings tree (lambda (p v) (push (cons p v) al)))
 231 ;;     al))
 232
 233 (defun radix-tree-count (tree)
 234   (let ((i 0))
 235     (radix-tree-iter-mappings tree (lambda (_k _v) (setq i (1+ i))))
 236     i))
 237
 238 (defun radix-tree-from-map (map)
 239   ;; Aka (cl-defmethod map-into (map (type (eql radix-tree)))) ...)
 240   (require 'map)
 241   (let ((rt nil))
 242     (map-apply (lambda (k v) (setq rt (radix-tree-insert rt k v))) map)
 243     rt))
 244
 245 (provide 'radix-tree)
 246 ;;; radix-tree.el ends here