plum

module std
import std/List
import std/Option
import std/Array
import std/Bool
import std/Number
import std/Str

# Any type usable as a `Map` key needs a `hash`, so keys landing in different
# buckets (`hash(k) % BUCKET_COUNT`) can be told apart in O(1) without a full
# linear scan of the whole map — see `Str.hash`/`Int.hash` in `str.plum`/
# `int.plum` for the two implementations `libs/std` provides today. Declared
# here (rather than alongside `ToStr`/`Comparable` in `str.plum`) since `Map`
# is its only real consumer; like `Comparable`/`Readable`/`Writable`
# elsewhere in `libs/std`, `Str`/`Int` don't formally list it in their own
# `implements` clause (that would need them to import `map.plum`, an
# unwanted dependency direction for such foundational types) — this is
# checked the same "unenforced, trust the method exists" way those already
# are (see the README's note on `checkTraitConformance`).
trait Hashable =
  hash() -> Int

# A Pair is a grouping of a key with a value.
#
# `toStr` (unlike `Option`'s — see `libs/std/option.plum`'s note on
# `toOptionStr`) IS safe as an ordinary method here: `Pair[K, V]` is only
# ever used as `List[Pair[K, V]]`'s element type inside `Map[K, V]`, and
# `List[T: ToStr]` already requires its own `T` to support `toStr` —
# so a `Pair` used as a `List` element always has `K`/`V: ToStr` too,
# and `List`'s own eagerly-compiled `join`/`toStr` (which call `.toStr()` on
# every element) need this to exist for every `Map` specialization anyway.
enum Pair[K, V](ToStr) =
  | Pair(key: K, val: V)

  fun toStr(self) -> Str =
    return "{self.key.toStr()}: {self.val.toStr()}"

# Number of buckets every `Map` allocates at construction — fixed for the
# map's whole lifetime, no rehashing/resizing as it grows. A real hash table
# would grow this as `size` outpaces it (keeping average bucket length, and
# so average `get`/`set` cost, roughly constant); that's future work; for now
# a big-enough fixed count keeps small-to-medium maps fast in practice.
BUCKET_COUNT = 16

# A Map is a hash table: `BUCKET_COUNT` buckets, each an association-list
# chain (`List[Pair[K, V]]`) of every key hashing into it — O(1) average
# `get`/`set`/`remove` (versus the single, whole-map linear scan a plain
# association list costs), degrading to O(n) only if many keys collide into
# the same bucket.
enum Map[K: Hashable, V] =
  | Map(buckets: Array[List[Pair[K, V]]], size: Int)

  # `Map[Str, Int]()` — explicit generics required (nothing in a no-arg call
  # pins down `K`/`V`), builds `BUCKET_COUNT` empty bucket lists up front via
  # `Array`'s own `push`-to-grow (see `libs/std/array.plum`'s header comment
  # on why `Array` itself has no direct "make me length N" constructor).
  fun init() -> Map[K, V] =
    b := Array[List[Pair[K, V]]]()
    i := 0
    while i < BUCKET_COUNT
      b = b.push(List[Pair[K, V]]())
      i = i + 1
    return Map(buckets: b, size: 0)

  # `k`'s bucket index — `hash`'s sign is unconstrained (a plain `%` on a
  # negative `Int` stays negative in Plum, matching wasm's signed
  # `i64.rem_s`), so this folds a negative result back into `[0,
  # BUCKET_COUNT)` rather than indexing `buckets` out of range.
  fun bucketIndex(self, k: K) -> Int =
    m := k.hash() % BUCKET_COUNT
    return m < 0 ? m + BUCKET_COUNT : m

  # Gets the value for `k`, or `None` if absent.
  fun get(self, k: K) -> Option[V] =
    current := self.buckets.get(self.bucketIndex(k)).head
    while current != None
      match current
        Some(Node(Pair(key, val), _, next)) =>
          if key == k
            return Some(val)
          current = next
        None =>
          break
    return None

  fun has(self, k: K) -> Bool =
    return self.get(k) != None

  # Sets `k`'s value to `v` — replacing an existing entry for `k` in place
  # (mutating the stored `Pair`'s `val` field) rather than appending a
  # duplicate, so a repeated `set` on the same key doesn't grow its bucket
  # unboundedly or leave a stale value reachable by `get`.
  fun set(self, k: K, v: V) -> Unit =
    bucket := self.buckets.get(self.bucketIndex(k))
    current := bucket.head
    while current != None
      match current
        Some(Node(pair, _, next)) =>
          if pair.key == k
            pair.val = v
            return
          current = next
        None =>
          break
    bucket.add(Pair(key: k, val: v))
    self.size = self.size + 1

  # Sets `k`'s value to `v` only if `k` isn't already present.
  fun putIfAbsent(self, k: K, v: V) -> Unit =
    if !self.has(k)
      self.set(k, v)

  # Removes `k`'s entry, if present.
  fun remove(self, k: K) -> Unit =
    bucket := self.buckets.get(self.bucketIndex(k))
    current := bucket.head
    while current != None
      match current
        Some(node) =>
          if node.value.key == k
            bucket.unlink(node)
            self.size = self.size - 1
            return
          current = node.next
        None =>
          break

  fun length(self) -> Int =
    return self.size

  fun isEmpty(self) -> Bool =
    return self.size == 0

  fun clear(self) -> Unit =
    i := 0
    while i < self.buckets.length()
      self.buckets.get(i).clear()
      i = i + 1
    self.size = 0

  # Calls `cb` with each key and value in the map. Order is bucket order,
  # then insertion order within a bucket — NOT overall insertion order (that
  # would need a separate side list to track, which nothing here needs today).
  fun each(self, cb: fn(K, V)) -> Unit =
    i := 0
    while i < self.buckets.length()
      current := self.buckets.get(i).head
      while current != None
        match current
          Some(Node(Pair(key, val), _, next)) =>
            cb(key, val)
            current = next
          None =>
            break
      i = i + 1

  fun keys(self) -> List[K] =
    result := List(head: None, tail: None, size: 0)
    self.each(|k, v| result.add(k))
    return result

  fun values(self) -> List[V] =
    result := List(head: None, tail: None, size: 0)
    self.each(|k, v| result.add(v))
    return result

  fun map(self, cb: fn(K, V) -> Pair[X, Y]) -> Map[X, Y] =
    b := Array[List[Pair[X, Y]]]()
    i := 0
    while i < BUCKET_COUNT
      b = b.push(List[Pair[X, Y]]())
      i = i + 1
    result := Map(buckets: b, size: 0)
    j := 0
    while j < self.buckets.length()
      current := self.buckets.get(j).head
      while current != None
        match current
          Some(Node(Pair(key, val), _, next)) =>
            pair := cb(key, val)
            result.set(pair.key, pair.val)
            current = next
          None =>
            break
      j = j + 1
    return result

