Rewrite apply_cont_rewrite

middle_end/flambda2/simplify/apply_cont_rewrite.ml

@@ -14,135 +14,218 @@
 (*                                                                        *)
 (**************************************************************************)
 
-module EA = Continuation_extra_params_and_args.Extra_arg
-module BP = Bound_parameter
+module EPA = Continuation_extra_params_and_args
+module EA = EPA.Extra_arg
 module Id = Apply_cont_rewrite_id
 
 type used =
+  | Unused
   | Used
   | Used_as_invariant
-  | Unused
 
 type t =
-  { original_params : Bound_parameters.t;
-    used_params : BP.Set.t;
-    used_extra_params : Bound_parameters.t;
-    used_extra_invariant_params : Bound_parameters.t;
-    invariant_params : BP.Set.t;
-    extra_args : (EA.t * used) list Id.Map.t;
+  {
+    original_params_usage : used list ;
+    extra_params_usage : used list ;
+    extra_args : EA.t list Id.Map.t ;
+    original_params : Bound_parameters.t ;
+    extra_params : Bound_parameters.t ;
   }
 
 let print_used ppf = function
-  | Used -> ()
-  | Used_as_invariant -> Format.fprintf ppf "@ invariant"
-  | Unused -> Format.fprintf ppf "@ unused"
+  | Unused -> Format.fprintf ppf "unused"
+  | Used -> Format.fprintf ppf "used"
+  | Used_as_invariant -> Format.fprintf ppf "invariant"
 
-let print_ea_used ppf t =
+let print_params_used ppf (params, usage) =
   Format.fprintf ppf "(%a)"
-    (Format.pp_print_list ~pp_sep:Format.pp_print_space (fun ppf (ea, used) ->
-         Format.fprintf ppf "%a%a" EA.print ea print_used used))
-    t
+    (Format.pp_print_list ~pp_sep:Format.pp_print_space (fun ppf (param, used) ->
+         Format.fprintf ppf "(%a:@ %a)" Bound_parameter.print param print_used used))
+    (List.combine (Bound_parameters.to_list params) usage)
 
 let [@ocamlformat "disable"] print ppf
-  { original_params; used_params; used_extra_params; used_extra_invariant_params; invariant_params; extra_args; } =
+  { original_params_usage; extra_params_usage; extra_args;
+    original_params; extra_params } =
   Format.fprintf ppf "@[<hov 1>(\
       @[<hov 1>(original_params@ (%a))@]@ \
-      @[<hov 1>(used_params@ %a)@]@ \
-      @[<hov 1>(used_extra_params@ (%a))@]@ \
-      @[<hov 1>(used_extra_invariant_params@ (%a))@]@ \
-      @[<hov 1>(invariant_params@ %a)@]@ \
+      @[<hov 1>(extra_params@ %a)@]@ \
       @[<hov 1>(extra_args@ %a)@]\
       )@]"
-  Bound_parameters.print original_params
-    BP.Set.print used_params
-    Bound_parameters.print used_extra_params
-    Bound_parameters.print used_extra_invariant_params
-    BP.Set.print invariant_params
-    (Id.Map.print print_ea_used) extra_args
+    print_params_used (original_params, original_params_usage)
+    print_params_used (extra_params, extra_params_usage)
+    (Id.Map.print
+       (Format.pp_print_list ~pp_sep:Format.pp_print_space EA.print))
+    extra_args
 
 let does_nothing t =
-  Bound_parameters.cardinal t.original_params = BP.Set.cardinal t.used_params
-  && Bound_parameters.is_empty t.used_extra_params
-  && Bound_parameters.is_empty t.used_extra_invariant_params
-  && BP.Set.is_empty t.invariant_params
+  List.for_all (function | Used -> true | Used_as_invariant | Unused -> false)
+    t.original_params_usage &&
+  List.for_all (function | Unused -> true | Used | Used_as_invariant -> false)
+    t.extra_params_usage
 
-let create ~original_params ~used_params ~invariant_params ~extra_params ~extra_args
-    ~used_extra_params =
+let create ~original_params ~extra_params_and_args ~decide_param_usage =
+  let extra_params = EPA.extra_params extra_params_and_args in
+  let extra_args = EPA.extra_args extra_params_and_args in
   Bound_parameters.check_no_duplicates original_params;
   Bound_parameters.check_no_duplicates extra_params;
-  if Bound_parameters.cardinal original_params < BP.Set.cardinal used_params
-  then
-    Misc.fatal_errorf
-      "Must have at least as many [original_params] (%a)@ as [used_params] (%a)"
-      Bound_parameters.print original_params BP.Set.print used_params;
-  if Bound_parameters.cardinal extra_params < BP.Set.cardinal used_extra_params
-  then
-    Misc.fatal_errorf
-      "Must have at least as many [extra_params] (%a)@ as [used_extra_params] \
-       (%a)"
-      Bound_parameters.print extra_params BP.Set.print used_extra_params;
-  let extra_args =
-    if Bound_parameters.is_empty extra_params
-    then Id.Map.empty
-    else
-      let num_extra_params = Bound_parameters.cardinal extra_params in
-      Id.Map.map
-        (fun extra_args ->
-          if num_extra_params <> List.length extra_args
-          then
-            Misc.fatal_errorf
-              "Lengths of [extra_params] (%a)@ and all [extra_args] (e.g. %a) \
-               should be equal"
-              Bound_parameters.print extra_params
-              Continuation_extra_params_and_args.Extra_arg.List.print extra_args;
-          List.map2
-            (fun extra_param extra_arg ->
-              ( extra_arg,
-                if BP.Set.mem extra_param used_extra_params
-                then if BP.Set.mem extra_param invariant_params then Used_as_invariant else Used
-                else Unused ))
-            (Bound_parameters.to_list extra_params)
-            extra_args)
-        extra_args
-  in
-  let used_extra_params =
-    Bound_parameters.filter
-      (fun extra_param -> BP.Set.mem extra_param used_extra_params)
-      extra_params
-  in
-  let used_extra_invariant_params, used_extra_params =
-    List.partition (fun extra_param -> BP.Set.mem extra_param invariant_params)
-      (Bound_parameters.to_list used_extra_params)
+  let original_params_usage =
+    List.map decide_param_usage (Bound_parameters.to_list original_params)
   in
-  let used_extra_invariant_params, used_extra_params =
-    Bound_parameters.create used_extra_invariant_params, Bound_parameters.create used_extra_params
+  let extra_params_usage =
+    List.map decide_param_usage (Bound_parameters.to_list extra_params)
   in
-  let invariant_params =
-    BP.Set.inter invariant_params (BP.Set.of_list (Bound_parameters.to_list original_params))
-  in
-  { original_params; used_params; used_extra_params; used_extra_invariant_params; invariant_params; extra_args }
+  { original_params_usage; extra_params_usage; extra_args;
+    original_params; extra_params }
 
-let original_params t = t.original_params
+let original_params_arity t =
+  Bound_parameters.arity_with_subkinds t.original_params
 
-let used_params t = t.used_params
+let rec partition_used l usage =
+  match l, usage with
+  | [], [] -> [], []
+  | [], _ :: _ | _ :: _, [] ->
+    Misc.fatal_error "Apply_cont_rewrite.partition_used: list and usage must have the same length"
+  | _ :: l, Unused :: usage ->
+    partition_used l usage
+  | x :: l, Used :: usage ->
+    let invariant, normal = partition_used l usage in
+    invariant, x :: normal
+  | x :: l, Used_as_invariant :: usage ->
+    let invariant, normal = partition_used l usage in
+    x :: invariant, normal
 
-let used_extra_params t = t.used_extra_params
+let get_used_params rewrite =
+  let invariant_params, variant_params =
+    partition_used (Bound_parameters.to_list rewrite.original_params)
+      rewrite.original_params_usage
+  in
+  let invariant_extra_params, variant_extra_params =
+    partition_used (Bound_parameters.to_list rewrite.extra_params)
+      rewrite.extra_params_usage
+  in
+  Bound_parameters.create (invariant_params @ invariant_extra_params),
+  Bound_parameters.create (variant_params @ variant_extra_params)
 
-let used_extra_invariant_params t = t.used_extra_invariant_params
+type rewrite_apply_cont_ctx =
+  | Apply_cont
+  | Apply_expr of Simple.t list
 
-let invariant_params t = t.invariant_params
+let extra_args_list rewrite id =
+  try Id.Map.find id rewrite.extra_args
+  with Not_found ->
+  match rewrite.extra_params_usage with
+  | [] -> []
+  | _ :: _ ->
+    Misc.fatal_errorf
+      "Apply_cont_rewrite.extra_args_list:@ \
+       Could not find extra args but extra params were not empty"
 
-let extra_args t id =
-  match Id.Map.find id t.extra_args with
-  | exception Not_found ->
-    if not (Bound_parameters.is_empty t.used_extra_params)
-    then
-      Misc.fatal_errorf
-        "This [Apply_cont_rewrite] does not have any@ extra arguments for use \
-         ID %a, but it has@ >= 1 extra parameter:@ %a"
-        Id.print id print t;
-    []
-  | extra_args -> extra_args
+let make_rewrite rewrite ~ctx id args =
+  let invariant_args, args = partition_used args rewrite.original_params_usage in
+  let extra_args_list = extra_args_list rewrite id in
+  let extra_invariant_args_rev, extra_args_rev, extra_lets, _ =
+    List.fold_left2
+      (fun (extra_invariant_args_rev, extra_args_rev, extra_lets, required_by_other_extra_args)
+           (arg : EA.t) used ->
+        (* Some extra_args computation can depend on other extra args. But those
+           required extra args might not be needed as argument to the
+           continuation. But we want to keep the let bindings.
+           [required_by_other_extra_args] tracks that dependency. It is the set
+           of free variables of [extra_args_rev] *)
+        let extra_arg, extra_let, free_names, defined_names =
+          match arg with
+          | Already_in_scope simple ->
+            simple, [], Simple.free_names simple, Name_occurrences.empty
+          | New_let_binding (temp, prim) ->
+            let extra_let =
+              ( Bound_var.create temp Name_mode.normal,
+                Code_size.prim prim,
+                Flambda.Named.create_prim prim Debuginfo.none )
+            in
+            Simple.var temp, [extra_let], Flambda_primitive.free_names prim,
+            Name_occurrences.singleton_variable temp Name_mode.normal
+          | New_let_binding_with_named_args (temp, gen_prim) ->
+            let prim =
+              match (ctx : rewrite_apply_cont_ctx) with
+              | Apply_expr function_return_values ->
+                gen_prim function_return_values
+              | Apply_cont ->
+                Misc.fatal_errorf
+                  "Apply_cont rewrites should not need to name arguments, \
+                   since they are already named."
+            in
+            let extra_let =
+              ( Bound_var.create temp Name_mode.normal,
+                Code_size.prim prim,
+                Flambda.Named.create_prim prim Debuginfo.none )
+            in
+            Simple.var temp, [extra_let], Flambda_primitive.free_names prim,
+            Name_occurrences.singleton_variable temp Name_mode.normal
+        in
+        let required_let, extra_invariant_args_rev, extra_args_rev =
+          match used with
+          | Used ->
+            true, extra_invariant_args_rev, extra_arg :: extra_args_rev
+          | Used_as_invariant ->
+            true, extra_arg :: extra_invariant_args_rev, extra_args_rev
+          | Unused ->
+            Name_occurrences.inter_domain_is_non_empty defined_names
+              required_by_other_extra_args,
+            extra_invariant_args_rev, extra_args_rev
+        in
+        if required_let
+        then
+          ( extra_invariant_args_rev, extra_args_rev,
+            extra_let @ extra_lets,
+            Name_occurrences.union free_names required_by_other_extra_args )
+        else extra_invariant_args_rev, extra_args_rev, extra_lets, required_by_other_extra_args)
+      ([], [], [], Name_occurrences.empty)
+      extra_args_list rewrite.extra_params_usage
+  in
+  extra_lets,
+  invariant_args @ List.rev_append extra_invariant_args_rev args @ List.rev extra_args_rev
 
-let original_params_arity t =
-  Bound_parameters.arity_with_subkinds (original_params t)
+let rewrite_exn_continuation rewrite id exn_cont =
+  let exn_cont_arity = Exn_continuation.arity exn_cont in
+  if not
+       (Flambda_arity.With_subkinds.equal
+          exn_cont_arity
+          (original_params_arity rewrite))
+  then
+    Misc.fatal_errorf
+      "Arity of exception continuation %a does not match@ [original_params] \
+       (%a)"
+      Exn_continuation.print exn_cont Bound_parameters.print rewrite.original_params;
+  assert (Flambda_arity.With_subkinds.cardinal exn_cont_arity >= 1);
+  if List.hd rewrite.original_params_usage <> Used then
+    Misc.fatal_errorf "The usage of the exn parameter of the continuation rewrite should be [Used]: %a"
+      print rewrite;
+  if List.exists (fun x -> x = Used_as_invariant)
+      (rewrite.original_params_usage @ rewrite.extra_params_usage) then
+    Misc.fatal_errorf "An exception continuation should never have invariant parameters: %a"
+      print rewrite;
+  let _, extra_args0 = partition_used (Exn_continuation.extra_args exn_cont)
+      (List.tl rewrite.original_params_usage)
+  in
+  let _, extra_args1 =
+    let extra_args_list =
+      List.map2 (fun (arg : EA.t) extra_param ->
+          match arg with
+          | Already_in_scope simple ->
+            simple, Bound_parameter.kind extra_param
+          | New_let_binding _ | New_let_binding_with_named_args _ ->
+            (* Note: this is unsupported for now. If we choose to support it in
+               the future, we must take care of not introducing a wrapper continuation,
+               which would come with its own pushtrap/poptrap. *)
+            Misc.fatal_error
+              "[New_let_binding] are currently forbidden for exn continuation \
+               rewrites"
+        ) (extra_args_list rewrite id)
+        (Bound_parameters.to_list rewrite.extra_params)
+    in
+    partition_used extra_args_list rewrite.extra_params_usage
+  in
+  let extra_args = extra_args0 @ extra_args1 in
+  Exn_continuation.create
+    ~exn_handler:(Exn_continuation.exn_handler exn_cont)
+    ~extra_args

middle_end/flambda2/simplify/apply_cont_rewrite.mli

@@ -21,41 +21,40 @@
 
 type t
 
-type used = private
+type used =
+  | Unused
   | Used
   | Used_as_invariant
-  | Unused
 
 val print : Format.formatter -> t -> unit
 
 (** [extra_args] (and hence [extra_params]) must be given in order: later
     extra-args may refer to earlier extra-args, but not vice-versa. *)
 val create :
   original_params:Bound_parameters.t ->
-  used_params:Bound_parameter.Set.t ->
-  invariant_params:Bound_parameter.Set.t ->
-  extra_params:Bound_parameters.t ->
-  extra_args:
-    Continuation_extra_params_and_args.Extra_arg.t list
-    Apply_cont_rewrite_id.Map.t ->
-  used_extra_params:Bound_parameter.Set.t ->
+  extra_params_and_args:Continuation_extra_params_and_args.t ->
+  decide_param_usage:(Bound_parameter.t -> used) ->
   t
 
-val original_params : t -> Bound_parameters.t
-
-val used_params : t -> Bound_parameter.Set.t
+val does_nothing : t -> bool
 
-val used_extra_params : t -> Bound_parameters.t
+val get_used_params : t -> Bound_parameters.t * Bound_parameters.t
 
-val used_extra_invariant_params : t -> Bound_parameters.t
+val original_params_arity : t -> Flambda_arity.With_subkinds.t
 
-val invariant_params : t -> Bound_parameter.Set.t
+type rewrite_apply_cont_ctx =
+  | Apply_cont
+  | Apply_expr of Simple.t list
 
-val extra_args :
+val make_rewrite :
   t ->
+  ctx:rewrite_apply_cont_ctx ->
   Apply_cont_rewrite_id.t ->
-  (Continuation_extra_params_and_args.Extra_arg.t * used) list
+  Simple.t list ->
+  (Bound_var.t * Code_size.t * Flambda.Named.t) list * Simple.t list
 
-val original_params_arity : t -> Flambda_arity.With_subkinds.t
-
-val does_nothing : t -> bool
+val rewrite_exn_continuation :
+  t ->
+  Apply_cont_rewrite_id.t ->
+  Exn_continuation.t ->
+  Exn_continuation.t