Add replay history to match names between two downwards traversals

middle_end/flambda2/algorithms/lmap.ml

@@ -61,6 +61,8 @@ module type S = sig
 
   val exists : (key -> 'a -> bool) -> 'a t -> bool
 
+  val choose : 'a t -> key * 'a
+
   val add_seq : (key * 'a) Seq.t -> 'a t -> 'a t
 
   val of_seq : (key * 'a) Seq.t -> 'a t
@@ -108,6 +110,8 @@ module Make (T : Thing) : S with type key = T.t = struct
 
   let exists f m = List.exists (fun (k, v) -> f k v) m
 
+  let choose = function [] -> raise Not_found | res :: _ -> res
+
   let keys m = List.map fst m
 
   let data m = List.map snd m

middle_end/flambda2/algorithms/lmap.mli

@@ -105,6 +105,10 @@ module type S = sig
 
   val exists : (key -> 'a -> bool) -> 'a t -> bool
 
+  (** Returns an unspecified binding from the map
+      @raise Not_found if the map is empty *)
+  val choose : 'a t -> key * 'a
+
   (** Keys in the sequence must be distinct from each other and from keys
       already in the map; neither of these conditions is checked. *)
   val add_seq : (key * 'a) Seq.t -> 'a t -> 'a t

middle_end/flambda2/bound_identifiers/bound_continuations.ml

@@ -50,6 +50,9 @@ let ids_for_export t =
 
 let rename t = List.map Continuation.rename t
 
+let is_renamed_version_of t t' =
+  Misc.Stdlib.List.equal Continuation.is_renamed_version_of t t'
+
 let renaming t1 ~guaranteed_fresh:t2 =
   try List.fold_left2 Renaming.add_continuation Renaming.empty t1 t2
   with Invalid_argument _ ->

middle_end/flambda2/bound_identifiers/bound_for_function.ml

@@ -180,6 +180,16 @@ let rename
     my_depth = Variable.rename my_depth
   }
 
+let is_renamed_version_of t t' =
+  Continuation.is_renamed_version_of t.return_continuation
+    t'.return_continuation
+  && Continuation.is_renamed_version_of t.exn_continuation t'.exn_continuation
+  && Bound_parameters.is_renamed_version_of t.params t'.params
+  && Variable.is_renamed_version_of t.my_closure t'.my_closure
+  && Variable.is_renamed_version_of t.my_region t'.my_region
+  && Variable.is_renamed_version_of t.my_ghost_region t'.my_ghost_region
+  && Variable.is_renamed_version_of t.my_depth t'.my_depth
+
 let renaming
     { return_continuation = return_continuation1;
       exn_continuation = exn_continuation1;

middle_end/flambda2/bound_identifiers/bound_parameter.ml

@@ -56,6 +56,10 @@ let with_kind t kind = { t with kind }
 
 let rename t = { t with param = Variable.rename t.param }
 
+let is_renamed_version_of t t' =
+  Flambda_kind.With_subkind.equal t.kind t'.kind
+  && Variable.is_renamed_version_of t.param t'.param
+
 let equal_kinds t1 t2 = Flambda_kind.With_subkind.equal t1.kind t2.kind
 
 let free_names ({ param = _; kind = _ } as t) =

middle_end/flambda2/bound_identifiers/bound_parameter.mli

@@ -40,6 +40,8 @@ val equal_kinds : t -> t -> bool
 
 val rename : t -> t
 
+val is_renamed_version_of : t -> t -> bool
+
 include Container_types.S with type t := t
 
 include Contains_names.S with type t := t

middle_end/flambda2/bound_identifiers/bound_parameters.ml

@@ -60,6 +60,9 @@ let var_set t = Variable.Set.of_list (vars t)
 
 let rename t = List.map (fun t -> BP.rename t) t
 
+let is_renamed_version_of t t' =
+  Misc.Stdlib.List.equal BP.is_renamed_version_of t t'
+
 let arity t =
   List.map
     (fun t -> Flambda_arity.Component_for_creation.Singleton (BP.kind t))

middle_end/flambda2/bound_identifiers/bound_pattern.ml

@@ -88,6 +88,21 @@ let rename t =
     Set_of_closures bound_vars
   | Static _ -> t
 
+let is_renamed_version_of t t' =
+  match t, t' with
+  | Singleton bound_var, Singleton bound_var' ->
+    Bound_var.is_renamed_version_of bound_var bound_var'
+  | Set_of_closures bound_vars, Set_of_closures bound_vars' ->
+    Misc.Stdlib.List.equal Bound_var.is_renamed_version_of bound_vars
+      bound_vars'
+  | Static _bound_static, Static _bound_static' ->
+    (* CR gbury/ncourant: We should try and compare the bound statics here *)
+    true
+  | Singleton _, (Set_of_closures _ | Static _)
+  | Set_of_closures _, (Singleton _ | Static _)
+  | Static _, (Singleton _ | Set_of_closures _) ->
+    false
+
 let renaming t1 ~guaranteed_fresh:t2 =
   match t1, t2 with
   | Singleton bound_var1, Singleton bound_var2 ->

middle_end/flambda2/bound_identifiers/bound_var.ml

@@ -37,6 +37,10 @@ let with_name_mode t name_mode = { t with name_mode }
 
 let rename t = with_var t (Variable.rename t.var)
 
+let is_renamed_version_of t t' =
+  Name_mode.equal t.name_mode t'.name_mode
+  && Variable.is_renamed_version_of t.var t'.var
+
 let apply_renaming t renaming =
   with_var t (Renaming.apply_variable renaming t.var)
 

middle_end/flambda2/compare/compare.ml

@@ -419,7 +419,7 @@ and subst_let_cont env (let_cont_expr : Let_cont_expr.t) =
       ~f:(fun ~invariant_params ~body handlers ->
         let body = subst_expr env body in
         let handlers =
-          Continuation.Map.map_sharing (subst_cont_handler env)
+          Continuation.Lmap.map_sharing (subst_cont_handler env)
             (handlers |> Continuation_handlers.to_map)
         in
         Let_cont_expr.create_recursive handlers ~invariant_params ~body)
@@ -1283,7 +1283,7 @@ and let_cont_exprs env (let_cont1 : Let_cont.t) (let_cont2 : Let_cont.t) :
                  ~free_names_of_body:Unknown))
   | Recursive handlers1, Recursive handlers2 ->
     let compare_handler_maps env map1 map2 :
-        Continuation_handler.t Continuation.Map.t Comparison.t =
+        Continuation_handler.t Continuation.Lmap.t Comparison.t =
       lists
         ~f:(fun env (cont, handler1) (_cont, handler2) ->
           cont_handlers env handler1 handler2
@@ -1292,15 +1292,15 @@ and let_cont_exprs env (let_cont1 : Let_cont.t) (let_cont2 : Let_cont.t) :
           |> Comparison.map ~f:(fun handler1' -> cont, handler1'))
         ~subst:(fun env (cont, handler) -> cont, subst_cont_handler env handler)
         ~subst_snd:false env
-        (map1 |> Continuation.Map.bindings)
-        (map2 |> Continuation.Map.bindings)
-      |> Comparison.map ~f:Continuation.Map.of_list
+        (map1 |> Continuation.Lmap.bindings)
+        (map2 |> Continuation.Lmap.bindings)
+      |> Comparison.map ~f:Continuation.Lmap.of_list
     in
     Recursive_let_cont_handlers.pattern_match_pair handlers1 handlers2
       ~f:(fun ~invariant_params ~body1 ~body2 cont_handlers1 cont_handlers2 ->
         pairs ~f1:exprs ~f2:compare_handler_maps
           ~subst2:(fun env map ->
-            Continuation.Map.map_sharing (subst_cont_handler env) map)
+            Continuation.Lmap.map (subst_cont_handler env) map)
           env
           (body1, cont_handlers1 |> Continuation_handlers.to_map)
           (body2, cont_handlers2 |> Continuation_handlers.to_map)

middle_end/flambda2/from_lambda/closure_conversion_aux.ml

@@ -1138,7 +1138,7 @@ module Let_cont_with_acc = struct
     in
     let expr = Let_cont.create_recursive ~invariant_params handlers ~body in
     let acc =
-      Continuation.Map.fold
+      Continuation.Lmap.fold
         (fun cont _ acc -> Acc.remove_continuation_from_free_names cont acc)
         handlers acc
     in
@@ -1158,9 +1158,9 @@ module Let_cont_with_acc = struct
           ( Name_occurrences.union free_names handler_free_names,
             Cost_metrics.( + ) costs cost_metrics_of_handler,
             acc,
-            Continuation.Map.add cont handler handlers ))
+            Continuation.Lmap.add cont handler handlers ))
         handlers
