(define-library (csc hash-map)
  (export
    alist->map
    map-for-each
    hash-bytevector
    insert
    key-not-found-error?
    lookup
    make-map
    map->alist
    map?
    merge)
  (import (scheme base)
          (only (csc format) sprintf)
          (only (csc match) match))
  (begin


    (define-record-type <key-hash>
      (make-key-hash k hash)
      key-hash?
      (k key-hash-value)
      (hash key-hash-hash))


    (define (key-hash<? k1 k2 cmp)
      (define h1 (key-hash-hash k1))
      (define h2 (key-hash-hash k2))
      (cond
        ((< h1 h2)
          #t)
        ((> h1 h2)
          #f)
        (else
          (cmp (key-hash-value k1) (key-hash-value k2)))))


    (define (key-hash=? k1 k2 cmp)
      (define v1 (key-hash-value k1))
      (define v2 (key-hash-value k2))
      (and (= (key-hash-hash k1) (key-hash-hash k2))
           (not (cmp v1 v2))
           (not (cmp v2 v1))))


    (define-record-type <node>
      (make-node color key-hash val left right)
      node?
      (color node-color)
      (key-hash node-key)
      (val node-value)
      (left node-left)
      (right node-right))


    (define (red? n)
      (and (not (null? n))
           (eq? 'red (node-color n))))


    (define (black? n)
      (or (null? n)
          (eq? 'black (node-color n))))


    (define (rebalance-left m)
      (define p (node-left m))
      (define u (node-right m))
      (cond
        ((or (and (red? p)
                  (red? (node-left p))
                  (red? u))
             (and (red? p)
                  (red? (node-right p))
                  (red? u)))
          ;     b          r
          ;    / \        / \
          ;   r   r =>   b   b
          ;  /          /
          ; r          r

          ;   b        r
          ;  / \      / \
          ; r   r => b   b
          ;  \        \
          ;   r        r
          (make-node 'red (node-key m) (node-value m)
                     (make-node 'black (node-key p) (node-value p) (node-left p) (node-right p))
                     (make-node 'black (node-key u) (node-value u) (node-left u) (node-right u))))
        ((and (red? p)
              (red? (node-right p))
              (black? u))
          ;   b        b
          ;  / \      / \
          ; r   b => r   r
          ;  \            \
          ;   r            b
          (let ((n (node-right p)))
            (make-node 'black (node-key n) (node-value n)
                       (make-node 'red (node-key p) (node-value p) (node-left p) (node-left n))
                       (make-node 'red (node-key m) (node-value m) (node-right n) u))))
        ((and (red? p)
              (red? (node-left p))
              (black? u))
          ;     b        b
          ;    / \      / \
          ;   r   b => r   r
          ;  /              \
          ; r                b
          (make-node 'black (node-key p) (node-value p)
                     (node-left p)
                     (make-node 'red (node-key m) (node-value m) (node-right p) u)))
        (else m)))


    (define (rebalance-right m)
      (define u (node-left m))
      (define p (node-right m))
      (cond
        ((or (and (red? u)
                  (red? p)
                  (red? (node-left p)))
             (and (red? u)
                  (red? p)
                  (red? (node-right p))))
          ;   b          r
          ;  / \        / \
          ; r   r   => b   b
          ;      \          \
          ;       r          r

          ;   b        r
          ;  / \      / \
          ; r   r => b   b
          ;    /        /
          ;   r        r
          (make-node 'red (node-key m) (node-value m)
                     (make-node 'black (node-key u) (node-value u) (node-left u) (node-right u))
                     (make-node 'black (node-key p) (node-value p) (node-left p) (node-right p))))
        ((and (black? u)
              (red? p)
              (red? (node-left p)))
          ;   b          b
          ;  / \        / \
          ; b   r =>   r   r
          ;    /      /
          ;   r      b
          (let ((n (node-left p)))
            (make-node 'black (node-key n) (node-value n)
                       (make-node 'red (node-key m) (node-value m) u (node-left n))
                       (make-node 'red (node-key p) (node-value p) (node-right n) (node-right p)))))
        ((and (black? u)
              (red? p)
              (red? (node-right p)))
          ;   b            b
          ;  / \          / \
          ; b   r   =>   r   r
          ;      \      /
          ;       r    b
          (make-node 'black (node-key p) (node-value p)
                     (make-node 'red (node-key m) (node-value m) u (node-left p))
                     (node-right p)))
        (else m)))


    (define (insert-node m k v cmp)
      (cond ((null? m) (make-node 'red k v '() '()))
            ((key-hash<? k (node-key m) cmp)
             (rebalance-left
               (make-node (node-color m) (node-key m) (node-value m)
                          (insert-node (node-left m) k v cmp)
                          (node-right m))))
            ((key-hash=? k (node-key m) cmp) (make-node (node-color m) k v (node-left m) (node-right m)))
            (else
              (rebalance-right
                (make-node (node-color m) (node-key m) (node-value m)
                           (node-left m)
                           (insert-node (node-right m) k v cmp))))))


    (define-record-type <map>
      (construct-map hash cmp root)
      map?
      (hash map-raw-hash)
      (cmp map-cmp)
      (root map-root))


    (define (make-map hash cmp)
      (construct-map hash cmp '()))


    (define (shuffle n)
      (truncate-remainder
        (* #x9e3779b97f4a7c55 n)
        #x10000000000000000))


    (define (map-hash m)
      (lambda (k)
        (shuffle ((map-raw-hash m) k))))


    (define (insert m k v)
      (let ((res (insert-node (map-root m) (make-key-hash k ((map-hash m) k)) v (map-cmp m))))
        (construct-map
          (map-raw-hash m)
          (map-cmp m)
          (make-node 'black (node-key res) (node-value res) (node-left res) (node-right res)))))


    (define-record-type <key-not-found-error>
      (make-key-not-found-error)
      key-not-found-error?)


    (define (lookup m k)
      (letrec ((k* (make-key-hash k ((map-hash m) k)))
               (lookup
                 (lambda (n)
                   (cond ((null? n) (raise (make-key-not-found-error)))
                         ((key-hash<? k* (node-key n) (map-cmp m)) (lookup (node-left n)))
                         ((key-hash=? k* (node-key n) (map-cmp m)) (node-value n))
                         (else (lookup (node-right n)))))))
              (lookup (map-root m))))


    (define (map-for-each f m)
      (letrec ((node-foreach
                 (lambda (n)
                   (unless (null? n)
                     (node-foreach (node-left n))
                     (f (key-hash-value (node-key n)) (node-value n))
                     (node-foreach (node-right n))))))
              (node-foreach (map-root m))))


    (define (map->alist m)
      (let ((alist '()))
        (map-for-each
          (lambda (k v)
            (set! alist (cons (cons k v) alist)))
          m)
        alist))


    (define (alist->map hash cmp alist)
      (let loop ((alist alist)
                 (m (make-map hash cmp)))
        (if (null? alist)
            m
          (loop (cdr alist) (insert m (caar alist) (cdar alist))))))


    (define (hash-bytevector b)
      (let loop ((i 0)
                 (hash 0))
        (if (>= i (bytevector-length b))
            hash
          (loop (+ 1 i) (+ (* hash #x100) (bytevector-u8-ref b i))))))


    (define (merge2 m1 m2)
      (let ((m1 m1))
        (map-for-each
          (lambda (k v)
            (set! m1 (insert m1 k v)))
          m2)
        m1))


    (define (merge m . m*)
      (let loop ((m* m*)
                 (m m))
        (match m*
          ('() m)
          ((head . tail) (loop tail (merge2 m head))))))))