diff --git a/map/map.mbt b/map/map.mbt index bd7aef43a..f8f745628 100644 --- a/map/map.mbt +++ b/map/map.mbt @@ -212,6 +212,70 @@ pub fn map_with_key[K, X, Y](self : Map[K, X], f : (K, X) -> Y) -> Map[K, Y] { } } +/// Fold the values in the map. +/// O(n). +pub fn fold[K, V, T](self : Map[K, V], f : (T, V) -> T, ~init : T) -> T { + foldl_with_key(self, fn(acc, _k, v) { f(acc, v) }, ~init) +} + +/// Post-order fold. +/// O(n). +pub fn foldr_with_key[K, V, T]( + self : Map[K, V], + f : (T, K, V) -> T, + ~init : T +) -> T { + fn go(m : Map[K, V], acc) { + match m { + Empty => acc + Tree(k, v, _, l, r) => go(l, f(go(r, acc), k, v)) + } + } + + go(self, init) +} + +/// Pre-order fold. +/// O(n). +pub fn foldl_with_key[K, V, T]( + self : Map[K, V], + f : (T, K, V) -> T, + ~init : T +) -> T { + fn go(m : Map[K, V], acc) { + match m { + Empty => acc + Tree(k, v, _, l, r) => go(r, f(go(l, acc), k, v)) + } + } + + go(self, init) +} + +/// Filter values that satisfy the predicate +pub fn filter[K : Compare, V]( + self : Map[K, V], + pred : (V) -> Bool +) -> Map[K, V] { + filter_with_key(self, fn(_k, v) { pred(v) }) +} + +/// Filter key-value pairs that satisfy the predicate +pub fn filter_with_key[K : Compare, V]( + self : Map[K, V], + pred : (K, V) -> Bool +) -> Map[K, V] { + match self { + Empty => Empty + Tree(k, v, _, l, r) => + if pred(k, v) { + balance(k, v, filter_with_key(l, pred), filter_with_key(r, pred)) + } else { + glue(filter_with_key(l, pred), filter_with_key(r, pred)) + } + } +} + /// The ratio between the sizes of the left and right subtrees. let ratio = 5 @@ -409,6 +473,39 @@ test "map_with_key" { )? } +test "fold" { + let m = Map::[("a", 1), ("b", 2), ("c", 3), ("d", 4), ("e", 5)] + @assertion.assert_eq(m.fold(fn(acc, v) { acc + v }, ~init=0), 15)? +} + +test "foldl_with_key" { + let m = Map::[("a", 1), ("b", 2), ("c", 3), ("d", 4), ("e", 5)] + @assertion.assert_eq( + m.foldl_with_key(fn(acc, k, v) { acc + k + v.to_string() }, ~init=""), + "a1b2c3d4e5", + )? +} + +test "foldr_with_key" { + let m = Map::[("a", 1), ("b", 2), ("c", 3), ("d", 4), ("e", 5)] + @assertion.assert_eq( + m.foldr_with_key(fn(acc, k, v) { acc + k + v.to_string() }, ~init=""), + "e5d4c3b2a1", + )? +} + +test "filter" { + let m = Map::[(3, "three"), (8, "eight"), (1, "one"), (2, "two"), (0, "zero")] + let fm = m.filter(fn(v) { v.length() > 3 }) + @assertion.assert_eq(fm.debug_tree(), "(3,three,(0,zero,_,_),(8,eight,_,_))")? +} + +test "filter_with_key" { + let m = Map::[(3, "three"), (8, "eight"), (1, "one"), (2, "two"), (0, "zero")] + let fm = m.filter_with_key(fn(k, v) { k > 3 && v.length() > 3 }) + @assertion.assert_eq(fm.debug_tree(), "(8,eight,_,_)")? +} + test "singleton" { let m = singleton(3, "three") @assertion.assert_eq(m.debug_tree(), "(3,three,_,_)")?