-        (Name_occurrences.empty, Cost_metrics.zero, acc, Continuation.Map.empty)
+        (Name_occurrences.empty, Cost_metrics.zero, acc, Continuation.Lmap.empty)
     in
     let body_free_names, acc, body = Acc.eval_branch_free_names acc ~f:body in
     let acc =

middle_end/flambda2/identifiers/continuation.ml

@@ -133,6 +133,11 @@ let rename t =
   let { Data.name; sort; name_stamp = _; compilation_unit = _ } = find_data t in
   create ~sort ~name ()
 
+let is_renamed_version_of t t' =
+  let data = find_data t in
+  let data' = find_data t' in
+  Sort.equal data.sort data'.sort && String.equal data.name data'.name
+
 let name t = (find_data t).name
 
 let name_stamp t = (find_data t).name_stamp
@@ -165,6 +170,7 @@ end
 module Tree = Patricia_tree.Make (T)
 module Set = Tree.Set
 module Map = Tree.Map
+module Lmap = Lmap.Make (T)
 
 let export t = find_data t
 

middle_end/flambda2/identifiers/continuation.mli

@@ -22,6 +22,8 @@ type exported
 
 include Container_types.S with type t := t
 
+module Lmap : Lmap.S with type key := t
+
 module Sort : sig
   type t =
     | Normal_or_exn
