Add parent regions as parameters to Begin_region

lthls · Jan 22, 2025 · 080c8ad · 080c8ad
1 parent b1ffd3d
commit 080c8ad
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 9 deletions.
diff --git a/middle_end/flambda2/from_lambda/closure_conversion.ml b/middle_end/flambda2/from_lambda/closure_conversion.ml
@@ -1096,13 +1096,13 @@ let close_named acc env ~let_bound_ids_with_kinds (named : IR.named)
     in
     Lambda_to_flambda_primitives_helpers.bind_recs acc None ~register_const0
       prim Debuginfo.none k
-  | Begin_region { is_try_region; ghost } ->
+  | Begin_region { is_try_region; ghost; parent_region } ->
     let prim : Lambda_to_flambda_primitives_helpers.expr_primitive =
       Variadic
         ( (if is_try_region
           then Begin_try_region { ghost }
           else Begin_region { ghost }),
-          [] )
+          [Simple (find_simple_from_id env parent_region)] )
     in
     Lambda_to_flambda_primitives_helpers.bind_recs acc None ~register_const0
       prim Debuginfo.none k

diff --git a/middle_end/flambda2/from_lambda/closure_conversion_aux.ml b/middle_end/flambda2/from_lambda/closure_conversion_aux.ml
@@ -37,7 +37,8 @@ module IR = struct
     | Get_tag of Ident.t
     | Begin_region of
         { ghost : bool;
-          is_try_region : bool
+          is_try_region : bool;
+          parent_region : Ident.t
         }
     | End_region of
         { is_try_region : bool;
@@ -98,7 +99,7 @@ module IR = struct
     | Simple (Var id) -> Ident.print ppf id
     | Simple (Const cst) -> Printlambda.structured_constant ppf cst
     | Get_tag id -> fprintf ppf "@[<2>(Gettag %a)@]" Ident.print id
-    | Begin_region { is_try_region; ghost } ->
+    | Begin_region { is_try_region; ghost; parent_region = _ } ->
       if is_try_region
       then fprintf ppf "Begin_try_region"
       else fprintf ppf "Begin_region";

diff --git a/middle_end/flambda2/from_lambda/closure_conversion_aux.mli b/middle_end/flambda2/from_lambda/closure_conversion_aux.mli
@@ -39,7 +39,8 @@ module IR : sig
     | Get_tag of Ident.t (* Intermediary primitive for block switch *)
     | Begin_region of
         { ghost : bool;
-          is_try_region : bool
+          is_try_region : bool;
+          parent_region : Ident.t
         }
     | End_region of
         { is_try_region : bool;

diff --git a/middle_end/flambda2/from_lambda/lambda_to_flambda.ml b/middle_end/flambda2/from_lambda/lambda_to_flambda.ml
@@ -845,16 +845,26 @@ let rec cps acc env ccenv (lam : L.lambda) (k : cps_continuation)
                { is_try_region = true; region = ghost_region; ghost = true })
             ~body:(fun acc ccenv -> cps_tail acc env ccenv handler k k_exn))
     in
+    let region_stack_elt = Env.current_region env in
     let begin_try_region body =
       CC.close_let acc ccenv
         [region, Flambda_kind.With_subkind.region]
         Not_user_visible
-        (Begin_region { is_try_region = true; ghost = false })
+        (Begin_region
+           { is_try_region = true;
+             ghost = false;
+             parent_region = Env.Region_stack_element.region region_stack_elt
+           })
         ~body:(fun acc ccenv ->
           CC.close_let acc ccenv
             [ghost_region, Flambda_kind.With_subkind.region]
             Not_user_visible
-            (Begin_region { is_try_region = true; ghost = true })
+            (Begin_region
+               { is_try_region = true;
+                 ghost = true;
+                 parent_region =
+                   Env.Region_stack_element.ghost_region region_stack_elt
+               })
             ~body)
     in
     begin_try_region (fun acc ccenv ->
@@ -964,19 +974,29 @@ let rec cps acc env ccenv (lam : L.lambda) (k : cps_continuation)
        continuation for the code after the body. *)
     let region = Ident.create_local "region" in
     let ghost_region = Ident.create_local "ghost_region" in
+    let parent_stack_elt = Env.current_region env in
     let region_stack_elt =
       Env.Region_stack_element.create ~region ~ghost_region
     in
     let dbg = Debuginfo.none in
     CC.close_let acc ccenv
       [region, Flambda_kind.With_subkind.region]
       Not_user_visible
-      (Begin_region { is_try_region = false; ghost = false })
+      (Begin_region
+         { is_try_region = false;
+           ghost = false;
+           parent_region = Env.Region_stack_element.region parent_stack_elt
+         })
       ~body:(fun acc ccenv ->
         CC.close_let acc ccenv
           [ghost_region, Flambda_kind.With_subkind.region]
           Not_user_visible
-          (Begin_region { is_try_region = false; ghost = true })
+          (Begin_region
+             { is_try_region = false;
+               ghost = true;
+               parent_region =
+                 Env.Region_stack_element.ghost_region parent_stack_elt
+             })
           ~body:(fun acc ccenv ->
             maybe_insert_let_cont "body_return" layout k acc env ccenv
               (fun acc env ccenv k ->