add removal API for immutable hash map (#110)

Co-authored-by: Yorkin <[email protected]>

immutable_hashmap/HAMT.mbt

@@ -148,6 +148,68 @@ test "HAMT" {
   }
 }
 
+/// Remove an element from a map
+pub fn remove[K : Eq + Hash, V](self : Map[K, V], key : K) -> Map[K, V] {
+  self.remove_with_hash(key, 0, key.hash())
+}
+
+fn remove_with_hash[K : Eq, V](
+  self : Map[K, V],
+  key : K,
+  depth : Int,
+  hash : Int
+) -> Map[K, V] {
+  match self {
+    Empty => self
+    Leaf(old_key, _) => if key == old_key { Empty } else { self }
+    Collision(bucket) =>
+      match bucket.remove(key) {
+        None => Empty
+        Some(Just_One(k, v)) => Leaf(k, v)
+        Some(new_bucket) => Collision(new_bucket)
+      }
+    Branch(children) => {
+      let idx = hash.land((1).lsl(segment_length) - 1)
+      match children[idx] {
+        None => self
+        Some(child) => {
+          let new_child = child.remove_with_hash(
+            key,
+            depth + segment_length,
+            hash.lsr(segment_length),
+          )
+          if children.size() == 1 {
+            match new_child {
+              Empty => Empty
+              _ => Branch(children.replace(idx, new_child))
+            }
+          } else {
+            Branch(children.replace(idx, new_child))
+          }
+        }
+      }
+    }
+  }
+}
+
+test "HAMT::remove" {
+  let map = loop 0, Map::make() {
+    100, map => map
+    i, map => continue i + 1, map.add(i, i)
+  }
+  for i = 0; i < 100; i = i + 1 {
+    @assertion.assert_eq((i, map.find(i)), (i, Some(i)))?
+  }
+  let map = loop 0, map {
+    100, map => map.remove(100) // test for removing non-existing element
+    i, map => continue i + 2, map.remove(i)
+  }
+  for i = 0; i < 100; i = i + 2 {
+    @assertion.assert_eq(map.find(i), None)?
+    @assertion.assert_eq(map.find(i + 1), Some(i + 1))?
+  }
+}
+
 /// Calculate the size of a map.
 ///
 /// WARNING: this operation is `O(N)` in map size

immutable_hashmap/README.md

@@ -12,6 +12,3 @@ test {
   @assert.assert_eq(h2.find(3), Some(3))
 }
 ```
-
-## TODO
-More API, in particular removal.

immutable_hashmap/bucket.mbt

@@ -51,6 +51,22 @@ fn add[K : Eq, V](self : Bucket[K, V], key : K, value : V) -> Bucket[K, V] {
   }
 }
 
+/// Remove a key from a bucket
+fn remove[K : Eq, V](self : Bucket[K, V], key : K) -> Option[Bucket[K, V]] {
+  match self {
+    Just_One(old_key, _) => if key == old_key { None } else { Some(self) }
+    More(old_key, old_value, rest) =>
+      if key == old_key {
+        Some(rest)
+      } else {
+        match rest.remove(key) {
+          None => Some(Just_One(old_key, old_value))
+          Some(rest) => Some(More(old_key, old_value, rest))
+        }
+      }
+  }
+}
+
 /// Get the size of a bucket
 fn size[K, V](self : Bucket[K, V]) -> Int {
   match self {
@@ -72,6 +88,20 @@ test "Bucket" {
   @assertion.assert_eq(b.find(0), Some(2))?
   @assertion.assert_eq(b.find(1), Some(1))?
   @assertion.assert_eq(b.size(), 2)?
+  match b.remove(0) {
+    None => @assertion.assert_true(false)?
+    Some(b1) => {
+      @assertion.assert_eq(b1.find(0), None)?
+      @assertion.assert_eq(b1.find(1), Some(1))?
+    }
+  }
+  match b.remove(1) {
+    None => @assertion.assert_true(false)?
+    Some(b1) => {
+      @assertion.assert_eq(b1.find(0), Some(2))?
+      @assertion.assert_eq(b1.find(1), None)?
+    }
+  }
 }
 
 /// Iterate through elements of a bucket

immutable_hashmap/sparse_array.mbt

@@ -102,6 +102,13 @@ test "SparseArray" {
   @assertion.assert_eq(arr[2], Some(2))?
 }
 
+/// `size(self: SparseArray[X]) -> Int`
+///
+/// Return the size of a sparse array
+fn size[X](self : SparseArray[X]) -> Int {
+  self.data.length()
+}
+
 /// `iter(self: SparseArray[X], f: (X) -> Unit) -> Unit`
 ///
 /// Iterate through elements in a sparse array