From 1f09efdda33c0dc8e3e3d167ef2b29de9a47b17a Mon Sep 17 00:00:00 2001
From: Luke Maurer <lmaurer@janestreet.com>
Date: Mon, 25 Nov 2024 19:04:29 +0000
Subject: [PATCH] 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.
---
 .../flambda2/algorithms/container_types.ml    | 16 +++++++++
 .../algorithms/container_types_intf.ml        |  2 ++
 .../flambda2/algorithms/patricia_tree.ml      | 35 ++++++++++++++-----
 .../tests/algorithms/patricia_tree_tests.ml   | 13 +++++++
 4 files changed, 58 insertions(+), 8 deletions(-)

diff --git a/middle_end/flambda2/algorithms/container_types.ml b/middle_end/flambda2/algorithms/container_types.ml
index 0eebe5781e5..a9aa73d960d 100644
--- a/middle_end/flambda2/algorithms/container_types.ml
+++ b/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]
 
diff --git a/middle_end/flambda2/algorithms/container_types_intf.ml b/middle_end/flambda2/algorithms/container_types_intf.ml
index 16fd523fdab..9813552233a 100644
--- a/middle_end/flambda2/algorithms/container_types_intf.ml
+++ b/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
diff --git a/middle_end/flambda2/algorithms/patricia_tree.ml b/middle_end/flambda2/algorithms/patricia_tree.ml
index 0f8aa88bbb4..7561bb9313d 100644
--- a/middle_end/flambda2/algorithms/patricia_tree.ml
+++ b/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
 
diff --git a/middle_end/flambda2/tests/algorithms/patricia_tree_tests.ml b/middle_end/flambda2/tests/algorithms/patricia_tree_tests.ml
index 749cb67da8c..2ab49f0cbb5 100644
--- a/middle_end/flambda2/tests/algorithms/patricia_tree_tests.ml
+++ b/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 () =