Add fold and filter to Map (#83)

* Add `fold` to Map

* Add `filter` to Map

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,_,_)")?