(definterface hash (Fn [(Ref a)] Int))

(defmodule String
  (defn rehash [k l]
    (let-do [a 31415
             b 27183
             vh 0]
      (for [x 0 (length k)]
        (do
          (set! vh (+ (* a (* vh l)) (Char.to-int (char-at k x))))
          (set! a (* a b))
          (set! x (Int.inc x))))
      (Int.abs vh)))

  (defn hash [k]
    (rehash k 1))
)

(defmodule Int
  (defn hash [k] (the Int @k))
)

(defmodule Long
  (defn hash [k] (to-int (the Long @k)))
)

(defmodule Bool
  (defn hash [k] (if (the Bool @k) 1 0))
)

(defmodule Char
  (defn hash [k] (to-int (the Char @k)))
)

(defmodule Float
  (defn hash [k] (to-bytes @k))
)

(defmodule Double
  (defn hash [k] (Long.to-int (to-bytes @k)))
)

(deftype (Bucket a b) [entries (Array (Pair a b))])

(defmodule Bucket
  (defn empty []
    (Bucket.init []))

  (defn grow [b e]
    (set-entries @b (Array.push-back @(entries b) e)))

  (defn get [b k]
    (let-do [e &(zero)
             l (Array.length (Bucket.entries b))
             es (entries b)]
      (for [i 0 l]
        (when (= (Pair.a (Array.nth es i)) k)
          (do
            (set! e (Pair.b (Array.nth es i)))
            (break))))
       @e))

  (defn contains? [b k]
    (let-do [e false
             l (Array.length (Bucket.entries b))
             es (entries b)]
      (for [i 0 l]
        (when (= (Pair.a (Array.nth es i)) k)
          (do
            (set! e true)
            (break))))
       e))

  (defn remove [entries k]
    (let-do [nentries (the (Array (Pair a b)) [])]
      (for [i 0 (Array.length entries)]
        (let [e (Array.nth entries i)]
          (unless (= (Pair.a e) k)
             (set! nentries (Array.push-back nentries @e)))))
      nentries))

  (defn shrink [b k]
    (if (contains? b k)
      (set-entries @b (remove (entries b) k))
      @b))
)

(deftype (Map a b) [n-buckets Int buckets (Array (Bucket a b))])

(defmodule Map
  (hidden dflt-len)
  (def dflt-len 256)

  (doc create "Create an empty map.")
  (defn create []
    (init dflt-len (Array.repeat dflt-len Bucket.empty)))

  (doc create-with-len "Create an empty map with a given number of buckets. High numbers reduce the possibility of hash collisions while increasing the memory footprint.")
  (defn create-with-len [len]
    (init len (Array.repeat len Bucket.empty)))

  (doc put "Put a a value v into map m, using the key k.")
  (defn put [m k v]
    (let [idx (Int.mod (hash k) @(n-buckets m))
          b (buckets m)]
      (set-buckets @m (Array.aset @b
                                  idx
                                  (Bucket.grow (Array.nth b idx) (Pair.init @k @v))))))

  (doc get "Get the value for the key k from map m. If it isn’t found, a zero element for the value type is returned.")
  (defn get [m k]
    (let [idx (Int.mod (hash k) @(n-buckets m))]
      (Bucket.get (Array.nth (buckets m) idx) k)))

  (doc length "Get the length of the map m.")
  (defn length [m]
    (let-do [c 0]
      (for [i 0 @(n-buckets m)]
        (set! c (+ c (Array.length (Bucket.entries (Array.nth (buckets m) i))))))
      c))

  (doc empty "Check whether the map m is empty.")
  (defn empty? [m]
    (= (length m) 0))

  (doc contains? "Check whether the map m contains the key k.")
  (defn contains? [m k]
    (let [idx (Int.mod (hash k) @(n-buckets m))]
      (Bucket.contains? (Array.nth (buckets m) idx) k)))

  (doc remove "Remove the value under the key k from the map m.")
  (defn remove [m k]
    (let [idx (Int.mod (hash k) @(n-buckets &m))
          b (buckets &m)]
      (set-buckets m (Array.aset @b
                                 idx
                                 (Bucket.shrink (Array.nth b idx) k)))))

  (doc for-each "Execute the binary function f for all keys and values in the map m.")
  (defn for-each [m f]
    (for [i 0 @(n-buckets m)]
      (let [bucket (Array.nth (buckets m) i)
            len (Array.length (Bucket.entries bucket))
            entries (Bucket.entries bucket)]
        (for [j 0 len]
          (let [e (Array.nth entries j)]
            (f (Pair.a e) (Pair.b e)))))))

  (doc from-array "Create a map from the array a containing key-value pairs.")
  (defn from-array [a]
    (let-do [m (create)]
      (for [i 0 (Array.length a)]
        (let [e (Array.nth a i)
              k (Pair.a e)
              v (Pair.b e)]
          (set! m (put &m k v))))
      m))

  (defn str [m]
    (let-do [res @"{"]
      (for [i 0 @(n-buckets m)]
        (let [bucket (Array.nth (buckets m) i)
              len (Array.length (Bucket.entries bucket))
              entries (Bucket.entries bucket)]
          (for [j 0 len]
            (let [e (Array.nth entries j)]
              (set! res (String.join @"" &[res @" " (prn @(Pair.a e)) @" " (prn @(Pair.b e))]))))))
      (String.append &res " }")))
)

