Fix computation of extra_args during CPS conversion of Lstaticraise (o…

…caml-flambda#126)

middle_end/flambda2/from_lambda/lambda_to_flambda.ml

@@ -894,12 +894,12 @@ let rec cps_non_tail acc env ccenv (lam : L.lambda)
       k k_exn
   | Lstaticraise (static_exn, args) ->
     let continuation = Env.get_static_exn_continuation env static_exn in
-    let extra_args =
-      List.map (fun var : IR.simple -> Var var)
-        (Env.extra_args_for_continuation env continuation)
-    in
     cps_non_tail_list acc env ccenv args
       (fun acc env ccenv args ->
+         let extra_args =
+           List.map (fun var : IR.simple -> Var var)
+             (Env.extra_args_for_continuation env continuation)
+         in
          compile_staticfail acc env ccenv
            ~continuation ~args:(args @ extra_args)
       ) k_exn
@@ -1251,12 +1251,12 @@ and cps_tail acc env ccenv (lam : L.lambda) (k : Continuation.t)
       k k_exn
   | Lstaticraise (static_exn, args) ->
     let continuation = Env.get_static_exn_continuation env static_exn in
-    let extra_args =
-      List.map (fun var : IR.simple -> Var var)
-        (Env.extra_args_for_continuation env continuation)
-    in
     cps_non_tail_list acc env ccenv args
       (fun acc env ccenv args ->
+         let extra_args =
+           List.map (fun var : IR.simple -> Var var)
+             (Env.extra_args_for_continuation env continuation)
+         in
          compile_staticfail acc env ccenv
            ~continuation ~args:(args @ extra_args)
       ) k_exn

middle_end/flambda2/tests/mlexamples/extra_args_and_staticraise.ml

@@ -0,0 +1,30 @@
+(* After [Base.Hashtbl.fold].
+   This example produces the following lambda code:
+
+    (let
+      (f/9 =
+         (function x/11[int] init/12[int] : int
+           (let (m/13 =v[int] init/12)
+             (catch
+               (try (exit 1 (for i/14 0 to x/11 (assign m/13 (+ m/13 i/14))))
+                with exn/23 0)
+              with (1 val/22[int]) m/13))))
+      (makeblock 0 f/9))
+
+   which is interesting because of the for-loop as an argument to
+   [exit], in conjunction with the mutable variable [m]. *)
+
+type 'a ref = { mutable contents : 'a }
+external ref : 'a -> 'a ref = "%makemutable"
+external ( ! ) : 'a ref -> 'a = "%field0"
+external ( := ) : 'a ref -> 'a -> unit = "%setfield0"
+external ( + ) : int -> int -> int = "%addint"
+let f x init =
+  let m = ref init in
+  match
+    for i = 0 to x do
+      m := !m + i
+    done
+  with
+  | () -> !m
+  | exception _ -> 0