Implement fast filter_map and filter_map_sharing for Patricia trees

We're seeing cases where compilation time is bottlenecked on a
`Variable.Map.filter_map` in `To_cmm_env`. This has a stub O(n log n)
implementation, which can easily be improved to O(n). Furthermore, the original
element is often being returned unchanged, so there's hopefully some further
improvement to be gained by returning the (sub)tree unchanged in those cases.

middle_end/flambda2/algorithms/container_types.ml

@@ -135,6 +135,22 @@ module Make_map (T : Thing) (Set : Set_plus_stdlib with type elt = T.t) = struct
         t
     in
     if not !changed then t else t'
+
+  let filter_map_sharing f t =
+    let changed = ref false in
+    let t' =
+      filter_map
+        (fun k v ->
+          let v' = f k v in
+          let () =
+            match v' with
+            | Some v' -> if not (v == v') then changed := true
+            | None -> changed := true
+          in
+          v')
+        t
+    in
+    if not !changed then t else t'
 end
 [@@inline always]
 

middle_end/flambda2/algorithms/container_types_intf.ml

@@ -158,6 +158,8 @@ module type Map = sig
 
   val filter_map : (key -> 'a -> 'b option) -> 'a t -> 'b t
 
+  val filter_map_sharing : (key -> 'a -> 'a option) -> 'a t -> 'a t
+
   val partition : (key -> 'a -> bool) -> 'a t -> 'a t * 'a t
 
   val cardinal : 'a t -> int

middle_end/flambda2/algorithms/patricia_tree.ml

@@ -308,6 +308,10 @@ module Tree_operations (Tree : Tree) : sig
 
   val mapi : 'b is_value -> ('a, 'b) Callback.t -> 'a t -> 'b t
 
+  val filter_map : 'b is_value -> (key -> 'a -> 'b option) -> 'a t -> 'b t
+
+  val filter_map_sharing : (key -> 'a -> 'a option) -> 'a t -> 'a t
+
   val to_seq : 'a t -> 'a Binding.t Seq.t
 
   val of_list : 'a is_value -> 'a Binding.t list -> 'a t
@@ -845,6 +849,28 @@ end = struct
     | Branch (prefix, bit, t0, t1) ->
       branch_non_empty prefix bit (mapi iv f t0) (mapi iv f t1)
 
+  let rec filter_map iv f t =
+    match descr t with
+    | Empty -> empty iv
+    | Leaf (k, d) -> (
+      match f k d with None -> empty iv | Some d' -> leaf iv k d')
+    | Branch (prefix, bit, t0, t1) ->
+      branch prefix bit (filter_map iv f t0) (filter_map iv f t1)
+
+  let rec filter_map_sharing f t =
+    let iv = is_value_of t in
+    match descr t with
+    | Empty -> t
+    | Leaf (k, d) -> (
+      match f k d with
+      | None -> empty iv
+      | Some d' when d == d' -> t
+      | Some d' -> leaf iv k d')
+    | Branch (prefix, bit, t0, t1) ->
+      let t0' = filter_map_sharing f t0 in
+      let t1' = filter_map_sharing f t1 in
+      if t0' == t0 && t1' == t1 then t else branch prefix bit t0' t1'
+
   let to_seq t =
     let rec aux acc () =
       match acc with
@@ -988,14 +1014,7 @@ module Map = struct
 
   let mapi f t = Ops.mapi Any f t
 
-  (* CR-someday lmaurer: Implement this in [Ops] in O(n) time rather than O(n
-     log n). Should be able to implement [filter] in terms of it, though
-     Flambda2 currently has trouble preserving sharing (unnecessary control flow
-     obscures a CSE opportunity). *)
-  let filter_map f t =
-    fold
-      (fun id v map -> match f id v with None -> map | Some r -> add id r map)
-      t empty
+  let filter_map f t = Ops.filter_map Any f t
 
   let of_list l = Ops.of_list Any l
 

middle_end/flambda2/tests/algorithms/patricia_tree_tests.ml

@@ -651,6 +651,14 @@ module Map_specs (V : Value) = struct
     Map.equal V.equal (Map.map_sharing f m) (Map.map f m)
 
   let map_sharing_id m = Map.map_sharing (fun v -> v) m == m
+
+  let filter_map_sharing_valid f m = Map.valid (Map.filter_map_sharing f m)
+
+  let filter_map_sharing_vs_filter_map f m =
+    Map.equal V.equal (Map.filter_map_sharing f m) (Map.filter_map f m)
+
+  let filter_map_sharing_id m =
+    Map.filter_map_sharing (fun _k v -> Some v) m == m
 end
 
 (* CR-someday lmaurer: Move the [Abitrary.t] for perms into a separate module
@@ -1132,6 +1140,11 @@ let () =
     c "map_sharing is valid" map_sharing_valid [value_to_value; map];
     c "map_sharing vs. map" map_sharing_vs_map [value_to_value; map];
     c "map_sharing of id" map_sharing_id [map];
+    c "filter_map_sharing is valid" filter_map_sharing_valid
+      [key_and_value_to_value_option; map];
+    c "filter_map_sharing vs. filter_map" filter_map_sharing_vs_filter_map
+      [key_and_value_to_value_option; map];
+    c "filter_map_sharing of id" filter_map_sharing_id [map];
     ()
   in
   let () =