Add set operations to @immut/hash{map, set} and @internal/sparse_array Summary

immut/hashmap/HAMT.mbt

@@ -329,6 +329,133 @@ pub fn[K : Eq + Hash, V] union(self : T[K, V], other : T[K, V]) -> T[K, V] {
   }
 }
 
+///|
+/// Union two hashmaps with a function
+pub fn[K : Eq + Hash, V] union_with(
+  self : T[K, V],
+  other : T[K, V],
+  f : (K, V, V) -> V
+) -> T[K, V] {
+  match (self, other) {
+    (_, Empty) => self
+    (Empty, _) => other
+    (_, Leaf(k, v)) =>
+      match self.get(k) {
+        Some(v1) => self.add(k, f(k, v1, v))
+        None => self.add(k, v)
+      }
+    (Leaf(k, v), _) =>
+      match other.get(k) {
+        Some(v2) => other.add(k, f(k, v, v2))
+        None => other.add(k, v)
+      }
+    (Branch(sa1), Branch(sa2)) =>
+      Branch(sa1.union(sa2, fn(m1, m2) { m1.union_with(m2, f) }))
+    (_, _) =>
+      self
+      .iter()
+      .fold(init=other, fn(m, kv) {
+        match m.get(kv.0) {
+          Some(v2) => m.add(kv.0, f(kv.0, kv.1, v2))
+          None => m.add(kv.0, kv.1)
+        }
+      })
+  }
+}
+
+///|
+/// Intersect two hashmaps
+pub fn[K : Eq + Hash, V] intersection(
+  self : T[K, V],
+  other : T[K, V]
+) -> T[K, V] {
+  match (self, other) {
+    (_, Empty) => Empty
+    (Empty, _) => Empty
+    (Leaf(k, v), _) =>
+      match other.get(k) {
+        Some(_) => Leaf(k, v)
+        None => Empty
+      }
+    (_, Leaf(k, _)) =>
+      match self.get(k) {
+        Some(v) => Leaf(k, v)
+        None => Empty
+      }
+    (Branch(sa1), Branch(sa2)) =>
+      Branch(sa1.intersection(sa2, fn(m1, m2) { m1.intersection(m2) }))
+    (_, _) =>
+      self
+      .iter()
+      .fold(init=Empty, fn(m, kv) {
+        if other.get(kv.0) is Some(_) {
+          m.add(kv.0, kv.1)
+        } else {
+          m
+        }
+      })
+  }
+}
+
+///|
+/// Intersection two hashmaps with a function
+pub fn[K : Eq + Hash, V] intersection_with(
+  self : T[K, V],
+  other : T[K, V],
+  f : (K, V, V) -> V
+) -> T[K, V] {
+  match (self, other) {
+    (_, Empty) => Empty
+    (Empty, _) => Empty
+    (Leaf(k, v), _) =>
+      match other.get(k) {
+        Some(v2) => Leaf(k, f(k, v, v2))
+        None => Empty
+      }
+    (_, Leaf(k, v2)) =>
+      match self.get(k) {
+        Some(v1) => Leaf(k, f(k, v1, v2))
+        None => Empty
+      }
+    (Branch(sa1), Branch(sa2)) =>
+      Branch(sa1.intersection(sa2, fn(m1, m2) { m1.intersection_with(m2, f) }))
+    (_, _) =>
+      self
+      .iter()
+      .fold(init=Empty, fn(m, kv) {
+        match other.get(kv.0) {
+          Some(v2) => m.add(kv.0, f(kv.0, kv.1, v2))
+          None => m
+        }
+      })
+  }
+}
+
+///|
+/// Difference of two hashmaps: elements in `self` but not in `other`
+pub fn[K : Eq + Hash, V] difference(self : T[K, V], other : T[K, V]) -> T[K, V] {
+  match (self, other) {
+    (Empty, _) => Empty
+    (_, Empty) => self
+    (Leaf(k, v), _) =>
+      match other.get(k) {
+        Some(_) => Empty
+        None => Leaf(k, v)
+      }
+    (Branch(sa1), Branch(sa2)) => Branch(sa1.difference(sa2))
+    (_, _) =>
+      self
+      .iter()
+      .fold(init=Empty, fn(m, kv) {
+        if other.get(kv.0) is None {
+          m.add(kv.0, kv.1)
+        } else {
+          m
+        }
+      })
+  }
+}
+
 ///|
 /// Iterate through the elements in a hash map
 pub fn[K, V] each(self : T[K, V], f : (K, V) -> Unit) -> Unit {

immut/hashmap/HAMT_test.mbt

@@ -580,6 +580,99 @@ test "HAMT::union leaf to non-overlapping map" {
   assert_eq(u.get(2), Some(2))
 }
 
+///|
+test "HAMT::intersection with empty" {
+  let m1 = @hashmap.of([(1, 1)])
+  let m2 = @hashmap.new()
+  assert_eq!(m1.intersection(m2).size(), 0)
+  assert_eq!(m2.intersection(m1).size(), 0)
+}
+
+///|
+test "HAMT::intersection with leaf" {
+  let m1 = @hashmap.of([(1, 1), (2, 2)])
+  let m2 = @hashmap.singleton(2, 2)
+  let m3 = @hashmap.singleton(3, 3)
+  assert_eq!(m1.intersection(m2).get(2), Some(2))
+  assert_eq!(m1.intersection(m3).get(3), None)
+  assert_eq!(m2.intersection(m1).get(2), Some(2))
+}
+
+///|
+test "HAMT::intersection with branch" {
+  let m1 = @hashmap.of([(1, 1), (2, 2), (3, 3)])
+  let m2 = @hashmap.of([(2, 2), (3, 30), (4, 4)])
+  let inter = m1.intersection(m2)
+  assert_eq!(inter.get(1), None)
+  assert_eq!(inter.get(2), Some(2))
+  assert_eq!(inter.get(3), Some(3))
+  assert_eq!(inter.get(4), None)
+}
+
+///|
+test "HAMT::intersection_with with empty" {
+  let m1 = @hashmap.of([(1, 1)])
+  let m2 = @hashmap.new()
+  assert_eq!(m1.intersection_with(m2, fn(_k, v1, v2) { v1 + v2 }).size(), 0)
+  assert_eq!(m2.intersection_with(m1, fn(_k, v1, v2) { v1 + v2 }).size(), 0)
+}
+
+///|
+test "HAMT::intersection_with with leaf" {
+  let m1 = @hashmap.of([(1, 1), (2, 2)])
+  let m2 = @hashmap.singleton(2, 20)
+  let m3 = @hashmap.singleton(3, 30)
+  assert_eq!(
+    m1.intersection_with(m2, fn(_k, v1, v2) { v1 + v2 }).get(2),
+    Some(22),
+  )
+  assert_eq!(m1.intersection_with(m3, fn(_k, v1, v2) { v1 + v2 }).get(3), None)
+  assert_eq!(
+    m2.intersection_with(m1, fn(_k, v1, v2) { v1 + v2 }).get(2),
+    Some(22),
+  )
+}
+
+///|
+test "HAMT::intersection_with with branch" {
+  let m1 = @hashmap.of([(1, 1), (2, 2), (3, 3)])
+  let m2 = @hashmap.of([(2, 20), (3, 30), (4, 4)])
+  let inter = m1.intersection_with(m2, fn(_k, v1, v2) { v1 * v2 })
+  assert_eq!(inter.get(1), None)
+  assert_eq!(inter.get(2), Some(40))
+  assert_eq!(inter.get(3), Some(90))
+  assert_eq!(inter.get(4), None)
+}
+
+///|
+test "HAMT::difference with empty" {
+  let m1 = @hashmap.of([(1, 1)])
+  let m2 = @hashmap.new()
+  assert_eq!(m1.difference(m2), m1)
+  assert_eq!(m2.difference(m1).size(), 0)
+}
+
+///|
+test "HAMT::difference with leaf" {
+  let m1 = @hashmap.of([(1, 1), (2, 2)])
+  let m2 = @hashmap.singleton(2, 2)
+  let m3 = @hashmap.singleton(3, 3)
+  assert_eq!(m1.difference(m2).get(2), None)
+  assert_eq!(m1.difference(m3).get(1), Some(1))
+  assert_eq!(m2.difference(m1).size(), 0)
+}
+
+///|
+test "HAMT::difference with branch" {
+  let m1 = @hashmap.of([(1, 1), (2, 2), (3, 3)])
+  let m2 = @hashmap.of([(2, 2), (3, 30), (4, 4)])
+  let diff = m1.difference(m2)
+  assert_eq!(diff.get(1), Some(1))
+  assert_eq!(diff.get(2), None)
+  assert_eq!(diff.get(3), None)
+  assert_eq!(diff.get(4), None)
+}
+
 ///|
 test "HAMT::each" {
   let empty = @hashmap.new()

immut/hashmap/hashmap.mbti

@@ -9,6 +9,8 @@ fn[K : Eq + Hash, V] add(T[K, V], K, V) -> T[K, V]
 
 fn[K : Eq + Hash, V] contains(T[K, V], K) -> Bool
 
+fn[K : Eq + Hash, V] difference(T[K, V], T[K, V]) -> T[K, V]
+
 fn[K, V] each(T[K, V], (K, V) -> Unit) -> Unit
 
 fn[K, V] elems(T[K, V]) -> Iter[V]
@@ -28,6 +30,10 @@ fn[K : Eq + Hash, V] from_iter(Iter[(K, V)]) -> T[K, V]
 
 fn[K : Eq + Hash, V] get(T[K, V], K) -> V?
 
+fn[K : Eq + Hash, V] intersection(T[K, V], T[K, V]) -> T[K, V]
+
+fn[K : Eq + Hash, V] intersection_with(T[K, V], T[K, V], (K, V, V) -> V) -> T[K, V]
+
 fn[K, V] iter(T[K, V]) -> Iter[(K, V)]
 
 fn[K, V] iter2(T[K, V]) -> Iter2[K, V]
@@ -55,11 +61,14 @@ fn[K, V] to_array(T[K, V]) -> Array[(K, V)]
 
 fn[K : Eq + Hash, V] union(T[K, V], T[K, V]) -> T[K, V]
 
+fn[K : Eq + Hash, V] union_with(T[K, V], T[K, V], (K, V, V) -> V) -> T[K, V]
+
 // Types and methods
 type T[K, V]
 impl T {
   add[K : Eq + Hash, V](Self[K, V], K, V) -> Self[K, V]
   contains[K : Eq + Hash, V](Self[K, V], K) -> Bool
+  difference[K : Eq + Hash, V](Self[K, V], Self[K, V]) -> Self[K, V]
   each[K, V](Self[K, V], (K, V) -> Unit) -> Unit
   elems[K, V](Self[K, V]) -> Iter[V]
   filter[K : Eq + Hash, V](Self[K, V], (V) -> Bool) -> Self[K, V]
@@ -68,6 +77,8 @@ impl T {
   fold[K, V, A](Self[K, V], init~ : A, (A, V) -> A) -> A
   fold_with_key[K, V, A](Self[K, V], init~ : A, (A, K, V) -> A) -> A
   get[K : Eq + Hash, V](Self[K, V], K) -> V?
+  intersection[K : Eq + Hash, V](Self[K, V], Self[K, V]) -> Self[K, V]
+  intersection_with[K : Eq + Hash, V](Self[K, V], Self[K, V], (K, V, V) -> V) -> Self[K, V]
   iter[K, V](Self[K, V]) -> Iter[(K, V)]
   iter2[K, V](Self[K, V]) -> Iter2[K, V]
   keys[K, V](Self[K, V]) -> Iter[K]
@@ -79,6 +90,7 @@ impl T {
   size[K, V](Self[K, V]) -> Int
   to_array[K, V](Self[K, V]) -> Array[(K, V)]
   union[K : Eq + Hash, V](Self[K, V], Self[K, V]) -> Self[K, V]
+  union_with[K : Eq + Hash, V](Self[K, V], Self[K, V], (K, V, V) -> V) -> Self[K, V]
 }
 impl[K : Eq + Hash, V : Eq] Eq for T[K, V]
 impl[K : Hash, V : Hash] Hash for T[K, V]

immut/hashset/HAMT.mbt

@@ -192,6 +192,56 @@ pub fn[K : Eq + Hash] union(self : T[K], other : T[K]) -> T[K] {
   }
 }
 
+///|
+/// Intersect two hashsets
+pub fn[K : Eq + Hash] intersection(self : T[K], other : T[K]) -> T[K] {
+  match (self, other) {
+    (_, Empty) => Empty
+    (Empty, _) => Empty
+    (Leaf(k), _) => if other.contains(k) { Leaf(k) } else { Empty }
+    (_, Leaf(k)) => if self.contains(k) { Leaf(k) } else { Empty }
+    (Branch(sa1), Branch(sa2)) => {
+      let res = sa1.intersection(sa2, fn(m1, m2) { m1.intersection(m2) })
+      if res.size() == 0 {
+        Empty
+      } else {
+        Branch(res)
+      }
+    }
+    (_, _) =>
+      self
+      .iter()
+      .fold(init=Empty, fn(m, k) {
+        if other.contains(k) {
+          m.add(k)
+        } else {
+          m
+        }
+      })
+  }
+}
+
+///|
+/// Difference of two hashsets: elements in `self` but not in `other`
+pub fn[K : Eq + Hash] difference(self : T[K], other : T[K]) -> T[K] {
+  match (self, other) {
+    (Empty, _) => Empty
+    (_, Empty) => self
+    (Leaf(k), _) => if other.contains(k) { Empty } else { Leaf(k) }
+    (Branch(sa1), Branch(sa2)) => Branch(sa1.difference(sa2))
+    (_, _) =>
+      self
+      .iter()
+      .fold(init=Empty, fn(m, k) {
+        if other.contains(k) {
+          m
+        } else {
+          m.add(k)
+        }
+      })
+  }
+}
+
 ///|
 /// Returns true if the hash set is empty.
 pub fn[A] is_empty(self : T[A]) -> Bool {

immut/hashset/HAMT_test.mbt

@@ -217,3 +217,43 @@ test "union 2 hashsets" {
   let set4 = @hashset.of([1, 2, 3, 4, 5])
   inspect(set3 == set4, content="true")
 }
+
+///|
+test "@hashset.intersection with overlapping sets" {
+  let set1 = @hashset.of([1, 2, 3, 4])
+  let set2 = @hashset.of([3, 4, 5, 6])
+  let result = set1.intersection(set2)
+  inspect!(result, content="@immut/hashset.of([3, 4])")
+}
+
+///|
+test "@hashset.intersection with disjoint sets" {
+  let set1 = @hashset.of([1, 2])
+  let set2 = @hashset.of([3, 4])
+  let result = set1.intersection(set2)
+  inspect!(result.is_empty(), content="true")
+}
+
+///|
+test "@hashset.intersection with one empty set" {
+  let set1 = @hashset.of([1, 2, 3])
+  let set2 = @hashset.new()
+  let result = set1.intersection(set2)
+  inspect!(result.is_empty(), content="true")
+}
+
+///|
+test "@hashset.intersection with identical sets" {
+  let set1 = @hashset.of([1, 2, 3])
+  let set2 = @hashset.of([1, 2, 3])
+  let result = set1.intersection(set2)
+  inspect!(result == set1, content="true")
+}
+
+///|
+test "@hashset.intersection with subset" {
+  let set1 = @hashset.of([1, 2, 3, 4])
+  let set2 = @hashset.of([2, 3])
+  let result = set1.intersection(set2)
+  inspect!(result == @hashset.of([2, 3]), content="true")
+}