feat: add set operations to @immut/hash{map, set} and @internal/sparse_array (#2145)

* feat: add four new functions to HAMT and their tests

* fix(union_with): fix union_with for HAMT,now it can handle branch

* feat(sparse_array): add intersection and difference methods

* fix(HAMT): fix union_with, intersection, intersection_with and difference methods, now they can handle branchs

* feat(hashset): add intersection and difference methods to hashset

* commit other files

* feat: add four new functions to HAMT and their tests

* fix(union_with): fix union_with for HAMT,now it can handle branch

* feat(sparse_array): add intersection and difference methods

* fix(HAMT): fix union_with, intersection, intersection_with and difference methods, now they can handle branchs

* feat(hashset): add intersection and difference methods to hashset

* style: change the position of some function declarations

* fix: fix formatting of the code

* feat: Update the function declarations of hash tables and sparse arrays

* refactor:update mbti

* refactor: update hashmap and hashset function signatures to include type prefix

---------

Co-authored-by: 东灯 <[email protected]>

Author: Asterless

immut/hashmap/HAMT.mbt

@@ -311,7 +311,7 @@ pub fn[K, V] size(self : T[K, V]) -> Int {
 
 ///|
 /// Union two hashmaps
-pub fn[K : Eq + Hash, V] union(self : T[K, V], other : T[K, V]) -> T[K, V] {
+pub fn[K : Eq + Hash, V] T::union(self : T[K, V], other : T[K, V]) -> T[K, V] {
   match (self, other) {
     (_, Empty) => self
     (Empty, _) => other
@@ -329,6 +329,136 @@ 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] T::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] T::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] T::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] T::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

@@ -54,14 +54,13 @@ fn[K, V] size(T[K, V]) -> Int
 
 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, V] values(T[K, V]) -> Iter[V]
 
 // Types and methods
 type T[K, V]
 fn[K : Eq + Hash, V] T::add(Self[K, V], K, V) -> Self[K, V]
 fn[K : Eq + Hash, V] T::contains(Self[K, V], K) -> Bool
+fn[K : Eq + Hash, V] T::difference(Self[K, V], Self[K, V]) -> Self[K, V]
 fn[K, V] T::each(Self[K, V], (K, V) -> Unit) -> Unit
 #deprecated
 fn[K, V] T::elems(Self[K, V]) -> Iter[V]
@@ -71,6 +70,8 @@ fn[K : Eq + Hash, V] T::find(Self[K, V], K) -> V?
 fn[K, V, A] T::fold(Self[K, V], init~ : A, (A, V) -> A) -> A
 fn[K, V, A] T::fold_with_key(Self[K, V], init~ : A, (A, K, V) -> A) -> A
 fn[K : Eq + Hash, V] T::get(Self[K, V], K) -> V?
+fn[K : Eq + Hash, V] T::intersection(Self[K, V], Self[K, V]) -> Self[K, V]
+fn[K : Eq + Hash, V] T::intersection_with(Self[K, V], Self[K, V], (K, V, V) -> V) -> Self[K, V]
 fn[K, V] T::iter(Self[K, V]) -> Iter[(K, V)]
 fn[K, V] T::iter2(Self[K, V]) -> Iter2[K, V]
 fn[K, V] T::keys(Self[K, V]) -> Iter[K]
@@ -82,6 +83,7 @@ fn[K : Eq + Hash, V] T::remove(Self[K, V], K) -> Self[K, V]
 fn[K, V] T::size(Self[K, V]) -> Int
 fn[K, V] T::to_array(Self[K, V]) -> Array[(K, V)]
 fn[K : Eq + Hash, V] T::union(Self[K, V], Self[K, V]) -> Self[K, V]
+fn[K : Eq + Hash, V] T::union_with(Self[K, V], Self[K, V], (K, V, V) -> V) -> Self[K, V]
 fn[K, V] T::values(Self[K, V]) -> Iter[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

@@ -180,7 +180,7 @@ pub fn[A] size(self : T[A]) -> Int {
 
 ///|
 /// Union two hashsets
-pub fn[K : Eq + Hash] union(self : T[K], other : T[K]) -> T[K] {
+pub fn[K : Eq + Hash] T::union(self : T[K], other : T[K]) -> T[K] {
   match (self, other) {
     (_, Empty) => self
     (Empty, _) => other
@@ -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] T::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] T::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")
+}