(deftype (SetBucket a) [entries (Array a)])

(defmodule SetBucket
  (defn empty []
    (SetBucket.init []))

  (defn grow [b e]
    (set-entries @b (Array.push-back @(entries b) e)))

  (defn contains? [b k]
    (let-do [e false
             es (entries b)
             l (Array.length es)]
      (for [i 0 l]
        (when (= (Array.nth es i) k)
          (do
            (set! e true)
            (break))))
       e))

  (defn remove [entries k]
    (let-do [nentries []]
      (for [i 0 (Array.length entries)]
        (let [e (Array.nth entries i)]
          (unless (= e k)
             (set! nentries (Array.push-back nentries @e)))))
      nentries))

  (defn shrink [b k]
    (if (contains? b k)
      (set-entries @b (remove (entries b) k))
      @b))
)

(deftype (Set a) [n-buckets Int buckets (Array (SetBucket a))])

(defmodule Set
  (hidden dflt-len)
  (def dflt-len 256)

  (doc create "Create an empty set.")
  (defn create []
    (init dflt-len (Array.repeat dflt-len SetBucket.empty)))

  (doc create-with-len "Create an empty set with a given number of buckets. Higher numbers decrease the probability of hash collisions while increasing the memory footprint.")
  (defn create-with-len [len]
    (init len (Array.repeat len SetBucket.empty)))

  (doc put "Put a a key k into the set s.")
  (defn put [s k]
    (let [idx (Int.mod (hash k) @(n-buckets s))
          b (buckets s)]
      (set-buckets @s (Array.aset @b
                                  idx
                                  (SetBucket.grow (Array.nth b idx) @k)))))

  (doc length "Get the length of set s.")
  (defn length [s]
    (let-do [c 0]
      (for [i 0 @(n-buckets s)]
        (set! c (+ c (Array.length (SetBucket.entries (Array.nth (buckets s) i))))))
      c))

  (doc empty? "Check whether the set s is empty.")
  (defn empty? [s]
    (= (length s) 0))

  (doc contains? "Check whether the set s contains the key k.")
  (defn contains? [s k]
    (let [idx (Int.mod (hash k) @(n-buckets s))]
      (SetBucket.contains? (Array.nth (buckets s) idx) k)))

  (doc remove "Remove the key k from the set s.")
  (defn remove [s k]
    (let [idx (Int.mod (hash k) @(n-buckets &s))
          b (buckets &s)]
      (set-buckets s (Array.aset @b
                                 idx
                                 (SetBucket.shrink (Array.nth b idx) k)))))


  (doc for-each "Execute the unary function f for each element in the set s.")
  (defn for-each [s f]
    (for [i 0 @(n-buckets s)]
      (let [bucket (Array.nth (buckets s) i)
            len (Array.length (SetBucket.entries bucket))
            entries (SetBucket.entries bucket)]
        (for [j 0 len]
          (let [e (Array.nth entries j)]
            (f e))))))

  (doc from-array "Create a set from the values in array a.")
  (defn from-array [a]
    (let-do [s (create)]
      (for [i 0 (Array.length a)]
        (let [e (Array.nth a i)]
          (set! s (put &s e))))
      s))

  (defn str [s]
    (let-do [res @"{"]
      (for [i 0 @(n-buckets s)]
        (let [bucket (Array.nth (buckets s) i)
              len (Array.length (SetBucket.entries bucket))
              entries (SetBucket.entries bucket)]
          (for [j 0 len]
            (let [e (Array.nth entries j)]
              (set! res (String.join @"" &[res @" " (prn e)]))))))
      (String.append &res " }")))
)