Trie Data Structure in R7RS Scheme

Question

Making the transition from R6RS Scheme to R7RS has been very educational and fun. Chibi Scheme is my favorite because of its close compliance with the standard.

When starting a new project, I wanted to re-use a library of Trie data procedures I had developed in Chez Scheme (R6RS). The conversion was painless at first. However, when starting to work on some of the units tests, performance was not great. A test that reads in a large (~172,000) dictionary of words and counts them took hundreds of times longer. It was essentially instant in Chez, adding a barely noticable delay to the beginning of the tests.

Running in convenient R7RS distributions gave much poorer execution times.

Here is a small example. It is a much reduced version of the original library. It loads a dictionary and counts the words in the dictionary. (The file used to build the dictionary was obtained from: https://www.wordgamedictionary.com/enable/)

;;;
;;; trie-problem.scm -- illustrate performance problem with trie library.
;;;

(import (scheme base)
        (scheme file)
        (scheme write)
        ;; Depending on Scheme implementation used, you may have to 
        ;; un-comment one of the following.
        ;;(chibi time)
        ;;(gauche time)
        )

(define alphabet-size 256)

;; In a trie, each node consists of an "end-of-word" marker and
;; and a vector of pointers to child nodes.
(define-record-type trie
  (mk-trie c end-of-word? children)
  trie?
  (c trie-c trie-c-set!)
  (end-of-word? trie-end-of-word? trie-end-of-word-set!)
  (children trie-children trie-children-set!))

;; Return a new, initialized trie node.
(define (trie-new-node)
  (let ((node (mk-trie #\null #f (make-vector alphabet-size '()))))
    node))

;; Return the child node associated with the index into the child vector.
(define (child-for-index node idx)
  (vector-ref (trie-children node) idx))

;; Return the child node associated with the character in the child vector.
(define (child-for-char node c)
  (child-for-index node (char->integer c)))

;; Add a word to the trie node.
(define (trie-add-word! node word)
  (let loop ((crawler node)
             (lst  (string->list word)))
    (if (null? lst)
        (trie-end-of-word-set! crawler #t)
        (let* ((c (car lst))
               (idx (char->integer c)))
          (when (null? (child-for-index crawler idx))
            (let ((new-node (trie-new-node)))
              (trie-c-set! new-node c)
              (vector-set! (trie-children crawler) idx new-node)))
          (loop (child-for-index crawler idx) (cdr lst))))))

;; Convenience function to convert a vector index to the corresponding
;; letter and append it to the string.
(define (append-char-for-index s i)
  (string-append s (string (integer->char i))))

;; Do an in-order traversal of the tree and, for each complete
;; word, invoke the callback function with the word.
(define (trie-process-helper node s callback)
  (when (trie-end-of-word? node)
    (callback s))
  (let loop ((idx 0))
    (when (< idx alphabet-size)
      (let ((child (child-for-index node idx)))
        (when (not (null? child))
          (trie-process-helper child (append-char-for-index s idx) callback))
        (loop (+ 1 idx))))))

;; Return the number of words stored in the trie.
(define (trie-count-words node)
  (let* ((cntr 0)
         (count-fn (lambda (s) (set! cntr (+ 1 cntr)))))
    (trie-process-helper node "" count-fn)
    cntr))

;; A dedicated word-counting procedure. Uses no "helper" function
;; and does not contain code to generalize to additional abilities.
(define (faster-counter node)
  (letrec ((word-count 0)
           (cf (lambda (nd)
                 (when (trie-end-of-word? nd)
                   (set! word-count (+ 1 word-count)))
                 (let loop ((idx 0))
                   (when (< idx alphabet-size)
                     (let ((child (child-for-index nd idx)))
                       (when (not (null? child))
                         (cf child))
                       (loop (+ 1 idx))))))))
    (cf node)
    word-count))

;; Return an ordered list of lines in the named file.
(define (readlines filename)
  (call-with-input-file filename
    (lambda (p)
      (let loop ((line (read-line p))
                 (result '()))
        (if (eof-object? line)
            result
            ;;(reverse result)
            (loop (read-line p) (cons line result)))))))

(let ((root (trie-new-node)))

  ;; Open the dictionary file, containing one word per line, and
  ;; build a trie from the contents.
  (define (trie-build-dictionary! word-list-file-name root-node)
    (display "Enter 'trie-build-dictionary!") (newline)
    (with-input-from-file word-list-file-name
      ;; Process the file by adding words from the
      ;; file, one word per line, to the trie root-node.
      (lambda ()
        (let loop ((line (read-line (current-input-port))))
          (when (not (eof-object? line))
            (trie-add-word! root-node line)
            (loop (read-line (current-input-port)))))))
    (display "Exit 'trie-build-dictionary!") (newline))

  (time "Build dictionary"
        ;; Use the small dictionary file for debugging.
        ;;(trie-build-dictionary! "./word-lists/subset_list.txt" root))
        (trie-build-dictionary! "./word-lists/ENABLE_word_list.txt" root))

  (display "Done loading dictionary. Counting...") (newline)
  (time "Counting words"
        (begin
          (display "Words in trie: ") (display (trie-count-words root))
          (newline)))

    (time "Fast counting words"
        (begin
          (display "Fast count of words: ")
          (display (faster-counter root)) (newline))))

As I mentioned, running this code, which is part of the startup routine of a program written in Chez, is nearly instantaneous. But in Chibi I got:

;; Run with chibi
➜  david in trees chibi-scheme trie-problem.scm
Enter 'trie-build-dictionary!
Exit 'trie-build-dictionary!
Build dictionary: 356168 ms
Done loading dictionary. Counting...
Words in trie: 172820
Counting words: 238451 ms
Fast count of words: 172820
Fast counting words: 30125 ms

(All measurements on a 2017 iMac Pro running macOS 13.3 x86)

I have seen that the developers of Chibi are not as concerned with performance as other aspects of the development effort but I was still surprised at the difference. Six minutes to load the file and build the trie-based dictionary. Four minutes to count the number of words in the dictionary.

This code uses the trie-count-words procedure to do the counting. Like several other procedures in the full library, trie-count-words pushes off the real work to a "helper", trie-process-helper. This procedure takes a couple of argument to specify its behavior. In this case, only one is used. It just increments a counter when a word is found.

Thinking that the generalized processing routine might not be optimal, I added a second procedure dedicated to counting, faster-counter. The results show a significant speed increase, but still 30 seconds to do the counting.

Maybe Chibi is not the best choice? So I tried using a couple other R7RS implementations, Gauche Scheme and Chicken Scheme. Here are some results.

;; Run with the Gauche interpreter
➜  david in trees gosh -r7 trie-problem.scm
Enter 'trie-build-dictionary!
Exit 'trie-build-dictionary!
;(time "Build dictionary" (trie-build-dictionary! "./word-lists/ENABLE_w"...
; real   1.554
; user   2.080
; sys    0.460
Done loading dictionary. Counting...
Words in trie: 172820
;(time "Counting words" (begin (display "Words in trie: ") (display (tri ...
; real   8.902
; user   8.880
; sys    0.030
Fast count of words: 172820
;(time "Fast counting words" (begin (display "Fast count of words: ") (d ...
; real   8.750
; user   8.750
; sys    0.000

;; Run with Chicken interpreter
➜  david in trees csi -R r7rs -q -batch trie-problem.scm
Enter 'trie-build-dictionary!
Exit 'trie-build-dictionary!
5.032s CPU time, 2.217s GC time (major), 1866429/163502 mutations (total/tracked), 12/29010 GCs (major/minor), maximum live heap: 775.96 MiB
Done loading dictionary. Counting...
Words in trie: 172820
57.799s CPU time, 0.182s GC time (major), 16193729/2250046 mutations (total/tracked), 1/786642 GCs (major/minor), maximum live heap: 775.95 MiB
Fast count of words: 172820
56.73s CPU time, 0.184s GC time (major), 15785834/2246632 mutations (total/tracked), 1/784933 GCs (major/minor), maximum live heap: 775.95 MiB

;; Run with Chicken compiler
➜  david in trees csc -R r7rs trie-problem.scm
➜  david in trees ./trie-problem
Enter 'trie-build-dictionary!
Exit 'trie-build-dictionary!
2.584s CPU time, 1.872s GC time (major), 387893/25139 mutations (total/tracked), 10/3388 GCs (major/minor), maximum live heap: 775.67 MiB
Done loading dictionary. Counting...
Words in trie: 172820
4.948s CPU time, 387884/135 mutations (total/tracked), 0/73428 GCs (major/minor), maximum live heap: 775.67 MiB
Fast count of words: 172820
4.773s CPU time, 4/0 mutations (total/tracked), 0/68512 GCs (major/minor), maximum live heap: 337.52 KiB

Again, much better, but still not very good compared to Chez.

There are some weirdnesses though. The "faster" counting procedure hardly made any difference. (Maybe better optimizers?). Also, the counting procedure takes longer than loading and building the dictionary. That doesn't seem right.

The question is if there are better ways to do this? Faster methods that haven't occurred to me?

Thanks for the help.

Answer 1

You'd like to know where the slow-down came from.

                     (let ((child (child-for-index nd idx)))
                       (when (not (null? child))
                         (cf child))
                       (loop (+ 1 idx))))))))

Could we maybe move that (incf idx) so it's not at the end? Perhaps by making idx an argument of the recursive cf call?

I speculate that your R6 compiler cleverly found some code motion to enable TCO while the R7 compilers are, reasonably enough, not jumping through quite as many hoops for you.

Examining the generated assembly code would confirm or reject this hypothesis.

Posting a comparison of per-line CPU profiling would also be helpful.

tl;dr: The "slow" code appears to be recursive rather than tail-recursive.

Answer 2

After reviewing the problem and question some more, it occurred to me that the timings from the Chez implementation had not been shown. I tried to write a version of the problem in Chez to gather those timings.

;; Run with Chez
➜  david in trees chez --script trie-problem.ss
Building dictionary...
(time (trie-build-dictionary! "./word-lists/ENABLE_word_list.txt" ...))
    101 collections
    1.337951594s elapsed cpu time, including 1.218552223s collecting
    1.337968000s elapsed real time, including 1.218811000s collecting
    847145312 bytes allocated, including 40059392 bytes reclaimed

Counting words...
Words in trie: 172820
(time (begin (display "Words in trie: ") ...))
    3 collections
    0.561774181s elapsed cpu time, including 0.008546659s collecting
    0.561770000s elapsed real time, including 0.008564000s collecting
    24809104 bytes allocated, including 23145280 bytes reclaimed

These were slightly slower than expected. Part of the reason is that the alphabet-size used for the example program is twice as large as was used with the program originally written in Chez. So, there went some of the performance.

One of the problems with a Trie data structure as originally implemented is that it takes up so much memory, much of which is usually empty. And the need to specify an alphabet size is ugly.

So I re-implemented the complete library (and the example program shown here) using hash tables from SRFI 69 to store the children. As a result, there is no hard-wired limitation of the alphabet used in the Trie and memory used is dependent on the characters stored, not the fixed length vectors that were used to hold pointers to children.

Here's a version of the example presented in the original question, but using hash tables instead of vectors to store pointers to the children of the current node.

;;;
;;; try-hash-problem.scm -- illustrate performance problem with trie library.
;;; This version is re-worked to use a hash-table as the basis for the
;;; trie children rather than a vector. Also eliminates the need to
;;; hard-wire an alphabet size into the implementation.
;;;

(import (scheme base)
        (scheme file)
        (scheme write)
        (srfi 69))

(cond-expand
 (gauche (import (gauche time)))
 (chibi (import (chibi time)))
 (else))

;; In a trie, each node consists of an "end-of-word" marker and
;; a hash-table of pointers to child nodes.
(define-record-type trie
  (mk-trie c end-of-word? children)
  trie?
  (c trie-c trie-c-set!)
  (end-of-word? trie-end-of-word? trie-end-of-word-set!)
  (children trie-children trie-children-set!))

;; Return a new, initialized trie node.
(define (trie-new-node)
  (let ((node (mk-trie #\null #f (make-hash-table))))
    node))

;; Add a word to the trie.
(define (trie-add-word! trie word)
  (let loop ((node trie)
             (chars (string->list word)))
    (cond ((null? chars)
           (trie-end-of-word-set! node #t))
          (else
           (let* ((c (car chars))
                  (children (trie-children node))
                  (child (hash-table-ref/default children c #f)))
             (unless child
               (set! child (trie-new-node))
               (hash-table-set! children c child))
             (loop child (cdr chars)))))))

;; Return the number of words stored in the trie.
(define (trie-count-words trie)
  (let ((children (trie-children trie)))
    (hash-table-fold
      children
      (lambda (k v count)
        (+ count (trie-count-words v)))
      (if (trie-end-of-word? trie) 1 0))))

;; Open the dictionary file, containing one word per line, and
;; build a trie from the contents.
(define (trie-build-dictionary! word-list-file-name root-node)
  (with-input-from-file word-list-file-name
    ;; Process the file by adding words from the
    ;; file, one word per line, to the trie root-node.
    (lambda ()
      (let loop ((line (read-line (current-input-port))))
        (when (not (eof-object? line))
          (trie-add-word! root-node line)
          (loop (read-line (current-input-port))))))))

;; Do some timing on building a trie from a file of words, then
;; counting the words in the trie.
(let ((root (trie-new-node)))

  (display "Building dictionary...") (newline)
  (time
   ;; Use the small dictionary file for debugging.
   ;;(trie-build-dictionary! "./word-lists/subset_list.txt" root))
   (trie-build-dictionary! "./word-lists/ENABLE_word_list.txt" root))

  (display "\nCounting words...") (newline)
  (time
   (begin
     (display "Words in trie: ")
     (display (trie-count-words root))
     (newline))))

The changes needed were almost trivial, even in the more complete version of the library (many more procedures in the complete library).

Running this program showed substantial improvements in the timings gathered. Of the implementations tested, the Gauche interpreter reported the best performance:

;; Run with Gauche
➜  david in trees gosh -r7 try-hash-problem.scm
Building dictionary...
;(time (trie-build-dictionary! "./word-lists/ENABLE_word_list.txt" root))
; real   0.705
; user   1.890
; sys    0.060

Counting words...
Words in trie: 172820
;(time (begin (display "Words in trie: ") (display (trie-count-words roo ...
; real   0.258
; user   0.590
; sys    0.010

And the example problem is fewer lines too.

Notice also that the trie-count-words procedure has been re-worked to use a type of fold, hash-table-fold from the SRFI.

With these changes, the timings for some of the implementations are coming into the range to make the library useful again.