@@ -36,6 +38,8 @@ val create : ?sort:Sort.t -> ?name:string -> unit -> t
 
 val rename : t -> t
 
+val is_renamed_version_of : t -> t -> bool
+
 val name : t -> string
 
 val sort : t -> Sort.t

middle_end/flambda2/identifiers/variable.ml

@@ -24,6 +24,10 @@ let rename ?append t =
   let user_visible = if user_visible t then Some () else None in
   create ?user_visible name
 
+let is_renamed_version_of t t' =
+  (* We only keep track of variables renamed with an empty {append} parameter *)
+  String.equal (name t) (name t')
+
 let raw_name = name
 
 let unique_name t = name t ^ string_of_int (name_stamp t)

middle_end/flambda2/identifiers/variable.mli

@@ -24,6 +24,8 @@ val create_with_same_name_as_ident : ?user_visible:unit -> Ident.t -> t
     the current unit, not the unit of the variable passed in. *)
 val rename : ?append:string -> t -> t
 
+val is_renamed_version_of : t -> t -> bool
+
 val unique_name : t -> string
 
 val raw_name : t -> string

middle_end/flambda2/nominal/bindable.ml

@@ -29,6 +29,17 @@ module type S = sig
   (** Freshen the given name. *)
   val rename : t -> t
 
+  (** Equivalence relation on renamed variables.
+
+      [is_renamed_version_of x y] is [true] if there exists a bindable [z]
+      such that [x] and [y] are renamed versions of [z].
+
+      Note: this function can return [true] in other cases (if there are some name collisions
+      for instance), this is (at least currently) only used for a sanity check, so
+      users should not rely too much on its expected semantics.
+      *)
+  val is_renamed_version_of : t -> t -> bool
+
   (** [renaming stale ~guaranteed_fresh:fresh] is to create a renaming that
       turns all occurrences of the name [stale] into [fresh] (in a
       capture-avoiding manner, but that is inherent in [Renaming]). *)

middle_end/flambda2/parser/fexpr_to_flambda.ml

@@ -700,7 +700,7 @@ let rec expr env (e : Fexpr.expr) : Flambda.Expr.t =
       Flambda.Let_cont.create_non_recursive name handler ~body
         ~free_names_of_body:Unknown
     | Recursive ->
-      let handlers = Continuation.Map.singleton name handler in
+      let handlers = Continuation.Lmap.singleton name handler in
       Flambda.Let_cont.create_recursive ~invariant_params:Bound_parameters.empty
         handlers ~body)
   | Let_cont _ -> failwith "TODO andwhere"

middle_end/flambda2/parser/flambda_to_fexpr.ml

@@ -1006,7 +1006,7 @@ and let_cont_expr env (lc : Flambda.Let_cont_expr.t) =
       ~f:(fun ~invariant_params:_ ~body handlers ->
         (* TODO support them *)
         let env =
-          Continuation.Set.fold
+          List.fold_right
             (fun c env ->
               let _, env = Env.bind_named_continuation env c in
               env)
@@ -1024,7 +1024,7 @@ and let_cont_expr env (lc : Flambda.Let_cont_expr.t) =
               in
               cont_handler env c sort handler)
             (handlers |> Flambda.Continuation_handlers.to_map
-           |> Continuation.Map.bindings)
+           |> Continuation.Lmap.bindings)
         in
         let body = expr env body in
         Fexpr.Let_cont { recursive = Recursive; bindings; body })
@@ -1227,7 +1227,7 @@ module Iter = struct
 
   and let_cont_rec f_c f_s conts body =
     let map = Continuation_handlers.to_map conts in
