;; A (mostly) direct port of ClojureScript's PersistentHashMap which
;; is based on Phil Bagwell's Hash Array Mapped Trie and implemented
;; and extended by Rich Hickey.
;;
;; Copyright (c) 2013, Moritz Heidkamp. All rights reserved.
;;
;; The use and distribution terms for this software are covered by the
;; Eclipse Public License 1.0
;; (http://opensource.org/licenses/eclipse-1.0.php) which can be found
;; in the file COPYING at the root of this distribution. By using this
;; software in any fashion, you are agreeing to be bound by the terms
;; of this license.
;;
;; You must not remove this notice, or any other,from this software.

(module persistent-hash-map

(persistent-map
 map?
 transient-map?
 map-add
 map-delete
 map-ref
 map-ref-in
 map-update-in
 map-size
 map-contains?
 map-add!
 map-delete!
 map->transient-map
 persist-map!
 map-reduce
 alist->map
 map->alist
 map-equal?
 map-each
 map-collect
 map-keys
 map-values
 map-merge)

(import (except chicken define-record-type)
        foreign
        scheme)

(use (only extras pp fprintf)
     (only ports with-output-to-string)
     (only srfi-69 equal?-hash)
     (only typed-records define-record-type)
     (rename fast-generic (define-type define-generic-type))
     clojurian-syntax)

(begin-for-syntax (require-library matchable srfi-1))
(import-for-syntax matchable srfi-1)

(define-syntax define-class
  (ir-macro-transformer
   (lambda (exp i c)
     (match exp
       ((_ (name prefix)
           (constructor (fields (? (lambda (o) (c o ':))) types) ...)
           (predicate self)
           (define-form (function args ...) body ...) ...)
        (let* ((getters (map (lambda (f)
                               (i (symbol-append (i prefix) (i f))))
                             fields))
               (setters (map (lambda (g)
                               (i (symbol-append (i g) '-set!)))
                             getters)))
          `(begin
             (define-record-type ,name
               (,constructor . ,fields)
               ,predicate
               . ,(map (lambda (field getter setter type)
                         (list field getter setter (i ':) type))
                       fields getters setters types))
             (define-type ,name (,(i 'struct) ,name))
             (define-generic-type ,name ,predicate)
             ,@(map (lambda (define-form function args body)
                      (cons define-form
                            (cons
                             (cons function
                                   (if (c 'define-generic define-form)
                                       (cons (list name self) args)
                                       (cons self args)))
                             body)))
                    define-form
                    function
                    args
                    body))))))))

(define-type box (struct box))

(define-record-type box
  (box val)
  box?
  (val box-ref box-set!))

(: inode-kv-reduce (vector (* * * -> *) * -> *))
(define (inode-kv-reduce vec f init)
  (let loop ((i 0) (init init))
    (if (< i (vector-length vec))
        (loop (+ i 2)
              (let ((k (vector-ref vec i)))
                (if k
                    (f k (vector-ref vec (+ 1 i)) init)
                    (let ((node (vector-ref vec (+ 1 i))))
                      (if node
                          (-kv-reduce node f init)
                          init)))))
        init)))

(: arithmetic-shift-right (number number --> number))
(define (arithmetic-shift-right n x)
  (arithmetic-shift n (- 0 x)))

(: mask (number number --> number))
(define (mask hash shift)
  (bitwise-and (arithmetic-shift-right hash shift) #x01f))

(: bitpos (number number --> number))
(define (bitpos hash shift)
  (arithmetic-shift 1 (mask hash shift)))

;; TODO: 64bit compat by doubling hex digits and -56 instead of -24
(: bit-count (number --> fixnum))

;; (define (bit-count v)
;;   (let* ((v (- v (bitwise-and (arithmetic-shift v -1) #x55555555)))
;;          (v (+ (bitwise-and v #x33333333) (bitwise-and (arithmetic-shift v -2) #x33333333))))
;;     (arithmetic-shift (* (bitwise-and (+ v (arithmetic-shift v -4)) #xF0F0F0F) #x1010101) -24)))

(define bit-count
  (foreign-lambda int __builtin_popcount unsigned-integer))

(: bitmap-indexed-node-index (number number --> number))
(define (bitmap-indexed-node-index bitmap bit)
  (bit-count (bitwise-and bitmap (- bit 1))))

(: vector-copy* (vector fixnum vector fixnum fixnum -> vector))
(define (vector-copy* from i to j len)
  (let loop ((i i) (j j) (len len))
    (if (zero? len)
        to
        (begin
          (vector-set! to j (vector-ref from i))
          (loop (+ 1 i) (+ 1 j) (- len 1))))))

(: vector-copy (vector -> vector))
(define (vector-copy vec)
  (let ((new (make-vector (vector-length vec))))
    (vector-copy! vec new)
    new))

(: vector-copy-downward (vector fixnum vector fixnum fixnum -> vector))
(define (vector-copy-downward from i to j len)
  (let loop ((i (+ i (- len 1)))
             (j (+ j (- len 1)))
             (len len))
    (if (zero? len)
        to
        (begin
          (vector-set! to j (vector-ref from i))
          (loop (- i 1) (- j 1) (- len 1))))))

(: vector-copy-and-set
   (or (vector fixnum * -> vector)
       (vector fixnum * fixnum * -> vector)))
(define vector-copy-and-set
  (case-lambda
   ((vec i a)
    (doto (vector-copy vec)
          (vector-set! i a)))
   ((vec i a j b)
    (doto (vector-copy vec)
          (vector-set! i a)
          (vector-set! j b)))))

(: key-test (* * -> boolean))
(define key-test equal?)

(: key-hash (* -> fixnum))
(define key-hash equal?-hash)

(: remove-pair (vector fixnum -> vector))
(define (remove-pair vec i)
  (let ((new-vec (make-vector (- (vector-length vec) 2) #f)))
    (vector-copy* vec 0 new-vec 0 (* 2 i))
    (vector-copy* vec (* 2 (+ i 1)) new-vec (* 2 i) (- (vector-length new-vec) (* 2 i)))
    new-vec))

(: hash-collision-node-find-index (vector fixnum * -> fixnum))
(define (hash-collision-node-find-index vec count key)
  (let ((lim (* 2 count)))
    (let loop ((i 0))
      (if (< i lim)
          (if (key-test key (vector-ref vec i))
              i
              (loop (+ i 2)))
          -1))))

(define-class (bitmap-indexed-node bi-)
  (make-bi-node (edit   : boolean)
                (bitmap : number)
                (vec    : vector))
  (bi-node? inode)

  (define-generic (-vec)
    (bi-vec inode))

  (define-generic (-inode-add shift hash key val added-leaf?)
    (let* ((bit (bitpos hash shift))
           (idx (bitmap-indexed-node-index (bi-bitmap inode) bit)))
      (if (zero? (bitwise-and (bi-bitmap inode) bit))
          (let ((n (bit-count (bi-bitmap inode))))
            (if (>= n 16)
                (let* ((nodes (make-vector 32 #f))
                       (jdx   (mask hash shift)))
                  (vector-set! nodes jdx (-inode-add +empty-bitmap-indexed-node+ (+ shift 5) hash key val added-leaf?))
                  (let loop ((i 0) (j 0))
                    (if (< i 32)
                        (if (zero? (bitwise-and (arithmetic-shift-right (bi-bitmap inode) i) 1))
                            (loop (+ 1 i) j)
                            (begin
                              (vector-set! nodes i
                                           (if (vector-ref (bi-vec inode) j)
                                               (-inode-add +empty-bitmap-indexed-node+
                                                           (+ shift 5)
                                                           (key-hash (vector-ref (bi-vec inode) j))
                                                           (vector-ref (bi-vec inode) j)
                                                           (vector-ref (bi-vec inode) (+ 1 j))
                                                           added-leaf?)
                                               (vector-ref (bi-vec inode) (+ 1 j))))
                              (loop (+ 1 i) (+ j 2))))))
                  (make-v-node #f (+ 1 n) nodes))
                (let ((new-vec (make-vector (* 2 (+ 1 n)) #f)))
                  (vector-copy* (bi-vec inode) 0 new-vec 0 (* 2 idx))
                  (vector-set! new-vec (* 2 idx) key)
                  (vector-set! new-vec (+ 1 (* 2 idx)) val)
                  (vector-copy* (bi-vec inode) (* 2 idx) new-vec (* 2 (+ 1 idx)) (* 2 (- n idx)))
                  (box-set! added-leaf? #t)
                  (make-bi-node #f (bitwise-ior (bi-bitmap inode) bit) new-vec))))
          (let ((maybe-key   (vector-ref (bi-vec inode) (* 2 idx)))
                (val-or-node (vector-ref (bi-vec inode) (+ 1 (* 2 idx)))))
            (cond ((not maybe-key)
                   (let ((n (-inode-add val-or-node (+ shift 5) hash key val added-leaf?)))
                     (if (eq? n val-or-node)
                         inode
                         (make-bi-node #f (bi-bitmap inode) (vector-copy-and-set (bi-vec inode) (+ 1 (* 2 idx)) n)))))
                  ((key-test key maybe-key)
                   (if (eq? val val-or-node)
                       inode
                       (make-bi-node #f (bi-bitmap inode) (vector-copy-and-set (bi-vec inode) (+ 1 (* 2 idx)) val))))
                  (else
                   (box-set! added-leaf? #t)
                   (make-bi-node #f (bi-bitmap inode)
                                 (vector-copy-and-set (bi-vec inode)
                                                      (* 2 idx) #f
                                                      (+ 1 (* 2 idx))
                                                      (create-node (+ shift 5) maybe-key val-or-node hash key val)))))))))

  (define-generic (-inode-without shift hash key)
    (let ((bit (bitpos hash shift)))
      (if (zero? (bitwise-and (bi-bitmap inode) bit))
          inode
          (let* ((idx         (bitmap-indexed-node-index (bi-bitmap inode) bit))
                 (maybe-key   (vector-ref (bi-vec inode) (* 2 idx)))
                 (val-or-node (vector-ref (bi-vec inode) (+ 1 (* 2 idx)))))
            (cond ((not maybe-key)
                   (let ((n (-inode-without val-or-node (+ shift 5) hash key)))
                     (cond ((eq? n val-or-node) inode)
                           (n (make-bi-node #f (bi-bitmap inode) (vector-copy-and-set (bi-vec inode) (+ 1 (* 2 idx)) n)))
                           ((= (bi-bitmap inode) bit) #f)
                           (else (make-bi-node #f (bitwise-xor (bi-bitmap inode) bit) (remove-pair (bi-vec inode) idx))))))
                  ((key-test key maybe-key)
                   (make-bi-node #f (bitwise-xor (bi-bitmap inode) bit) (remove-pair (bi-vec inode) idx)))
                  (else inode))))))

  (define-generic (-inode-lookup shift hash key not-found)
    (let ((bit (bitpos hash shift)))
      (if (zero? (bitwise-and (bi-bitmap inode) bit))
          not-found
          (let* ((idx (bitmap-indexed-node-index (bi-bitmap inode) bit))
                 (maybe-key (vector-ref (bi-vec inode) (* 2 idx)))
                 (val-or-node (vector-ref (bi-vec inode) (+ 1 (* 2 idx)))))
            (cond ((not maybe-key) (-inode-lookup val-or-node (+ shift 5) hash key not-found))
                  ((key-test key maybe-key) val-or-node)
                  (else not-found))))))

  ;; (define-generic (-inode-find shift hash key not-found)
  ;;   (let ((bit (bitpos hash shift)))
  ;;     (if (zero? (bitwise-and (bi-bitmap inode) bit))
  ;;         not-found
  ;;         (let* [(idx (bitmap-indexed-node-index (bi-bitmap inode) bit))
  ;;                (maybe-key (vector-ref (bi-vec inode) (* 2 idx)))
  ;;                (val-or-node (vector-ref (bi-vec inode) (+ 1 (* 2 idx))))]
  ;;           (cond ((not maybe-key)
  ;;                  (-inode-find val-or-node (+ shift 5) hash key not-found))
  ;;                 ((key-test key maybe-key)
  ;;                  (cons maybe-key val-or-node))
  ;;                 (else not-found))))))

  ;; (inode-seq [inode]
  ;;            (create-inode-seq (bi-vec inode)))

  (define (bi-ensure-editable e)
    (if (eq? e (bi-edit inode))
        inode
        (let* ((n (bit-count (bi-bitmap inode)))
               (new-vec (make-vector (if (negative? n) 4 (* 2 (+ 1 n))) #f)))
          (vector-copy* (bi-vec inode) 0 new-vec 0 (* 2 n))
          (make-bi-node e (bi-bitmap inode) new-vec))))

  (define-generic (-ensure-editable e)
    (bi-ensure-editable inode e))

  (define (-edit-and-remove-pair e bit i)
    (if (= (bi-bitmap inode) bit)
        #f
        (let* ((editable (bi-ensure-editable inode e))
               (evec     (bi-vec editable))
               (len      (vector-length evec)))
          (bi-bitmap-set! editable (bitwise-xor bit (bi-bitmap editable)))
          (vector-copy* evec (* 2 (+ 1 i))
                        evec (* 2 i)
                        (- len (* 2 (+ 1 i))))
          (vector-set! evec (- len 2) #f)
          (vector-set! evec (- len 1) #f)
          editable)))

  (define-generic (-inode-add! edit shift hash key val added-leaf?)
    (let* ((bit (bitpos hash shift))
           (idx (bitmap-indexed-node-index (bi-bitmap inode) bit)))
      (if (zero? (bitwise-and (bi-bitmap inode) bit))
          (let ((n (bit-count (bi-bitmap inode))))
            (cond
             ((< (* 2 n) (vector-length (bi-vec inode)))
              (let* ((editable (bi-ensure-editable inode edit))
                     (evec     (bi-vec editable)))
                (box-set! added-leaf? #t)
                (vector-copy-downward evec (* 2 idx)
                                      evec (* 2 (+ 1 idx))
                                      (* 2 (- n idx)))
                (vector-set! evec (* 2 idx) key)
                (vector-set! evec (+ 1 (* 2 idx)) val)
                (bi-bitmap-set! editable (bitwise-ior (bi-bitmap editable) bit))
                editable))
             
             ((>= n 16)
              (let ((nodes (make-vector 32 #f))
                    (jdx   (mask hash shift)))
                (vector-set! nodes jdx (-inode-add! +empty-bitmap-indexed-node+ edit (+ shift 5) hash key val added-leaf?))
                (let loop ((i 0) (j 0))
                  (if (< i 32)
                      (if (zero? (bitwise-and (arithmetic-shift-right (bi-bitmap inode) i) 1))
                          (loop (+ 1 i) j)
                          (begin
                            (vector-set! nodes i
                                         (if (vector-ref (bi-vec inode) j)
                                             (-inode-add! +empty-bitmap-indexed-node+
                                                          edit
                                                          (+ shift 5)
                                                          (key-hash (vector-ref (bi-vec inode) j))
                                                          (vector-ref (bi-vec inode) j)
                                                          (vector-ref (bi-vec inode) (+ 1 j))
                                                          added-leaf?)
                                             (vector-ref (bi-vec inode) (+ 1 j))))
                            (loop (+ 1 i) (+ j 2))))))
                (make-v-node edit (+ 1 n) nodes)))

             (else
              (let ((new-vec (make-vector (* 2 (+ n 4)) #f)))
                (vector-copy* (bi-vec inode) 0 new-vec 0 (* 2 idx))
                (vector-set! new-vec (* 2 idx) key)
                (vector-set! new-vec (+ 1 (* 2 idx)) val)
                (vector-copy* (bi-vec inode) (* 2 idx) new-vec (* 2 (+ 1 idx)) (* 2 (- n idx)))
                (box-set! added-leaf? #t)
                (let ((editable (bi-ensure-editable inode edit)))
                  (bi-vec-set! editable new-vec)
                  (bi-bitmap-set! editable (bitwise-ior (bi-bitmap editable) bit))
                  editable)))))
          (let ((maybe-key (vector-ref (bi-vec inode) (* 2 idx)))
                (val-or-node (vector-ref (bi-vec inode) (+ 1 (* 2 idx)))))
            (cond ((not maybe-key)
                   (let ((n (-inode-add! val-or-node edit (+ shift 5) hash key val added-leaf?)))
                     (if (eq? n val-or-node)
                         inode
                         (edit-and-set inode edit (+ 1 (* 2 idx)) n))))

                  ((key-test key maybe-key)
                   (if (eq? val val-or-node)
                       inode
                       (edit-and-set inode edit (+ 1 (* 2 idx)) val)))

                  (else
                   (box-set! added-leaf? #t)
                   (edit-and-set inode edit (* 2 idx) #f (+ 1 (* 2 idx))
                                 (create-node edit (+ shift 5) maybe-key val-or-node hash key val))))))))

  (define-generic (-inode-without! edit shift hash key removed-leaf?)
    (let ((bit (bitpos hash shift)))
      (if (zero? (bitwise-and (bi-bitmap inode) bit))
          inode
          (let* ((idx (bitmap-indexed-node-index (bi-bitmap inode) bit))
                 (maybe-key (vector-ref (bi-vec inode) (* 2 idx)))
                 (val-or-node (vector-ref (bi-vec inode) (+ 1 (* 2 idx)))))
            (cond ((not maybe-key)
                   (let ((node (-inode-without! val-or-node edit (+ shift 5) hash key removed-leaf?)))
                     (cond ((eq? node val-or-node)
                            inode)
                           (node
                            (edit-and-set inode edit (+ 1 (* 2 idx)) node))
                           ((= (bi-bitmap inode) bit)
                            #f)
                           (else
                            (-edit-and-remove-pair inode edit bit idx)))))
                  ((key-test key maybe-key)
                   (box-set! removed-leaf? #t)
                   (-edit-and-remove-pair inode edit bit idx))
                  (else inode))))))

  (define-generic (-kv-reduce f init)
    (inode-kv-reduce (bi-vec inode) f init)))

(: -edit-and-remove-pair (bitmap-indexed-node boolean number fixnum -> bitmap-indexed-node))

(define-record-printer (bitmap-indexed-node vnode out)
  (fprintf out "#<bitmap-indexed-node ~S>" (-vec vnode)))

(define +empty-bitmap-indexed-node+
  (the bitmap-indexed-node (make-bi-node #f 0 (vector))))

(define-class (hash-collision-node hc-)
  (make-hc-node (edit           : boolean)
                (collision-hash : number)
                (count          : number)
                (vec            : vector))
  (hc-node? inode)

  (define-generic (-vec)
    (hc-vec inode))

  (define-generic (-inode-add shift hash key val added-leaf?)
    (if (= hash (hc-collision-hash inode))
        (let ((idx (hash-collision-node-find-index (hc-vec inode) (hc-count inode) key)))
          (if (= idx -1)
              (let* ((len (vector-length (hc-vec inode)))
                     (new-vec (make-vector (+ len 2) #f)))
                (vector-copy* (hc-vec inode) 0 new-vec 0 len)
                (vector-set! new-vec len key)
                (vector-set! new-vec (+ 1 len) val)
                (box-set! added-leaf? #t)
                (make-hc-node #f (hc-collision-hash inode) (+ 1 (hc-count inode)) new-vec))
              (if (= (vector-ref (hc-vec inode) idx) val)
                  inode
                  (make-hc-node #f (hc-collision-hash inode) (hc-count inode) (vector-copy-and-set (hc-vec inode) (+ 1 idx) val)))))
        (-inode-add (make-bi-node #f (bitpos (hc-collision-hash inode) shift) (vector #f inode))
                    shift hash key val added-leaf?)))

  (define-generic (-inode-without shift hash key)
    (let ((idx (hash-collision-node-find-index (hc-vec inode) (hc-count inode) key)))
      (cond ((= idx -1) inode)
            ((= (hc-count inode) 1) #f)
            (else (make-hc-node #f (hc-collision-hash inode) (- (hc-count inode) 1) (remove-pair (hc-vec inode) (quotient idx 2)))))))

  (define-generic (-inode-lookup shift hash key not-found)
    (let ((idx (hash-collision-node-find-index (hc-vec inode) (hc-count inode) key)))
      (cond ((< idx 0) not-found)
            ((key-test key (vector-ref (hc-vec inode) idx))
             (vector-ref (hc-vec inode) (+ 1 idx)))
            (else not-found))))

  ;; (define-generic (-inode-find shift hash key not-found)
  ;;   (let ((idx (hash-collision-node-find-index (hc-vec inode) (hc-count inode) key)))
  ;;     (cond ((< idx 0) not-found)
  ;;           ((key-test key (vector-ref (hc-vec inode) idx))
  ;;            (cons (vector-ref (hc-vec inode) idx) (vector-ref (hc-vec inode) (+ 1 idx))))
  ;;           (else not-found))))

  ;; (inode-seq [inode]
  ;;            (create-inode-seq (hc-vec inode)))

  (define (hc-ensure-editable e)
    (if (eq? e (hc-edit inode))
        inode
        (let ((new-vec (make-vector (* 2 (+ 1 (hc-count inode))) #f)))
          (vector-copy* (hc-vec inode) 0 new-vec 0 (* 2 (hc-count inode)))
          (make-hc-node e (hc-collision-hash inode) (hc-count inode) new-vec))))

  (define-generic (-ensure-editable e)
    (hc-ensure-editable inode e))

  (define (-ensure-editable-vector e count* vec*)
    (if (eq? e (hc-edit inode))
        (begin
          (hc-vec-set! inode vec*)
          (hc-count-set! inode count*)
          inode)
        (make-hc-node (hc-edit inode) (hc-collision-hash inode) count* vec*)))

  (define-generic (-inode-add! edit shift hash key val added-leaf?)
    (if (= hash (hc-collision-hash inode))
        (let ((idx (hash-collision-node-find-index (hc-vec inode) (hc-count inode) key)))
          (if (= idx -1)
              (if (> (vector-length (hc-vec inode)) (* 2 (hc-count inode)))
                  (let ((editable (edit-and-set inode edit (* 2 (hc-count inode)) key (+ 1 (* 2 (hc-count inode))) val)))
                    (box-set! added-leaf? #t)
                    (hc-count-set! editable (+ 1 (hc-count editable)))
                    editable)
                  (let* ((len (vector-length (hc-vec inode)))
                         (new-vec (make-vector (+ len 2) #f)))
                    (vector-copy* (hc-vec inode) 0 new-vec 0 len)
                    (vector-set! new-vec len key)
                    (vector-set! new-vec (+ 1 len) val)
                    (box-set! added-leaf? #t)
                    (-ensure-editable-vector inode edit (+ 1 (hc-count inode)) new-vec)))
              (if (eq? (vector-ref (hc-vec inode) (+ 1 idx)) val)
                  inode
                  (edit-and-set inode edit (+ 1 idx) val))))
        (-inode-add! (make-bi-node edit (bitpos (hc-collision-hash inode) shift) (vector #f inode #f #f))
                     edit shift hash key val added-leaf?)))

  (define-generic (-inode-without! edit shift hash key removed-leaf?)
    (let ((idx (hash-collision-node-find-index (hc-vec inode) (hc-count inode) key)))
      (if (= idx -1)
          inode
          (begin
            (box-set! removed-leaf? #t)
            (and (not (= (hc-count inode) 1))
                 (let* ((editable (hc-ensure-editable inode edit))
                        (evec     (hc-vec editable))
                        (count    (hc-count inode)))
                   (vector-set! evec idx (vector-ref evec (- (* 2 count) 2)))
                   (vector-set! evec (+ 1 idx) (vector-ref evec (- (* 2 count) 1)))
                   (vector-set! evec (- (* 2 count) 1) #f)
                   (vector-set! evec (- (* 2 count) 2) #f)
                   (hc-count-set! editable (- (hc-count editable) 1))
                   editable))))))

  (define-generic (-kv-reduce f init)
    (inode-kv-reduce (hc-vec inode) f init)))

(: -ensure-editable-vector (hash-collision-node boolean number vector -> hash-collision-node))

(define-record-printer (hash-collision-node hc-node out)
  (fprintf out "#<hash-collision-node ~S>" (hc-vec hc-node)))

;; type declaration below
(define create-node
  (case-lambda
   ((shift key1 val1 key2hash key2 val2)
    (let ((key1hash (key-hash key1)))
      (if (= key1hash key2hash)
          (make-hc-node #f key1hash 2 (vector key1 val1 key2 val2))
          (let ((added-leaf? (box #f)))
            (-> +empty-bitmap-indexed-node+
                (-inode-add shift key1hash key1 val1 added-leaf?)
                (-inode-add shift key2hash key2 val2 added-leaf?))))))
   ((edit shift key1 val1 key2hash key2 val2)
    (let ((key1hash (key-hash key1)))
      (if (= key1hash key2hash)
          (make-hc-node #f key1hash 2 (vector key1 val1 key2 val2))
          (let ((added-leaf? (box #f)))
            (-> +empty-bitmap-indexed-node+
                (-inode-add! edit shift key1hash key1 val1 added-leaf?)
                (-inode-add! edit shift key2hash key2 val2 added-leaf?))))))))

(define-class (vector-node v-)
  (make-v-node (edit  : boolean)
               (count : number)
               (vec   : vector))
  (v-node? inode)

  (define-generic (-vec)
    (v-vec inode))

  (define (pack-vector-node edit idx)
    (let* ((vec     (v-vec inode))
           (len     (* 2 (- (v-count inode) 1)))
           (new-vec (make-vector len #f)))
      (let loop ((i 0) (j 1) (bitmap 0))
        (if (< i len)
            (if (and (not (= i idx))
                     (vector-ref vec i))
                (begin
                  (vector-set! new-vec j (vector-ref vec i))
                  (loop (+ 1 i) (+ j 2) (bitwise-ior bitmap (arithmetic-shift 1 i))))
                (loop (+ 1 i) j bitmap))
            (make-bi-node edit bitmap new-vec)))))

  (define-generic (-inode-add shift hash key val added-leaf?)
    (let* ((idx (mask hash shift))
           (node (vector-ref (v-vec inode) idx)))
      (if (not node)
          (->> (-inode-add +empty-bitmap-indexed-node+ (+ shift 5) hash key val added-leaf?)
               (vector-copy-and-set (v-vec inode) idx)
               (make-v-node #f (+ 1 (v-count inode))))
          (let ((n (-inode-add node (+ shift 5) hash key val added-leaf?)))
            (if (eq? n node)
                inode
                (make-v-node #f (v-count inode) (vector-copy-and-set (v-vec inode) idx n)))))))

  (define-generic (-inode-without shift hash key)
    (let* ((idx (mask hash shift))
           (node (vector-ref (v-vec inode) idx)))
      (if (not node)
          inode
          (let ((n (-inode-without node (+ shift 5) hash key)))
            (cond
             ((eq? n node)
              inode)
             ((not n)
              (if (<= (v-count inode) 8)
                  (pack-vector-node inode #f idx)
                  (make-v-node #f (- (v-count inode) 1) (vector-copy-and-set (v-vec inode) idx n))))
             (else
              (make-v-node #f (v-count inode) (vector-copy-and-set (v-vec inode) idx n))))))))

  (define-generic (-inode-lookup shift hash key not-found)
    (let* ((idx (mask hash shift))
           (node (vector-ref (v-vec inode) idx)))
      (if (not node)
          not-found
          (-inode-lookup node (+ shift 5) hash key not-found))))

  ;; (define-generic (-inode-find shift hash key not-found)
  ;;   (let* ((idx (mask hash shift))
  ;;          (node (vector-ref (v-vec inode) idx)))
  ;;     (if (not node)
  ;;         not-found
  ;;         (-inode-find node (+ shift 5) hash key not-found))))

  ;; (inode-seq [inode]
  ;;            (create-vector-node-seq (v-vec inode)))

  (define-generic (-ensure-editable e)
    (if (eq? e (v-edit inode))
        inode
        (make-v-node e (v-count inode) (vector-copy (v-vec inode)))))

  (define-generic (-inode-add! edit shift hash key val added-leaf?)
    (let* ((idx  (mask hash shift))
           (node (vector-ref (v-vec inode) idx)))
      (if (not node)
          (let ((editable (edit-and-set inode edit idx (-inode-add! +empty-bitmap-indexed-node+ edit (+ shift 5) hash key val added-leaf?))))
            (v-count-set! editable (+ 1 (v-count editable)))
            editable)
          (let ((n (-inode-add! node edit (+ shift 5) hash key val added-leaf?)))
            (if (eq? n node)
                inode
                (edit-and-set inode edit idx n))))))

  (define-generic (-inode-without! edit shift hash key removed-leaf?)
    (let* ((idx  (mask hash shift))
           (node (vector-ref (v-vec inode) idx)))
      (if (not node)
          inode
          (let ((n (-inode-without! node edit (+ shift 5) hash key removed-leaf?)))
            (cond
             ((eq? n node)
              inode)
             ((not n)
              (if (<= (v-count inode) 8)
                  (pack-vector-node inode edit idx)
                  (let ((editable (edit-and-set inode edit idx n)))
                    (v-count-set! editable (- (v-count editable) 1))
                    editable)))
             (else
              (edit-and-set inode edit idx n)))))))

  (define-generic (-kv-reduce f init)
    (let* ((vec (v-vec inode))
           (len (vector-length vec)))
      (let loop ((i 0) (init init))
        (if (< i len)
            (loop (+ 1 i)
                  (let ((node (vector-ref vec i)))
                    (if node (-kv-reduce node f init) init)))
            init)))))


;; type declaration below
(define edit-and-set
  (case-lambda
   ((inode edit i a)
    (let ((editable (-ensure-editable inode edit)))
      (vector-set! (-vec editable) i a)
      editable))
   ((inode edit i a j b)
    (let* ((editable (-ensure-editable inode edit))
           (vec (-vec editable)))
      (vector-set! vec i a)
      (vector-set! vec j b)
      editable))))


(: pack-vector-node (vector-node boolean fixnum -> bitmap-indexed-node))

(define-record-printer (vector-node vnode out)
  (fprintf out "#<vector-node ~S>" (-vec vnode)))

(define-type node
  (or bitmap-indexed-node
      vector-node
      hash-collision-node
      boolean))

(: edit-and-set
   (or (node boolean fixnum * -> node)
       (node boolean fixnum * fixnum * -> node)))

(: create-node
   (or (fixnum * * fixnum * * -> node)
       (node fixnum * * fixnum * * -> node)))

(: bi-ensure-editable (bitmap-indexed-node boolean -> bitmap-indexed-node))
(: hc-ensure-editable (hash-collision-node boolean -> hash-collision-node))
(: -ensure-editable (forall (node* node) (node* boolean -> node*)))
(: -vec (node -> vector))
(: -kv-reduce (node (* * * -> *) * -> *))
(: -inode-add (node fixnum fixnum * * box -> node))
(: -inode-without (node fixnum fixnum * -> node))
(: -inode-lookup (node fixnum fixnum * * -> *))
(: -inode-without! (node node fixnum fixnum * box -> node))
(: -inode-add! (node node fixnum fixnum * * box -> node))
(: -inode-lookup (node fixnum fixnum * * -> *))

(define-class (transient-hash-map tm-)
  (make-transient-map (edit       : boolean)
                      (root       : node)
                      (count      : number)
                      (has-false? : boolean)
                      (false-val  : *))
  (transient-map? self)

  (define-generic (-count)
    (if (tm-edit self)
        (tm-count self)
        (error "map-size after persist-map!")))

  (define (add! k v)
    (if (tm-edit self)
        (if (not k)
            (begin
              (unless (eq? (tm-false-val self) v)
                (tm-false-val-set! self v))
              (unless (tm-has-false? self)
                (tm-count-set! self (+ 1 (tm-count self)))
                (tm-has-false?-set! self #t))
              self)
            (let* ((added-leaf? (box #f))
                   (node (-> (if (not (tm-root self))
                                 +empty-bitmap-indexed-node+
                                 (tm-root self))
                             (-inode-add! (tm-edit self) 0 (key-hash k) k v added-leaf?))))
              (unless (eq? node (tm-root self))
                (tm-root-set! self node))
              (when (box-ref added-leaf?)
                (tm-count-set! self (+ 1 (tm-count self))))
              self))
        (error "map-add! after persist-map!")))

  (define (delete! key)
    (if (tm-edit self)
        (if (not key)
            (if (tm-has-false? self)
                (begin
                  (tm-has-false?-set! self #f)
                  (tm-false-val-set! self #f)
                  (tm-count-set! self (- (tm-count self) 1))
                  self)
                self)
            (if (not (tm-root self))
                self
                (let* ((removed-leaf? (box #f))
                       (node (-inode-without! (tm-root self)
                                              (tm-edit self)
                                              0
                                              (key-hash key)
                                              key
                                              removed-leaf?)))
                  (unless (eq? node (tm-root self))
                    (tm-root-set! self node))
                  (when (box-ref removed-leaf?)
                    (tm-count-set! self (- (tm-count self) 1)))
                  self)))
        (error "map-delete! after persist-map!")))

  (define-generic (persist-map!)
    (if (tm-edit self)
        (begin
          (tm-edit-set! self #f)
          (make-map (tm-count self) (tm-root self) (tm-has-false? self) (tm-false-val self)))
        (error "persist-map! called twice")))

  (define-generic (-lookup k not-found)
    (if (not k)
        (if (tm-has-false? self)
            (tm-false-val self)
            not-found)
        (if (not (tm-root self))
            not-found
            (-inode-lookup (tm-root self) 0 (key-hash k) k not-found)))))

(define-class (persistent-hash-map pm-)
  (make-map (count      : number)
            (root       : node)
            (has-false? : boolean)
            (false-val  : *))
  (map? self)

  (define-generic (-count)
    (pm-count self))

  (define-generic (-lookup k not-found)
    (cond ((not k)
           (if (pm-has-false? self)
               (pm-false-val self)
               not-found))
          ((not (pm-root self))
           not-found)
          (else
           (-inode-lookup (pm-root self) 0 (key-hash k) k not-found))))

  (define (-add k v)
    (if (not k)
        (if (and (pm-has-false? self) (eq? v (pm-false-val self)))
            self
            (make-map (if (pm-has-false? self)
                          (pm-count self)
                          (+ 1 (pm-count self)))
                      (pm-root self)
                      #t
                      v))
        (let* ((added-leaf? (box #f))
               (new-root    (-> (or (pm-root self) +empty-bitmap-indexed-node+)
                                (-inode-add 0 (key-hash k) k v added-leaf?))))
          (if (eq? new-root (pm-root self))
              self
              (make-map (if (box-ref added-leaf?)
                            (+ 1 (pm-count self))
                            (pm-count self))
                        new-root
                        (pm-has-false? self)
                        (pm-false-val self))))))

  (define (-delete key)
    (cond ((not key)
           (if (pm-has-false? self)
               (make-map (- (pm-count self) 1) (pm-root self) #f #f)
               self))
          ((not (pm-root self))
           self)
          (else
           (let ((new-root (-inode-without (pm-root self) 0 (key-hash key) key)))
             (if (eq? new-root (pm-root self))
                 self
                 (make-map (- (pm-count self) 1)
                           new-root
                           (pm-has-false? self)
                           (pm-false-val self)))))))

  (define (kv-reduce f init)
    (let ((init (if (pm-has-false? self)
                    (f #f (pm-false-val self) init)
                    init)))
      (if (pm-root self)
          (-kv-reduce (pm-root self) f init)
          init)))

  (define (map->transient-map)
    (make-transient-map #t
                        (pm-root self)
                        (pm-count self)
                        (pm-has-false? self)
                        (pm-false-val self))))

(: map->transient-map (persistent-hash-map -> transient-hash-map))
(: persist-map! (transient-hash-map -> persistent-hash-map))

(define-type hash-map
  (or persistent-hash-map transient-hash-map))

(define +empty-persistent-hash-map+
  (the persistent-hash-map (make-map 0 #f #f #f)))

(define (->map in get-key get-value get-rest)
  (let loop ((in in)
             (out (map->transient-map +empty-persistent-hash-map+)))
    (if (null? in)
        (persist-map! out)
        (loop (get-rest in)
              (add! out (get-key in) (get-value in))))))

(: persistent-map (#!rest * -> persistent-hash-map))
(define (persistent-map . keyvals)
  (->map keyvals car cadr cddr))

(: alist->map ((list-of pair) -> persistent-hash-map))
(define (alist->map alist)
  (->map alist caar cdar cdr))

(define +not-found+
  (list 'not-found))

(: map-ref (hash-map * #!optional * -> *))
(define (map-ref map key #!optional not-found)
  (-lookup map key not-found))

(: map-ref-in (hash-map list #!optional * -> *))
(define (map-ref-in map keys #!optional not-found)
  (if (null? keys)
      map
      (let ((val (map-ref map (car keys) +not-found+)))
        (cond
         ((eq? +not-found+ val)
          not-found)
         ((null? (cdr keys))
          val)
         (else
          (map-ref-in val (cdr keys) not-found))))))

(define (map-keyvals-proc proc)
  (rec (self map key val . more)
       (let ((res (proc map key val)))
         (if (null? more)
             res
             (apply self res more)))))

(define (map-keys-proc proc)
  (lambda (map . keys)
    (foldl proc map keys)))

(: map-add (persistent-hash-map * * #!rest * -> persistent-hash-map))
(define map-add
  (map-keyvals-proc -add))

(: map-delete (persistent-hash-map * #!rest * -> persistent-hash-map))
(define map-delete
  (map-keys-proc -delete))

(: map-add! (transient-hash-map * * #!rest * -> transient-hash-map))
(define map-add!
  (map-keyvals-proc add!))

(: map-delete! (transient-hash-map * #!rest * -> transient-hash-map))
(define map-delete!
  (map-keys-proc delete!))

(: map-size (hash-map -> number))
(define map-size
  -count)

(: map-contains? (hash-map * -> boolean))
(define (map-contains? map key)
  (not (eq? +not-found+ (map-ref map key +not-found+))))

(: map-reduce ((* * * -> *) * hash-map -> *))
(define (map-reduce f init map)
  (kv-reduce map f init))

(: map-collect ((* * -> *) hash-map -> list))
(define (map-collect proc map)
  (map-reduce (lambda (k v r) (cons (proc k v) r)) '() map))

(: map-each ((* * -> *) hash-map ->))
(define (map-each proc map)
  (map-reduce (lambda (k v r) (proc k v) r) (void) map)
  (values))

(: map->alist (hash-map -> (list-of pair)))
(define (map->alist map)
  (map-collect cons map))

(: map-equal? (hash-map hash-map -> boolean))
(define (map-equal? x y)
  (and (= (map-size x) (map-size y))
       (call/cc
        (lambda (return)
          (map-reduce
           (lambda (key val result)
             (or ((if (map? val) map-equal? equal?)
                  (map-ref y key +not-found+)
                  val)
                 (return #f)))
           #t
           x)))))

(: map-keys (hash-map -> list))
(define (map-keys map)
  (map-collect (lambda (k v) k) map))

(: map-values (hash-map -> list))
(define (map-values map)
  (map-collect (lambda (k v) v) map))

(: map-merge (persistent-hash-map persistent-hash-map -> persistent-hash-map))
(define (map-merge m1 m2)
  (persist-map!
   (map-reduce (lambda (k v m)
                 (map-add! m k v))
               (map->transient-map m1)
               m2)))

(: map-update-in (persistent-hash-map list (* #!rest * -> *) #!rest * -> persistent-hash-map))
(define (map-update-in map keys proc . args)
  (when (null? keys)
    (error 'map-update-in "Need at least one key"))
  (let ((key  (car keys)))
    (if (null? (cdr keys))
        (map-add map key (apply proc (map-ref map key) args))
        (map-add map key (apply map-update-in
                                (map-ref map key +empty-persistent-hash-map+)
                                (cdr keys)
                                proc
                                args)))))

(define-record-printer (persistent-hash-map m out)
  (display "#<persistent-hash-map" out)
  (display " " out)
  (let ((s (with-output-to-string (lambda () (pp (map->alist m))))))
    (display (substring s 1 (- (string-length s) 2)) out))
  (display ">" out))

(define-record-printer (transient-hash-map m out)
  (fprintf out "#<transient-hash-map (~S)>" (map-size m)))

)