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