UNSLISP/trie.scm

;;; Copyright (C) Peter McGoron 2024
;;; 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, version 3 of the License.
;;; 
;;; 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, see <https://www.gnu.org/licenses/>.

;;; ;;;;;;;;;;;;;;
;;; Character maps
;;; ;;;;;;;;;;;;;;

(define char<=>
  (lambda (x y)
    (let ((x (char->integer x))
          (y (char->integer y)))
      (cond
        ((< x y) '<)
        ((= x y) '=)
        (else '>)))))


(define %charmap:<=> (set:<=>-to-map char<=>))
(define %charmap:update (set:update %charmap:<=>))

(define charmap:update (map:update %charmap:update))
(define charmap:insert (map:insert %charmap:update))
(define charmap:search (map:search %charmap:<=>))

;;; ;;;;;;;;;;
;;; Char tries
;;; 
;;; Tries are pairs, CAR = backing map from chars to trie nodes, and
;;; CDR = function tied to this part of the trie.
;;; ;;;;;;;;;;

;;; (TRIE:NEW TABLE FUNCTION)
(define trie:new cons)
(define trie:empty (trie:new '() #f))

;;; Get the function inside of trie NODE.
(define trie:function
  (lambda (node)
    (if (null? node)
        #f
        (cdr node))))

;;; Get the backing set inside of trie NODE.
(define trie:backing
  (lambda (node)
    (if (null? node)
        '()
        (car node))))

;;; Insert STRING-AS-LIST into NODE with value FUNCTION.
(define trie:insert
  (lambda (node string-as-list function)
    (if (null? string-as-list)
        (trie:new (trie:backing node) function)
        (let ((ch (car string-as-list))
              (string-as-list (cdr string-as-list)))
          (let ((newtree (charmap:update (trie:backing node) ch
                                         (lambda (node oldnode)
                                           (if oldnode
                                               (map:node-new-val
                                                oldnode
                                                ch
                                                (trie:insert (map:val oldnode)
                                                             string-as-list
                                                             function))
                                               (map:empty-node ch
                                                               (trie:insert
                                                                trie:empty
                                                                string-as-list
                                                                function)))))))
            (trie:new newtree (trie:function node)))))))

(define trie:insert-many
  (lambda (node lst)
    (fold (lambda (pair node)
            (let ((key (car pair)))
              (let ((key (cond
                           ((list? key) key)
                           ((string? key) (string->list key))
                           ((char? key) (list key)))))
                (trie:insert node key (cdr pair)))))
          node lst)))

;;; Search for CH in NODE.
(define trie:search-single
  (lambda (ch node)
    (if (null? node)
        '()
        (let ((node (charmap:search (trie:backing node) ch)))
          (if (null? node)
              '()
              (map:val node))))))

(define trie:search
  (lambda (node str)
    (fold trie:search-single node (string->list str))))

;;; ;;;;;
;;; Tests
;;; ;;;;;

(define %trie:tests
  (list
   (cons "insert with prefixes"
         (lambda ()
           (let ((trie
                  (trie:insert-many trie:empty
                                    (list
                                     (cons (string #\a #\b #\c #\d) 10)
                                     (cons (string #\a #\b #\c) 20)
                                     (cons (string #\a #\b #\c #\d #\e) 30)
                                     (cons (string #\b) 40)
                                     (cons (string #\b #\e #\g) 50))))
                 (t '()))
             (set! t (trie:search trie (string #\a #\b #\c)))
             (if (not (= (trie:function t) 20))
                 "abc not found"
                 (let ((t (trie:search-single #\d t)))
                   (if (not (= (trie:function t) 10))
                       "abcd not found"
                       (let ((t (trie:search-single #\e t)))
                         (if (not (= (trie:function t) 30))
                             "abcde not found"
                             (let ((t (trie:search-single #\b trie)))
                               (if (not (= (trie:function t) 40))
                                   "b not found"
                                   (let ((t (trie:search t (string #\e #\g))))
                                     (= (trie:function t) 50))))))))))))))