(module trie

(make-trie
 trie?
 trie-insert!
 trie-ref
 trie-ref*
 trie-value
 trie->list)

(import chicken scheme)
(use srfi-1 data-structures)

(define-record trie children value)

(define %make-trie make-trie)

(define (make-trie)
  (%make-trie (list) (list)))

(define (trie-ref* trie key)
  (alist-ref key (trie-children trie)))

(define (trie-ref trie key #!optional (default (constantly #f)))
  (let loop ((node trie)
             (key key))
    (if (null? key)
        (if (null? (trie-value node))
            (default)
            (car (trie-value node)))
        (let ((child (trie-ref* node (car key))))
          (if child
              (loop child (cdr key))
              (default))))))

(define (add-child! trie key child)
  (trie-children-set!
   trie (alist-cons key child (trie-children trie))))

(define (trie-insert! trie key val)
  (let loop ((node trie)
             (key key))
    (if (null? key)
        (trie-value-set! node (list val))
        (let* ((ckey (car key))
               (child (or (trie-ref* node ckey)
                          (let ((child (make-trie)))
                            (add-child! node ckey child)
                            child))))
          (loop child (cdr key))))))

(define (trie->list trie)
  (cons
   (let loop ((trie trie))
     (map (lambda (child)
            (cons (car child)
                  (trie->list (cdr child))))
          (trie-children trie)))
   (trie-value trie)))

)