Index lifted cont params by variables instead of custom indexes (#3303)

Index lifted cont params y variables instead of custom indexes

This should have no observable effect, but will simplify the
continuation specialization work, particularly when paired with the
replay histories feature, so that lifted cont params can be correctly
tracked between a first pass and subsequent passes during continuation
specialization.

Author: Gbury

middle_end/flambda2/simplify/lifted_cont_params.ml

@@ -13,82 +13,57 @@
 (*                                                                        *)
 (**************************************************************************)
 
-module Id : sig
-  type t
-
-  val fresh : unit -> t
-
-  val print : Format.formatter -> t -> unit
-
-  module Map : Container_types.Map with type key = t
-end = struct
-  type t = int
-
-  let print = Numbers.Int.print
-
-  let fresh =
-    let r = ref 0 in
-    fun () ->
-      incr r;
-      !r
-
-  module Tree = Patricia_tree.Make (struct
-    let print = print
-  end)
-
-  module Map = Tree.Map
-end
+module BP = Bound_parameter
 
 type t =
   { len : int;
-    new_params_indexed : Bound_parameter.t Id.Map.t
+    new_params_indexed : BP.t BP.Map.t
   }
 
 let print ppf { len = _; new_params_indexed } =
   Format.fprintf ppf "@[<hov 1>(@[<hov 1>(new_params_indexed@ %a)@])@]"
-    (Id.Map.print Bound_parameter.print)
-    new_params_indexed
+    (BP.Map.print BP.print) new_params_indexed
 
-let empty = { len = 0; new_params_indexed = Id.Map.empty }
+let empty = { len = 0; new_params_indexed = BP.Map.empty }
 
 let is_empty { len; new_params_indexed = _ } = len = 0
 
 let length { len; new_params_indexed = _ } = len
 
 let new_param t bound_param =
-  (* create a fresh var/bound_param to index the new parameter *)
-  let id = Id.fresh () in
-  let new_params_indexed = Id.Map.add id bound_param t.new_params_indexed in
+  let new_params_indexed =
+    BP.Map.add bound_param bound_param t.new_params_indexed
+  in
   { len = t.len + 1; new_params_indexed }
 
 let rename t =
-  let bindings = Id.Map.bindings t.new_params_indexed in
+  let bindings = BP.Map.bindings t.new_params_indexed in
   let keys, bound_param_list = List.split bindings in
   let bound_params = Bound_parameters.create bound_param_list in
   let new_bound_params = Bound_parameters.rename bound_params in
   let renaming =
     Bound_parameters.renaming bound_params ~guaranteed_fresh:new_bound_params
   in
   let new_params_indexed =
-    Id.Map.of_list
+    BP.Map.of_list
       (List.combine keys (Bound_parameters.to_list new_bound_params))
   in
   { t with new_params_indexed }, renaming
 
 let fold_aux ~init ~f { len = _; new_params_indexed } =
-  Id.Map.fold f new_params_indexed init
+  BP.Map.fold f new_params_indexed init
 
 let fold ~init ~f t = fold_aux ~init ~f:(fun _ param acc -> f param acc) t
 
-let rec find_arg id = function
+let rec find_arg bp = function
   | [] ->
     Misc.fatal_errorf
       "Missing lifted param id: %a not found in lifted_cont_params stack"
-      Id.print id
+      BP.print bp
   | { len = _; new_params_indexed } :: r -> (
-    match Id.Map.find_opt id new_params_indexed with
-    | Some param -> Bound_parameter.simple param
-    | None -> find_arg id r)
+    match BP.Map.find_opt bp new_params_indexed with
+    | Some param -> BP.simple param
+    | None -> find_arg bp r)
 
 (* NOTE about the order of the returned args/params for the {args} and
    {bound_parameters} functions: The exact order does not matter as long as both
@@ -99,8 +74,8 @@ let rec find_arg id = function
    order of the bindings in the Map, but that's fine since it is the case for
    both functions. *)
 let args ~callee_lifted_params ~caller_stack_lifted_params =
-  fold_aux callee_lifted_params ~init:[] ~f:(fun id _callee_param acc ->
-      find_arg id caller_stack_lifted_params :: acc)
+  fold_aux callee_lifted_params ~init:[] ~f:(fun bp_key _callee_param acc ->
+      find_arg bp_key caller_stack_lifted_params :: acc)
 
 let bound_parameters t =
   Bound_parameters.create

middle_end/flambda2/simplify/lifted_cont_params.mli

@@ -32,9 +32,7 @@
     mapping for each continuation. Then, for a given call site, for each new
     parameter, we can lookup the original variable that introduced it, and then
     look into the stack of parent continuations to see which one first defines
-    a new (lifting) param that maps back to the same original variable. To avoid
-    confusing the original variables and the renamed parameters, we instead use
-    unique ids (which are in fact integers). *)
+    a new (lifting) param that maps back to the same original variable. *)
 
 (** This type represents all of the new params for one lifted continuation.
     These added parameters are internally indexed by a unique identifier: when an individual