test "get/set/has round-trip and set replaces an existing key in place"
  m := Map[Str, Int]()
  assert m.has("a") == False
  m.set("a", 1)
  assert m.has("a") == True
  assert m.get("a").unwrap() == 1
  assert m.length() == 1
  m.set("a", 2)
  assert m.get("a").unwrap() == 2
  assert m.length() == 1

test "putIfAbsent only sets a key that isn't already present"
  m := Map[Str, Int]()
  m.putIfAbsent("a", 1)
  m.putIfAbsent("a", 2)
  assert m.get("a").unwrap() == 1

test "remove deletes an entry and clear empties the map"
  m := Map[Str, Int]()
  m.set("a", 1)
  m.set("b", 2)
  m.remove("a")
  assert m.has("a") == False
  assert m.length() == 1
  m.clear()
  assert m.isEmpty() == True

test "keys/values list every entry's key/value"
  m := Map[Str, Int]()
  m.set("a", 1)
  m.set("b", 2)
  assert m.length() == 2
  assert m.keys().contains("a")
  assert m.keys().contains("b")
  assert m.values().contains(1)
  assert m.values().contains(2)

test "each calls the callback with every key and value"
  m := Map[Str, Int]()
  m.set("a", 1)
  m.set("b", 2)
  # A closure captures its outer variables BY VALUE at creation time, so
  # reassigning a captured `Int` inside the callback wouldn't be visible out
  # here — mutate a captured `List` (a class, reference semantics) instead.
  seen := List[Int]()
  m.each(|k, v| seen.add(v))
  assert seen.length() == 2
  assert seen.contains(1)
  assert seen.contains(2)

test "map transforms every key/value pair into a new map, changing the value type"
  m := Map[Str, Int]()
  m.set("a", 1)
  m.set("b", 2)
  stringified := m.map(|k, v| Pair(key: k, val: v.toStr()))
  assert stringified.get("a").unwrap() == "1"
  assert stringified.get("b").unwrap() == "2"
  assert stringified.length() == 2

test "many keys landing in the same bucket still resolve correctly (collision chaining)"
  m := Map[Int, Int]()
  i := 0
  while i < 40
    m.set(i, i * i)
    i = i + 1
  assert m.length() == 40
  j := 0
  while j < 40
    assert m.get(j).unwrap() == j * j
    j = j + 1
  m.remove(5)
  assert m.has(5) == False
  assert m.get(4).unwrap() == 16
  assert m.get(21).unwrap() == 441
  assert m.length() == 39