Fix missing End_region primitives on switch arms (#898)

Author: lthls

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

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 =

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