/-
Copyright (c) 2017 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Author: Leonardo de Moura
-/
universes u v w

def bucket_array (α : Type u) (β : α → Type v) (n : nat) :=
array (list (Σ a, β a)) n

structure hash_map (α : Type u) [decidable_eq α] (β : α → Type v) :=
(hash_fn : α → nat)
(size nbuckets : nat)
(buckets : bucket_array α β nbuckets)

def mk_hash_map {α : Type u} [decidable_eq α] {β : α → Type v} (hash_fn : α → nat) (nbuckets := 8) : hash_map α β :=
{ hash_fn := hash_fn, size := 0, nbuckets := nbuckets,
  buckets := mk_array nbuckets [] }

namespace hash_map
variables {α : Type u} {β : α → Type v} {δ : Type w}

def reinsert_aux (hash_fn : α → nat) {n : nat} (data : bucket_array α β n) (a : α) (b : β a) : bucket_array α β n :=
let bidx := hash_fn a % n in
data^.write' bidx (⟨a, b⟩ :: data^.read' bidx)

def fold_buckets {n : nat} (data : bucket_array α β n) (d : δ) (f : δ → Π a, β a → δ) : δ :=
data^.foldl d (λ b d, list.foldl (λ r (p : Σ a, β a), f r p.1 p.2) d b)

variable [decidable_eq α]

def find_aux (a : α) : list (Σ a, β a) → option (β a)
| []          := none
| (⟨a',b⟩::t) :=
  if h : a' = a then some (eq.rec_on h b) else find_aux t

def contains_aux (a : α) (l : list (Σ a, β a)) : bool :=
(find_aux a l)^.is_some

def find (m : hash_map α β) (a : α) : option (β a) :=
match m with
| ⟨hash_fn, _, nbuckets, buckets⟩ :=
  find_aux a (buckets^.read' (hash_fn a % nbuckets))
end

def contains (m : hash_map α β) (a : α) : bool :=
(find m a)^.is_some

def fold [decidable_eq α] (m : hash_map α β) (d : δ) (f : δ → Π a, β a → δ) : δ :=
fold_buckets m^.buckets d f

def replace_aux (a : α) (b : β a) : list (Σ a, β a) → list (Σ a, β a)
| []            := []
| (⟨a', b'⟩::t) := if a' = a then ⟨a, b⟩::t else ⟨a', b'⟩ :: replace_aux t

def erase_aux (a : α) : list (Σ a, β a) → list (Σ a, β a)
| []            := []
| (⟨a', b'⟩::t) := if a' = a then t else ⟨a', b'⟩ :: erase_aux t

def insert (m : hash_map α β) (a : α) (b : β a) : hash_map α β :=
match m with
| ⟨hash_fn, size, nbuckets, buckets⟩ :=
  let bidx           := hash_fn a % nbuckets in
  let bkt            := buckets^.read' bidx in
  if contains_aux a bkt
  then ⟨hash_fn, size, nbuckets, buckets^.write' bidx (replace_aux a b bkt)⟩
  else let size'    := size + 1 in
       let buckets' := buckets^.write' bidx (⟨a, b⟩::bkt) in
       if size' <= nbuckets
       then ⟨hash_fn, size', nbuckets, buckets'⟩
       else let nbuckets' := nbuckets * 2 in
            ⟨hash_fn, size', nbuckets',
            fold_buckets buckets' (mk_array nbuckets' []) $
               λ r a b, reinsert_aux hash_fn r a b⟩
end

def erase (m : hash_map α β) (a : α) : hash_map α β :=
match m with
| ⟨hash_fn, size, nbuckets, buckets⟩ :=
  let bidx           := hash_fn a % nbuckets in
  let bkt            := buckets^.read' bidx in
  if contains_aux a bkt
  then ⟨hash_fn, size - 1, nbuckets, buckets^.write' bidx $ erase_aux a bkt⟩
  else m
end

section string
variables [has_to_string α] [∀ a, has_to_string (β a)]
open prod
private def key_data_to_string (a : α) (b : β a) (first : bool) : string :=
(if first then "" else ", ") ++ to_string a ++ " ← " ++ to_string b

private def to_string (m : hash_map α β) : string :=
"⟨" ++ (fst (fold m ("", tt) (λ p a b, (fst p ++ key_data_to_string a b (snd p), ff)))) ++ "⟩"

instance : has_to_string (hash_map α β) :=
⟨to_string⟩

end string

section format
open format prod
variables [has_to_format α] [∀ a, has_to_format (β a)]

private meta def format_key_data (a : α) (b : β a) (first : bool) : format :=
(if first then to_fmt "" else to_fmt "," ++ line) ++ to_fmt a ++ space ++ to_fmt "←" ++ space ++ to_fmt b

private meta def to_format (m : hash_map α β) : format :=
group $ to_fmt "⟨" ++ nest 1 (fst (fold m (to_fmt "", tt) (λ p a b, (fst p ++ format_key_data a b (snd p), ff)))) ++
        to_fmt "⟩"

meta instance : has_to_format (hash_map α β) :=
⟨to_format⟩
end format
end hash_map
