path: root/src/elmord-hashmap.sls

(library (elmord-hashmap)
  (export empty-sparsevec
          sparsevec-ref
          sparsevec-update
          sparsevec-remove

          hashmap
          hashmap?
          empty-hashmap
          hashmap-add
          hashmap-ref
          hashmap-remove
          hashmap-for-each
          debug-print-hashmap)
  (import (rnrs))

(define fx>> fxarithmetic-shift-right)
(define fx<< fxarithmetic-shift-left)

(define-record-type sparsevec
  (fields
    bitmap
    items))

(define empty-sparsevec (make-sparsevec 0 '#()))

(define (sparsevec-position bitmap index)
  (fxbit-count (fxand bitmap (fx- (fx<< 1 index) 1))))

(define (sparsevec-ref vec index default)
  (cond [(fxbit-set? (sparsevec-bitmap vec) index)
         (vector-ref (sparsevec-items vec)
                     (sparsevec-position (sparsevec-bitmap vec) index))]
        [else default]))

(define (vector-transfer! source source-start target target-start source-limit)
  (do ([source-i source-start (fx+ source-i 1)]
       [target-i target-start (fx+ target-i 1)])
      [(fx>=? source-i source-limit) target]
    (vector-set! target target-i (vector-ref source source-i))))

(define (sparsevec-update vec index value)
  (let* ([position (sparsevec-position (sparsevec-bitmap vec) index)]
         [new? (not (fxbit-set? (sparsevec-bitmap vec) index))]
         [old-items (sparsevec-items vec)]
         [old-length (vector-length old-items)]
         [new-items (make-vector (fx+ (vector-length old-items) (if new? 1 0)))]
         [new-bitmap (fxior (sparsevec-bitmap vec) (fx<< 1 index))])
    (vector-transfer! old-items 0 new-items 0 position)
    (vector-set! new-items position value)
    (vector-transfer! old-items (fx+ position (if new? 0 1))
                      new-items (fx+ position 1)
                      old-length)
    (make-sparsevec new-bitmap new-items)))

(define (sparsevec-remove vec index)
  (cond [(fxbit-set? (sparsevec-bitmap vec) index)
         (let* ([position (sparsevec-position (sparsevec-bitmap vec) index)]
                [old-items (sparsevec-items vec)]
                [old-length (vector-length old-items)]
                [new-items (make-vector (fx- old-length 1))]
                [new-bitmap (fxxor (sparsevec-bitmap vec) (fx<< 1 index))])
           (vector-transfer! old-items 0 new-items 0 position)
           (vector-transfer! old-items (fx+ position 1) new-items position old-length)
           (make-sparsevec new-bitmap new-items))]
        [else vec]))


;; ~~~

;; tree ::= sparsevec of subtrees
;;        | list of entries

(define-record-type hashmap
  (fields
    hash-function  ;; datum -> fixnum
    equal-function ;; datum -> bool
    root           ;; tree
    ))

(define-record-type entry
  (fields
    key
    value))

(define chunk-bits 5)
(define key-chunks 5)

(define chunk-mask (fx- (fx<< 1 chunk-bits) 1))

(define empty-hashmap (make-hashmap equal-hash equal? '()))
;;(define empty-hashmap (make-hashmap (lambda (x) x) equal? '()))

(define (hashmap-add hmap key value)
  (make-hashmap
   (hashmap-hash-function hmap)
   (hashmap-equal-function hmap)
   (tree-insert hmap
                (hashmap-root hmap)
                ((hashmap-hash-function hmap) key)
                0
                (make-entry key value))))

(define (tree-insert hmap tree hash depth entry)
  (cond
   [(sparsevec? tree)
    (let* ([index (fxand hash chunk-mask)]
           [subtree (sparsevec-ref tree index '())]
           [hash-rest (fx>> hash chunk-bits)])
      (sparsevec-update tree index
                        (tree-insert hmap subtree hash-rest (fx+ 1 depth) entry)))]
   [(null? tree)
    (list entry)]
   [(list? tree)
    (cond
     [(fx>=? depth key-chunks)
      ;; Reached bottom of the tree: append to collision list.
      ;; Remove any previous entry with the same key.
      (cons entry (filter (lambda (old-entry)
                            (not ((hashmap-equal-function hmap)
                                  (entry-key old-entry)
                                  (entry-key entry))))
                          tree))]
     [else
      ;; Element found but not at the bottom of the tree: turn this
      ;; node into a subtree.
      ;; There can only be more than one element in the collision list
      ;; if we're at the bottom of the tree
      (assert (null? (cdr tree)))
      (let* ([old-entry (car tree)])
        ;; Same key? Replace.
        (cond [((hashmap-equal-function hmap) (entry-key old-entry) (entry-key entry))
               (list entry)]
              [else
               (let* ([old-full-hash ((hashmap-hash-function hmap) (entry-key old-entry))]
                      [old-hash (fx>> old-full-hash (fx* chunk-bits depth))])
                 (tree-insert-conflicting hash entry old-hash old-entry depth))]))])]
   [else (error 'tree-insert "Invalid tree type" tree)]))

(define (tree-insert-conflicting hash1 entry1 hash2 entry2 depth)
  (cond [(fx>=? depth key-chunks) (list entry1 entry2)]
        [else
         (let ([index1 (fxand hash1 chunk-mask)]
               [index2 (fxand hash2 chunk-mask)])
           (cond [(fx=? index1 index2)
                  (sparsevec-update empty-sparsevec index1
                                    (tree-insert-conflicting (fx>> hash1 chunk-bits)
                                                             entry1
                                                             (fx>> hash2 chunk-bits)
                                                             entry2
                                                             (fx+ 1 depth)))]
                 [else
                  (let* ([tree (sparsevec-update empty-sparsevec index1 (list entry1))]
                         [tree (sparsevec-update tree index2 (list entry2))])
                    tree)]))]))

(define (hashmap-ref hmap key default)
  (tree-ref hmap
            (hashmap-root hmap)
            ((hashmap-hash-function hmap) key)
            key
            default))

(define (tree-ref hmap tree hash key default)
  (cond
   [(null? tree) default]
   [(pair? tree) (cond [(exists (lambda (entry)
                                  (and ((hashmap-equal-function hmap) key (entry-key entry))
                                       entry))
                                tree)
                        => entry-value]
                       [else default])]
   [(sparsevec? tree)
    (let ([index (fxand hash chunk-mask)])
      (tree-ref hmap (sparsevec-ref tree index '()) (fx>> hash chunk-bits) key default))]
   [else (error 'tree-ref "Invalid tree" tree)]))


(define (hashmap-remove hmap key)
  (make-hashmap
   (hashmap-hash-function hmap)
   (hashmap-equal-function hmap)
   (tree-remove hmap
                (hashmap-root hmap)
                ((hashmap-hash-function hmap) key)
                key)))

(define (tree-remove hmap tree hash key)
  (cond
   [(list? tree)
    (filter (lambda (entry)
              (not ((hashmap-equal-function hmap) key (entry-key entry))))
            tree)]
   [(sparsevec? tree)
    (let* ([index (fxand hash chunk-mask)]
           [subtree (sparsevec-ref tree index '())])
      (cond
       [(null? subtree) tree]
       [else (let* ([subtree* (tree-remove hmap subtree (fx>> hash chunk-bits) key)]
                    [tree* (if (null? subtree*)
                               (sparsevec-remove tree index)
                             (sparsevec-update tree index subtree*))])
               (if (and (= (vector-length (sparsevec-items tree*)) 1)
                        (list? (vector-ref (sparsevec-items tree*) 0)))
                   ;; If only one subtree remained, and it is a leaf, raise the leaf.
                   (vector-ref (sparsevec-items tree*) 0)
                 tree*))]))]
   [else (error 'tree-remove "Invalid tree" tree)]))


(define (hashmap-for-each proc hmap)
  (tree-for-each proc (hashmap-root hmap)))

(define (tree-for-each proc tree)
  (cond [(list? tree)
         (for-each (lambda (entry)
                     (proc (entry-key entry) (entry-value entry)))
                   tree)]
        [(sparsevec? tree)
         (vector-for-each (lambda (subtree)
                            (tree-for-each proc subtree))
                          (sparsevec-items tree))]
        [else (error 'hashmap-for-each "Invalid tree" tree)]))

(define (print-depth n)
  (do ([i 0 (+ i 1)])
      [(>= i n)]
    (display "    ")))


(define (debug-print-tree tree depth)
  (cond [(sparsevec? tree)
         (newline)
         (do ([i 0 (+ i 1)])
             [(>= i (fx<< 1 chunk-bits))]
           (let ([val (sparsevec-ref tree i '())])
             (when (not (null? val))
               (print-depth depth)
               (display (number->string i 2)) (display ": ")
               (debug-print-tree val (+ 1 depth)))))]
        [else (write tree)
              (newline)]))

(define (debug-print-hashmap hmap)
  (debug-print-tree (hashmap-root hmap) 0))


) ;; end library