From b2b19c32941e14ab01fec98b5821331a86d1b1e9 Mon Sep 17 00:00:00 2001
From: Vincent Laviron <vincent.laviron@gmail.com>
Date: Wed, 19 Oct 2022 14:31:29 +0200
Subject: [PATCH 1/2] Fix missing End_region primitives on switch arms

---
 .../flambda2/from_lambda/lambda_to_flambda.ml | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/middle_end/flambda2/from_lambda/lambda_to_flambda.ml b/middle_end/flambda2/from_lambda/lambda_to_flambda.ml
index 11c8986d95b..3279b3b8f08 100644
--- a/middle_end/flambda2/from_lambda/lambda_to_flambda.ml
+++ b/middle_end/flambda2/from_lambda/lambda_to_flambda.ml
@@ -805,6 +805,23 @@ let restore_continuation_context acc env ccenv cont ~close_early body =
           Continuation.print cont;
       body acc ccenv continuation_closing_region
 
+let restore_continuation_context_for_switch_arm env cont =
+  match Env.pop_regions_up_to_context env cont with
+  | None -> cont
+  | Some region ->
+    let ({ continuation_closing_region; continuation_after_closing_region }
+          : Env.region_closure_continuation) =
+      Env.region_closure_continuation env region
+    in
+    if not (Continuation.equal cont continuation_after_closing_region)
+    then
+      Misc.fatal_errorf
+        "The continuation %a following the region closure should be the \
+         current continuation %a"
+        Continuation.print continuation_after_closing_region Continuation.print
+        cont;
+    continuation_closing_region
+
 let apply_cont_with_extra_args acc env ccenv ~dbg cont traps args =
   let extra_args =
     List.map
@@ -1576,6 +1593,7 @@ and cps_switch acc env ccenv (switch : L.lambda_switch) ~condition_dbg
               (fun arg : IR.simple -> Var arg)
               (Env.extra_args_for_continuation env k)
           in
+          let k = restore_continuation_context_for_switch_arm env k in
           let consts_rev =
             (arm, k, None, IR.Var var :: extra_args) :: consts_rev
           in
@@ -1586,6 +1604,7 @@ and cps_switch acc env ccenv (switch : L.lambda_switch) ~condition_dbg
               (fun arg : IR.simple -> Var arg)
               (Env.extra_args_for_continuation env k)
           in
+          let k = restore_continuation_context_for_switch_arm env k in
           let consts_rev =
             (arm, k, None, IR.Const cst :: extra_args) :: consts_rev
           in

From bee93edbde3441af2a68be2251b3d51792636fda Mon Sep 17 00:00:00 2001
From: Vincent Laviron <vincent.laviron@gmail.com>
Date: Wed, 19 Oct 2022 14:49:13 +0200
Subject: [PATCH 2/2] Add test

---
 ocaml/testsuite/tests/typing-local/regions.ml        | 7 +++++++
 ocaml/testsuite/tests/typing-local/regions.reference | 1 +
 2 files changed, 8 insertions(+)

diff --git a/ocaml/testsuite/tests/typing-local/regions.ml b/ocaml/testsuite/tests/typing-local/regions.ml
index 44c73f23b68..57ceccb331c 100644
--- a/ocaml/testsuite/tests/typing-local/regions.ml
+++ b/ocaml/testsuite/tests/typing-local/regions.ml
@@ -114,6 +114,13 @@ let () = check_empty "toplevel rec binding"
  ());;
 let () = check_empty "toplevel eval"
 
+let () =
+  let f x b g =
+    let local_ p = x, x in
+    if b then () else (g p; ())
+  in f 0 true (fun _ -> ())
+let () = check_empty "constant switch arm"
+
 module type T = sig val x : int end
 let _ =
   let module M : T =
diff --git a/ocaml/testsuite/tests/typing-local/regions.reference b/ocaml/testsuite/tests/typing-local/regions.reference
index bea73d31191..d5ea082ec10 100644
--- a/ocaml/testsuite/tests/typing-local/regions.reference
+++ b/ocaml/testsuite/tests/typing-local/regions.reference
@@ -13,6 +13,7 @@
          toplevel binding: OK
      toplevel rec binding: OK
             toplevel eval: OK
+      constant switch arm: OK
           first class mod: OK
        class d definition: OK
         class definitions: OK