Attempting to calculate string-length of an x-expression with minimal allocation

xexpr-len.rkt

#lang racket/base

(require racket/contract/base
         racket/match
         racket/symbol
         txexpr
         xml)

;; Function for calculating the length in characters of an x-expression if it
;; we converted to a string --- without actually converting it into a string
;;
;; Discussion: https://racket.discourse.group/t/x-expression-string-length/550

(provide
 (contract-out [xexpr-string-length (-> xexpr? exact-nonnegative-integer?)]))

;; We need to convert symbols to strings to check their length, but we can
;; avoid most of the cost of doing so by using symbol->immutable-string.

(define (attrs-string-length lst)
  (for/sum ([attr-pair (in-list lst)])
    (define attr-str (symbol->immutable-string (car attr-pair)))
    (+ 4 (string-length attr-str) (string-length (cadr attr-pair)))))

(define (xexpr-string-length x)
  (cond
    [(string? x) (string-length x)]
    [else
     (define-values (tag attrs elems)
       (match x
         [(list tag) (values tag '() '())]
         [(list tag (? list? attrs) children ...)
          (values tag attrs children)]
         [(list tag children ...)
          (values tag '() children)]))
     (define tag-len (string-length (symbol->immutable-string tag)))
     (+ 3 ; </>
        (if (null? elems) 0 (+ 2 tag-len))
        tag-len
        (attrs-string-length attrs)
        (apply + (map xexpr-string-length elems)))]))

;; Validation and benchmarking ~~~~~~~~~~~~~~~~~~~
;;
(module+ test
  (require rackunit)
  
  (define test-txs '( (br)
                      (hr [[class "x"]])
                      (div "Hello")
                      (div [[id "top"]] "Hello")
                      (div [[class "x"]] (br))
                      (div [[class "x"]] (hr [[class "pause"]]) "hello")
                      (p "Hello " (b [[class "excl"]] "world!") " " (em "wow.")) ))
  (for ([test-tx (in-list test-txs)])
    (check-equal? (xexpr-string-length test-tx)
                  (string-length (xexpr->string test-tx))))

  (define naive-time
    (let* ([start (current-inexact-milliseconds)]
           [result (for/sum ([test-tx (in-list test-txs)])
                     (string-length (xexpr->string test-tx)))])
      (- (current-inexact-milliseconds) start)))

  (define hopefully-better
    (let* ([start (current-inexact-milliseconds)]
           [result (for/sum ([test-tx (in-list test-txs)])
                     (xexpr-string-length test-tx))])
      (- (current-inexact-milliseconds) start)))

  ;; After all this, we verify that the clever method...
  ;; is actually faster (whew)
  ;;
  ;; Naive nethod: 24.0 μs
  ;; Clever method: 5.0 μs
  
  (displayln (format "Naive nethod: ~a μs" (ceiling (* naive-time 1000))))
  (displayln (format "Clever method: ~a μs" (ceiling (* hopefully-better 1000)))))