-    Continuation.Map.iter (continuation_handler f_c f_s) map;
+    Continuation.Lmap.iter (continuation_handler f_c f_s) map;
     expr f_c f_s body
 
   and continuation_handler f_c f_s _ h =

middle_end/flambda2/reaper/rebuild.ml

@@ -518,7 +518,7 @@ and rebuild_holed (kinds : Flambda_kind.t Name.Map.t) (env : env)
            (Bound_parameters.to_list params))
     in
     let handlers =
-      Continuation.Map.mapi
+      Continuation.Lmap.mapi
         (fun cont handler ->
           let { bound_parameters; expr; is_exn_handler; is_cold } = handler in
           let bound_parameters = filter_params cont bound_parameters in
@@ -528,9 +528,7 @@ and rebuild_holed (kinds : Flambda_kind.t Name.Map.t) (env : env)
         handlers
     in
     let invariant_params =
-      filter_params
-        (fst (Continuation.Map.min_binding handlers))
-        invariant_params
+      filter_params (fst (Continuation.Lmap.choose handlers)) invariant_params
     in
     let let_cont_expr =
       RE.create_recursive_let_cont ~invariant_params handlers ~body:hole

middle_end/flambda2/reaper/rebuilt_expr.ml

@@ -19,7 +19,7 @@ type continuation_handler =
   }
 
 type continuation_handlers =
-  { handlers : Flambda.Continuation_handler.t Continuation.Map.t;
+  { handlers : Flambda.Continuation_handler.t Continuation.Lmap.t;
     free_names : Name_occurrences.t
   }
 
@@ -56,13 +56,14 @@ let create_continuation_handler bound_parameters ~handler ~is_exn_handler
   { handler; free_names }
 
 let create_continuation_handlers handlers =
-  Continuation.Map.fold
-    (fun cont (handler : continuation_handler) { handlers; free_names } ->
-      let handlers = Continuation.Map.add cont handler.handler handlers in
-      let free_names = Name_occurrences.union free_names handler.free_names in
-      { handlers; free_names })
-    handlers
-    { handlers = Continuation.Map.empty; free_names = Name_occurrences.empty }
+  let free_names, handlers =
+    Continuation.Lmap.fold_left_map
+      (fun free_names _cont (handler : continuation_handler) ->
+        let free_names = Name_occurrences.union free_names handler.free_names in
+        free_names, handler.handler)
+      Name_occurrences.empty handlers
+  in
+  { handlers; free_names }
 
 let create_non_recursive_let_cont cont (cont_handler : continuation_handler)
     ~body =
@@ -92,7 +93,7 @@ let create_recursive_let_cont ~invariant_params handlers0 ~body =
       (Name_occurrences.increase_counts handlers_free_names)
   in
   let free_names =
-    Continuation.Map.fold
+    Continuation.Lmap.fold
       (fun cont _ free_names ->
         Name_occurrences.remove_continuation free_names ~continuation:cont)
       handlers0 free_names

middle_end/flambda2/reaper/rebuilt_expr.mli

@@ -19,7 +19,7 @@ type continuation_handler =
   }
 
 type continuation_handlers =
-  { handlers : Flambda.Continuation_handler.t Continuation.Map.t;
+  { handlers : Flambda.Continuation_handler.t Continuation.Lmap.t;
     free_names : Name_occurrences.t
   }
 
@@ -38,14 +38,14 @@ val create_continuation_handler :
   continuation_handler
 
 val create_continuation_handlers :
-  continuation_handler Continuation.Map.t -> continuation_handlers
+  continuation_handler Continuation.Lmap.t -> continuation_handlers
 
 val create_non_recursive_let_cont :
   Continuation.t -> continuation_handler -> body:t -> t
 
 val create_recursive_let_cont :
   invariant_params:Bound_parameters.t ->
-  continuation_handler Continuation.Map.t ->
+  continuation_handler Continuation.Lmap.t ->
   body:t ->
   t
 

middle_end/flambda2/reaper/rev_expr.ml

@@ -33,7 +33,7 @@ type rev_expr_holed =
       }
   | Let_cont_rec of
       { invariant_params : Bound_parameters.t;
-        handlers : cont_handler Continuation.Map.t;
+        handlers : cont_handler Continuation.Lmap.t;
         parent : rev_expr_holed
       }
 

middle_end/flambda2/reaper/rev_expr.mli

@@ -33,7 +33,7 @@ type rev_expr_holed =
       }
   | Let_cont_rec of
       { invariant_params : Bound_parameters.t;
-        handlers : cont_handler Continuation.Map.t;
+        handlers : cont_handler Continuation.Lmap.t;
         parent : rev_expr_holed
       }