diff --git a/builtin/builtin.mbti b/builtin/builtin.mbti index 42e820729..6aaa3b5e1 100644 --- a/builtin/builtin.mbti +++ b/builtin/builtin.mbti @@ -257,6 +257,8 @@ fn[K : Hash + Eq, V] Map::contains(Self[K, V], K) -> Bool fn[K : Hash + Eq, V : Eq] Map::contains_kv(Self[K, V], K, V) -> Bool fn[K, V] Map::each(Self[K, V], (K, V) -> Unit?Error) -> Unit?Error fn[K, V] Map::eachi(Self[K, V], (Int, K, V) -> Unit?Error) -> Unit?Error +fn[K : Hash + Eq, V] Map::filter(Self[K, V], (K, V) -> Bool) -> Self[K, V] +fn[K : Hash + Eq, V, A] Map::fold(Self[K, V], A, (A, K, V) -> A) -> A fn[K : Hash + Eq, V] Map::from_array(Array[(K, V)]) -> Self[K, V] fn[K : Hash + Eq, V] Map::from_iter(Iter[(K, V)]) -> Self[K, V] fn[K : Hash + Eq, V] Map::get(Self[K, V], K) -> V? @@ -266,11 +268,14 @@ fn[K, V] Map::is_empty(Self[K, V]) -> Bool fn[K, V] Map::iter(Self[K, V]) -> Iter[(K, V)] fn[K, V] Map::iter2(Self[K, V]) -> Iter2[K, V] fn[K, V] Map::keys(Self[K, V]) -> Iter[K] +fn[K : Hash + Eq, V, V2] Map::map(Self[K, V], (V) -> V2) -> Self[K, V2] +fn[K : Hash + Eq, V, V2] Map::map_with_key(Self[K, V], (K, V) -> V2) -> Self[K, V2] fn[K, V] Map::new(capacity~ : Int = ..) -> Self[K, V] fn[K : Hash + Eq, V] Map::of(FixedArray[(K, V)]) -> Self[K, V] #deprecated fn[K : Hash + Eq, V] Map::op_get(Self[K, V], K) -> V? fn[K : Hash + Eq, V] Map::op_set(Self[K, V], K, V) -> Unit +fn[K : Hash + Eq, V] Map::partition(Self[K, V], (K, V) -> Bool) -> (Self[K, V], Self[K, V]) fn[K : Hash + Eq, V] Map::remove(Self[K, V], K) -> Unit fn[K : Hash + Eq, V] Map::set(Self[K, V], K, V) -> Unit fn[K, V] Map::size(Self[K, V]) -> Int diff --git a/builtin/linked_hash_map.mbt b/builtin/linked_hash_map.mbt index afbf52794..0b6b708a9 100644 --- a/builtin/linked_hash_map.mbt +++ b/builtin/linked_hash_map.mbt @@ -590,3 +590,85 @@ pub fn[K : Hash + Eq, V] Map::from_iter(iter : Iter[(K, V)]) -> Map[K, V] { pub impl[K, V] Default for Map[K, V] with default() { Map::new() } + +///| +/// Applies a function to each key-value pair in the map and returns a new map with the results. +pub fn[K : Hash + Eq, V, V2] Map::map( + self : Map[K, V], + f : (V) -> V2 +) -> Map[K, V2] { + self.map_with_key(fn(_, v) { f(v) }) +} + +///| +/// Applies a function to each key-value pair in the map and returns a new map with the results, using the original keys. +pub fn[K : Hash + Eq, V, V2] Map::map_with_key( + self : Map[K, V], + f : (K, V) -> V2 +) -> Map[K, V2] { + if self.is_empty() { + return Map::new() + } + let m = Map::new(capacity=self.size) + for e in self.iter() { + m.set(e.0, f(e.0, e.1)) + } + m +} + +///| +/// Filters the map by applying a predicate function to each key-value pair. +pub fn[K : Hash + Eq, V] Map::filter( + self : Map[K, V], + pred : (K, V) -> Bool +) -> Map[K, V] { + if self.is_empty() { + return Map::new() + } + let m = Map::new(capacity=self.size) + for e in self.iter() { + if pred(e.0, e.1) { + m.set(e.0, e.1) + } + } + m +} + +///| +/// Folds the map from an initial value by applying a function to each key-value pair and the accumulator. +/// The order of folding is not guaranteed. +pub fn[K : Hash + Eq, V, A] Map::fold( + self : Map[K, V], + init : A, + f : (A, K, V) -> A +) -> A { + if self.is_empty() { + return init + } + let mut acc = init + for e in self { + acc = f(acc, e.0, e.1) + } + acc +} + +///| +/// Partitions the map into two maps based on a predicate function. +pub fn[K : Hash + Eq, V] Map::partition( + self : Map[K, V], + pred : (K, V) -> Bool +) -> (Map[K, V], Map[K, V]) { + if self.is_empty() { + return (Map::new(), Map::new()) + } + let m1 = Map::new(capacity=self.size) + let m2 = Map::new(capacity=self.size) + for e in self { + if pred(e.0, e.1) { + m1.set(e.0, e.1) + } else { + m2.set(e.0, e.1) + } + } + (m1, m2) +} diff --git a/immut/hashmap/HAMT_test.mbt b/immut/hashmap/HAMT_test.mbt index 57712f6a7..3f41615ab 100644 --- a/immut/hashmap/HAMT_test.mbt +++ b/immut/hashmap/HAMT_test.mbt @@ -673,6 +673,99 @@ test "HAMT::difference with branch" { assert_eq!(diff.get(4), None) } +///| +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()