Some cleanup in bytecomp/bytegen.ml (#11613)

2 files changed, 170 insertions, 148 deletions

diff --git a/bytecomp/bytegen.ml b/bytecomp/bytegen.ml
index bccd38e717..6ea0b0766f 100644
--- a/bytecomp/bytegen.ml
+++ b/bytecomp/bytegen.ml
@@ -322,28 +322,43 @@ let add_pseudo_event loc modname c =
 
 (**** Compilation of a lambda expression ****)
 
-let try_blocks = ref []  (* list of stack size for each nested try block *)
+type stack_info = {
+  try_blocks : int list;
+  (* list of stack size for each nested try block *)
+  sz_static_raises : (int * (int * int * int list)) list;
+  (* association staticraise numbers -> (lbl,size of stack, try_blocks *)
+  max_stack_used : int ref;
+  (* Maximal stack size reached during the current function body *)
+}
+
+let create_stack_info () = {
+  try_blocks = [];
+  sz_static_raises = [];
+  max_stack_used = ref 0
+}
 
 (* association staticraise numbers -> (lbl,size of stack, try_blocks *)
 
-let sz_static_raises = ref []
+let push_static_raise stack_info i lbl_handler sz =
+  { stack_info
+    with
+      sz_static_raises = (i, (lbl_handler, sz, stack_info.try_blocks))
+                         :: stack_info.sz_static_raises
+  }
 
-let push_static_raise i lbl_handler sz =
-  sz_static_raises := (i, (lbl_handler, sz, !try_blocks)) :: !sz_static_raises
-
-let find_raise_label i =
+let find_raise_label stack_info i =
   try
-    List.assoc i !sz_static_raises
+    List.assoc i stack_info.sz_static_raises
   with
   | Not_found ->
       Misc.fatal_error
         ("exit("^Int.to_string i^") outside appropriated catch")
 
 (* Will the translation of l lead to a jump to label ? *)
-let code_as_jump l sz = match l with
+let code_as_jump stack_info l sz = match l with
 | Lstaticraise (i,[]) ->
-    let label,size,tb = find_raise_label i in
-    if sz = size && tb == !try_blocks then
+    let label,size,tb = find_raise_label stack_info i in
+    if sz = size && tb == stack_info.try_blocks then
       Some label
     else
       None
@@ -366,12 +381,9 @@ let functions_to_compile  = (Stack.create () : function_to_compile Stack.t)
 
 let compunit_name = ref ""
 
-(* Maximal stack size reached during the current function body *)
-
-let max_stack_used = ref 0
-
-let check_stack sz =
-  if sz > !max_stack_used then max_stack_used := sz
+let check_stack stack_info sz =
+  let curr = stack_info.max_stack_used in
+  if sz > !curr then curr := sz
 
 (* Sequence of string tests *)
 
@@ -386,8 +398,8 @@ let comp_bint_primitive bi suff args =
                 | Pint64 -> "caml_int64_" in
   Kccall(pref ^ suff, List.length args)
 
-let comp_primitive p sz args =
-  check_stack sz;
+let comp_primitive stack_info p sz args =
+  check_stack stack_info sz;
   match p with
     Pgetglobal id -> Kgetglobal id
   | Psetglobal id -> Ksetglobal id
@@ -403,7 +415,7 @@ let comp_primitive p sz args =
   | Pduprecord _ -> Kccall("caml_obj_dup", 1)
   | Pccall p -> Kccall(p.prim_name, p.prim_arity)
   | Pperform ->
-      check_stack (sz + 4);
+      check_stack stack_info (sz + 4);
       Kperform
   | Pnegint -> Knegint
   | Paddint -> Kaddint
@@ -550,8 +562,8 @@ module Storer =
    cont = list of instructions to execute afterwards
    Result = list of instructions that evaluate exp, then perform cont. *)
 
-let rec comp_expr env exp sz cont =
-  check_stack sz;
+let rec comp_expr stack_info env exp sz cont =
+  check_stack stack_info sz;
   match exp with
     Lvar id | Lmutvar id ->
       begin try
@@ -573,18 +585,19 @@ let rec comp_expr env exp sz cont =
   | Lapply{ap_func = func; ap_args = args} ->
       let nargs = List.length args in
       if is_tailcall cont then begin
-        comp_args env args sz
-          (Kpush :: comp_expr env func (sz + nargs)
+        comp_args stack_info env args sz
+          (Kpush :: comp_expr stack_info env func (sz + nargs)
             (Kappterm(nargs, sz + nargs) :: discard_dead_code cont))
       end else begin
         if nargs < 4 then
-          comp_args env args sz
-            (Kpush :: comp_expr env func (sz + nargs) (Kapply nargs :: cont))
+          comp_args stack_info env args sz
+            (Kpush ::
+             comp_expr stack_info env func (sz + nargs) (Kapply nargs :: cont))
         else begin
           let (lbl, cont1) = label_code cont in
           Kpush_retaddr lbl ::
-          comp_args env args (sz + 3)
-            (Kpush :: comp_expr env func (sz + 3 + nargs)
+          comp_args stack_info env args (sz + 3)
+            (Kpush :: comp_expr stack_info env func (sz + 3 + nargs)
                       (Kapply nargs :: cont1))
         end
       end
@@ -598,16 +611,16 @@ let rec comp_expr env exp sz cont =
         | _ -> (Kgetdynmet, met::obj::args)
       in
       if is_tailcall cont then
-        comp_args env args' sz
+        comp_args stack_info env args' sz
           (getmethod :: Kappterm(nargs, sz + nargs) :: discard_dead_code cont)
       else
         if nargs < 4 then
-          comp_args env args' sz
+          comp_args stack_info env args' sz
             (getmethod :: Kapply nargs :: cont)
         else begin
           let (lbl, cont1) = label_code cont in
           Kpush_retaddr lbl ::
-          comp_args env args' (sz + 3)
+          comp_args stack_info env args' (sz + 3)
             (getmethod :: Kapply nargs :: cont1)
         end
   | Lfunction{params; body; loc} -> (* assume kind = Curried *)
@@ -618,12 +631,12 @@ let rec comp_expr env exp sz cont =
         { params = List.map fst params; body = body; label = lbl;
           free_vars = fv; num_defs = 1; rec_vars = []; rec_pos = 0 } in
       Stack.push to_compile functions_to_compile;
-      comp_args env (List.map (fun n -> Lvar n) fv) sz
+      comp_args stack_info env (List.map (fun n -> Lvar n) fv) sz
         (Kclosure(lbl, List.length fv) :: cont)
   | Llet(_, _k, id, arg, body)
   | Lmutlet(_k, id, arg, body) ->
-      comp_expr env arg sz
-        (Kpush :: comp_expr (add_var id (sz+1) env) body (sz+1)
+      comp_expr stack_info env arg sz
+        (Kpush :: comp_expr stack_info (add_var id (sz+1) env) body (sz+1)
           (add_pop 1 cont))
   | Lletrec(decl, body) ->
       let ndecl = List.length decl in
@@ -645,10 +658,11 @@ let rec comp_expr env exp sz cont =
               lbl :: comp_fun (pos + 1) rem
           | _ -> assert false in
         let lbls = comp_fun 0 decl in
-        comp_args env (List.map (fun n -> Lvar n) fv) sz
+        comp_args stack_info env (List.map (fun n -> Lvar n) fv) sz
           (Kclosurerec(lbls, List.length fv) ::
-            (comp_expr (add_vars rec_idents (sz+1) env) body (sz + ndecl)
-                       (add_pop ndecl cont)))
+           (comp_expr stack_info
+              (add_vars rec_idents (sz+1) env) body (sz + ndecl)
+              (add_pop ndecl cont)))
       end else begin
         let decl_size =
           List.map (fun (id, exp) -> (id, exp, size_of_lambda Ident.empty exp))
@@ -685,14 +699,14 @@ let rec comp_expr env exp sz cont =
             :: rem ->
               comp_nonrec new_env sz (i-1) rem
           | (_id, exp, RHS_nonrec) :: rem ->
-              comp_expr new_env exp sz
+              comp_expr stack_info new_env exp sz
                 (Kassign (i-1) :: comp_nonrec new_env sz (i-1) rem)
         and comp_rec new_env sz i = function
-          | [] -> comp_expr new_env body sz (add_pop ndecl cont)
+          | [] -> comp_expr stack_info new_env body sz (add_pop ndecl cont)
           | (_id, exp, (RHS_block _ | RHS_infix _ |
                         RHS_floatblock _ | RHS_function _))
             :: rem ->
-              comp_expr new_env exp sz
+              comp_expr stack_info new_env exp sz
                 (Kpush :: Kacc i :: Kccall("caml_update_dummy", 2) ::
                  comp_rec new_env sz (i-1) rem)
           | (_id, _exp, RHS_nonrec) :: rem ->
@@ -701,55 +715,55 @@ let rec comp_expr env exp sz cont =
         comp_init env sz decl_size
       end
   | Lprim(Popaque, [arg], _) ->
-      comp_expr env arg sz cont
+      comp_expr stack_info env arg sz cont
   | Lprim(Pignore, [arg], _) ->
-      comp_expr env arg sz (add_const_unit cont)
+      comp_expr stack_info env arg sz (add_const_unit cont)
   | Lprim(Pnot, [arg], _) ->
       let newcont =
         match cont with
           Kbranchif lbl :: cont1 -> Kbranchifnot lbl :: cont1
         | Kbranchifnot lbl :: cont1 -> Kbranchif lbl :: cont1
         | _ -> Kboolnot :: cont in
-      comp_expr env arg sz newcont
+      comp_expr stack_info env arg sz newcont
   | Lprim(Psequand, [exp1; exp2], _) ->
       begin match cont with
         Kbranchifnot lbl :: _ ->
-          comp_expr env exp1 sz (Kbranchifnot lbl ::
-            comp_expr env exp2 sz cont)
+          comp_expr stack_info env exp1 sz (Kbranchifnot lbl ::
+            comp_expr stack_info env exp2 sz cont)
       | Kbranchif lbl :: cont1 ->
           let (lbl2, cont2) = label_code cont1 in
-          comp_expr env exp1 sz (Kbranchifnot lbl2 ::
-            comp_expr env exp2 sz (Kbranchif lbl :: cont2))
+          comp_expr stack_info env exp1 sz (Kbranchifnot lbl2 ::
+            comp_expr stack_info env exp2 sz (Kbranchif lbl :: cont2))
       | _ ->
           let (lbl, cont1) = label_code cont in
-          comp_expr env exp1 sz (Kstrictbranchifnot lbl ::
-            comp_expr env exp2 sz cont1)
+          comp_expr stack_info env exp1 sz (Kstrictbranchifnot lbl ::
+            comp_expr stack_info env exp2 sz cont1)
       end
   | Lprim(Psequor, [exp1; exp2], _) ->
       begin match cont with
         Kbranchif lbl :: _ ->
-          comp_expr env exp1 sz (Kbranchif lbl ::
-            comp_expr env exp2 sz cont)
+          comp_expr stack_info env exp1 sz (Kbranchif lbl ::
+            comp_expr stack_info env exp2 sz cont)
       | Kbranchifnot lbl :: cont1 ->
           let (lbl2, cont2) = label_code cont1 in
-          comp_expr env exp1 sz (Kbranchif lbl2 ::
-            comp_expr env exp2 sz (Kbranchifnot lbl :: cont2))
+          comp_expr stack_info env exp1 sz (Kbranchif lbl2 ::
+            comp_expr stack_info env exp2 sz (Kbranchifnot lbl :: cont2))
       | _ ->
           let (lbl, cont1) = label_code cont in
-          comp_expr env exp1 sz (Kstrictbranchif lbl ::
-            comp_expr env exp2 sz cont1)
+          comp_expr stack_info env exp1 sz (Kstrictbranchif lbl ::
+            comp_expr stack_info env exp2 sz cont1)
       end
   | Lprim(Praise k, [arg], _) ->
-      comp_expr env arg sz (Kraise k :: discard_dead_code cont)
+      comp_expr stack_info env arg sz (Kraise k :: discard_dead_code cont)
   | Lprim(Paddint, [arg; Lconst(Const_base(Const_int n))], _)
     when is_immed n ->
-      comp_expr env arg sz (Koffsetint n :: cont)
+      comp_expr stack_info env arg sz (Koffsetint n :: cont)
   | Lprim(Psubint, [arg; Lconst(Const_base(Const_int n))], _)
     when is_immed (-n) ->
-      comp_expr env arg sz (Koffsetint (-n) :: cont)
+      comp_expr stack_info env arg sz (Koffsetint (-n) :: cont)
   | Lprim (Poffsetint n, [arg], _)
     when not (is_immed n) ->
-      comp_expr env arg sz
+      comp_expr stack_info env arg sz
         (Kpush::
          Kconst (Const_base (Const_int n))::
          Kaddint::cont)
@@ -757,13 +771,15 @@ let rec comp_expr env exp sz cont =
       let cont = add_pseudo_event loc !compunit_name cont in
       begin match kind with
         Pintarray | Paddrarray ->
-          comp_args env args sz (Kmakeblock(List.length args, 0) :: cont)
+          comp_args stack_info env args sz
+            (Kmakeblock(List.length args, 0) :: cont)
       | Pfloatarray ->
-          comp_args env args sz (Kmakefloatblock(List.length args) :: cont)
+          comp_args stack_info env args sz
+            (Kmakefloatblock(List.length args) :: cont)
       | Pgenarray ->
           if args = []
           then Kmakeblock(0, 0) :: cont
-          else comp_args env args sz
+          else comp_args stack_info env args sz
                  (Kmakeblock(List.length args, 0) ::
                   Kccall("caml_make_array", 1) :: cont)
       end
@@ -771,30 +787,31 @@ let rec comp_expr env exp sz cont =
       let nargs = List.length args - 1 in
       assert (nargs = 2);
       (* Resume itself only pushes 3 words, but perform adds another *)
-      check_stack (sz + 4);
+      check_stack stack_info (sz + 4);
       if is_tailcall cont then
-        comp_args env args sz
+        comp_args stack_info env args sz
           (Kresumeterm(sz + nargs) :: discard_dead_code cont)
       else
-        comp_args env args sz (Kresume :: cont)
+        comp_args stack_info env args sz (Kresume :: cont)
   | Lprim(Preperform, args, _) ->
       let nargs = List.length args - 1 in
       assert (nargs = 2);
-      check_stack (sz + 3);
+      check_stack stack_info (sz + 3);
       if is_tailcall cont then
-        comp_args env args sz
+        comp_args stack_info env args sz
           (Kreperformterm(sz + nargs) :: discard_dead_code cont)
       else
         fatal_error "Reperform used in non-tail position"
   | Lprim (Pduparray (kind, mutability),
            [Lprim (Pmakearray (kind',_),args,_)], loc) ->
       assert (kind = kind');
-      comp_expr env (Lprim (Pmakearray (kind, mutability), args, loc)) sz cont
+      comp_expr stack_info env
+        (Lprim (Pmakearray (kind, mutability), args, loc)) sz cont
   | Lprim (Pduparray _, [arg], loc) ->
       let prim_obj_dup =
         Primitive.simple ~name:"caml_obj_dup" ~arity:1 ~alloc:true
       in
-      comp_expr env (Lprim (Pccall prim_obj_dup, [arg], loc)) sz cont
+      comp_expr stack_info env (Lprim (Pccall prim_obj_dup, [arg], loc)) sz cont
   | Lprim (Pduparray _, _, _) ->
       Misc.fatal_error "Bytegen.comp_expr: Pduparray takes exactly one arg"
 (* Integer first for enabling further optimization (cf. emitcode.ml)  *)
@@ -802,7 +819,8 @@ let rec comp_expr env exp sz cont =
       let p = Pintcomp (swap_integer_comparison c)
       and args = [k ; arg] in
       let nargs = List.length args - 1 in
-      comp_args env args sz (comp_primitive p (sz + nargs - 1) args :: cont)
+      comp_args stack_info env args sz
+        (comp_primitive stack_info p (sz + nargs - 1) args :: cont)
   | Lprim (Pfloatcomp cmp, args, _) ->
       let cont =
         match cmp with
@@ -817,16 +835,18 @@ let rec comp_expr env exp sz cont =
         | CFge -> Kccall("caml_ge_float", 2) :: cont
         | CFnge -> Kccall("caml_ge_float", 2) :: Kboolnot :: cont
       in
-      comp_args env args sz cont
+      comp_args stack_info env args sz cont
   | Lprim(Pmakeblock(tag, _mut, _), args, loc) ->
       let cont = add_pseudo_event loc !compunit_name cont in
-      comp_args env args sz (Kmakeblock(List.length args, tag) :: cont)
+      comp_args stack_info env args sz
+        (Kmakeblock(List.length args, tag) :: cont)
   | Lprim(Pfloatfield n, args, loc) ->
       let cont = add_pseudo_event loc !compunit_name cont in
-      comp_args env args sz (Kgetfloatfield n :: cont)
+      comp_args stack_info env args sz (Kgetfloatfield n :: cont)
   | Lprim(p, args, _) ->
       let nargs = List.length args - 1 in
-      comp_args env args sz (comp_primitive p (sz + nargs - 1) args :: cont)
+      comp_args stack_info env args sz
+        (comp_primitive stack_info p (sz + nargs - 1) args :: cont)
   | Lstaticcatch (body, (i, vars) , handler) ->
       let vars = List.map fst vars in
       let nvars = List.length vars in
@@ -836,27 +856,29 @@ let rec comp_expr env exp sz cont =
           let lbl_handler, cont2 =
             label_code
               (comp_expr
+                stack_info
                 (add_vars vars (sz+1) env)
                 handler (sz+nvars) (add_pop nvars cont1)) in
-          push_static_raise i lbl_handler (sz+nvars);
+          let stack_info =
+            push_static_raise stack_info i lbl_handler (sz+nvars) in
           push_dummies nvars
-            (comp_expr env body (sz+nvars)
+            (comp_expr stack_info env body (sz+nvars)
             (add_pop nvars (branch1 :: cont2)))
         end else begin (* small optimization for nvars = 1 *)
           let var = match vars with [var] -> var | _ -> assert false in
           let lbl_handler, cont2 =
             label_code
-              (Kpush::comp_expr
+              (Kpush::comp_expr stack_info
                 (add_var var (sz+1) env)
                 handler (sz+1) (add_pop 1 cont1)) in
-          push_static_raise i lbl_handler sz;
-          comp_expr env body sz (branch1 :: cont2)
+          let stack_info =
+            push_static_raise stack_info i lbl_handler sz in
+          comp_expr stack_info env body sz (branch1 :: cont2)
         end in
-      sz_static_raises := List.tl !sz_static_raises ;
       r
   | Lstaticraise (i, args) ->
       let cont = discard_dead_code cont in
-      let label,size,tb = find_raise_label i in
+      let label,size,tb = find_raise_label stack_info i in
       let cont = branch_to label cont in
       let rec loop sz tbb =
         if tb == tbb then add_pop (sz-size) cont
@@ -864,11 +886,11 @@ let rec comp_expr env exp sz cont =
         | [] -> assert false
         | try_sz :: tbb -> add_pop (sz-try_sz-4) (Kpoptrap :: loop try_sz tbb)
       in
-      let cont = loop sz !try_blocks in
+      let cont = loop sz stack_info.try_blocks in
       begin match args with
       | [arg] -> (* optim, argument passed in accumulator *)
-          comp_expr env arg sz cont
-      | _ -> comp_exit_args env args sz size cont
+          comp_expr stack_info env arg sz cont
+      | _ -> comp_exit_args stack_info env args sz size cont
       end
   | Ltrywith(body, id, handler) ->
       let (branch1, cont1) = make_branch cont in
@@ -876,33 +898,35 @@ let rec comp_expr env exp sz cont =
       let body_cont =
         Kpoptrap :: branch1 ::
         Klabel lbl_handler :: Kpush ::
-        comp_expr (add_var id (sz+1) env) handler (sz+1) (add_pop 1 cont1)
+        comp_expr
+          stack_info (add_var id (sz+1) env) handler (sz+1) (add_pop 1 cont1)
       in
-      try_blocks := sz :: !try_blocks;
-      let l = comp_expr env body (sz+4) body_cont in
-      try_blocks := List.tl !try_blocks;
+      let stack_info =
+        { stack_info with try_blocks = sz :: stack_info.try_blocks } in
+      let l = comp_expr stack_info env body (sz+4) body_cont in
       Kpushtrap lbl_handler :: l
   | Lifthenelse(cond, ifso, ifnot) ->
-      comp_binary_test env cond ifso ifnot sz cont
+      comp_binary_test stack_info env cond ifso ifnot sz cont
   | Lsequence(exp1, exp2) ->
-      comp_expr env exp1 sz (comp_expr env exp2 sz cont)
+      comp_expr stack_info env exp1 sz (comp_expr stack_info env exp2 sz cont)
   | Lwhile(cond, body) ->
       let lbl_loop = new_label() in
       let lbl_test = new_label() in
       Kbranch lbl_test :: Klabel lbl_loop :: Kcheck_signals ::
-        comp_expr env body sz
+        comp_expr stack_info env body sz
           (Klabel lbl_test ::
-            comp_expr env cond sz (Kbranchif lbl_loop :: add_const_unit cont))
+           comp_expr stack_info env cond sz
+             (Kbranchif lbl_loop :: add_const_unit cont))
   | Lfor(param, start, stop, dir, body) ->
       let lbl_loop = new_label() in
       let lbl_exit = new_label() in
       let offset = match dir with Upto -> 1 | Downto -> -1 in
       let comp = match dir with Upto -> Cgt | Downto -> Clt in
-      comp_expr env start sz
-        (Kpush :: comp_expr env stop (sz+1)
+      comp_expr stack_info env start sz
+        (Kpush :: comp_expr stack_info env stop (sz+1)
           (Kpush :: Kpush :: Kacc 2 :: Kintcomp comp :: Kbranchif lbl_exit ::
            Klabel lbl_loop :: Kcheck_signals ::
-           comp_expr (add_var param (sz+1) env) body (sz+2)
+           comp_expr stack_info (add_var param (sz+1) env) body (sz+2)
              (Kacc 1 :: Kpush :: Koffsetint offset :: Kassign 2 ::
               Kacc 1 :: Kintcomp Cne :: Kbranchif lbl_loop ::
               Klabel lbl_exit :: add_const_unit (add_pop 2 cont))))
@@ -937,7 +961,8 @@ let rec comp_expr env exp sz cont =
 *)
       let lbls = Array.make (Array.length acts) 0 in
       for i = Array.length acts-1 downto 0 do
-        let lbl,c1 = label_code (comp_expr env acts.(i) sz (branch :: !c)) in
+        let lbl,c1 =
+          label_code (comp_expr stack_info env acts.(i) sz (branch :: !c)) in
         lbls.(i) <- lbl ;
         c := discard_dead_code c1
       done ;
@@ -951,20 +976,19 @@ let rec comp_expr env exp sz cont =
       for i = sw.sw_numconsts - 1 downto 0 do
         lbl_consts.(i) <- lbls.(act_consts.(i))
       done;
-      comp_expr env arg sz (Kswitch(lbl_consts, lbl_blocks) :: !c)
+      comp_expr stack_info env arg sz (Kswitch(lbl_consts, lbl_blocks) :: !c)
   | Lstringswitch (arg,sw,d,loc) ->
-      comp_expr env (Matching.expand_stringswitch loc arg sw d) sz cont
+      comp_expr stack_info env
+        (Matching.expand_stringswitch loc arg sw d) sz cont
   | Lassign(id, expr) ->
       begin try
         let pos = Ident.find_same id env.ce_stack in
-        comp_expr env expr sz (Kassign(sz - pos) :: cont)
+        comp_expr stack_info env expr sz (Kassign(sz - pos) :: cont)
       with Not_found ->
         fatal_error "Bytegen.comp_expr: assign"
       end
   | Levent(lam, lev) ->
-      let ev_defname = match lev.lev_loc with
-        | Loc_unknown -> "??"
-        | Loc_known { loc = _; scopes } -> string_of_scopes scopes in
+      let ev_defname = string_of_scoped_location lev.lev_loc in
       let event kind info =
         { ev_pos = 0;                   (* patched in emitcode *)
           ev_module = !compunit_name;
@@ -992,15 +1016,15 @@ let rec comp_expr env exp sz cont =
       in
       begin match lev.lev_kind with
         Lev_before ->
-          let c = comp_expr env lam sz cont in
+          let c = comp_expr stack_info env lam sz cont in
           let ev = event Event_before Event_other in
           add_event ev c
       | Lev_function ->
-          let c = comp_expr env lam sz cont in
+          let c = comp_expr stack_info env lam sz cont in
           let ev = event Event_pseudo Event_function in
           add_event ev c
       | Lev_pseudo ->
-          let c = comp_expr env lam sz cont in
+          let c = comp_expr stack_info env lam sz cont in
           let ev = event Event_pseudo Event_other in
           add_event ev c
       | Lev_after ty ->
@@ -1011,7 +1035,7 @@ let rec comp_expr env exp sz cont =
           in
           if preserve_tailcall && is_tailcall cont then
             (* don't destroy tail call opt *)
-            comp_expr env lam sz cont
+            comp_expr stack_info env lam sz cont
           else begin
             let info =
               match lam with
@@ -1022,67 +1046,70 @@ let rec comp_expr env exp sz cont =
             in
             let ev = event (Event_after ty) info in
             let cont1 = add_event ev cont in
-            comp_expr env lam sz cont1
+            comp_expr stack_info env lam sz cont1
           end
       | Lev_module_definition _ ->
-          comp_expr env lam sz cont
+          comp_expr stack_info env lam sz cont
       end
   | Lifused (_, exp) ->
-      comp_expr env exp sz cont
+      comp_expr stack_info env exp sz cont
 
 (* Compile a list of arguments [e1; ...; eN] to a primitive operation.
    The values of eN ... e2 are pushed on the stack, e2 at top of stack,
    then e3, then ... The value of e1 is left in the accumulator. *)
 
-and comp_args env argl sz cont =
-  comp_expr_list env (List.rev argl) sz cont
+and comp_args stack_info env argl sz cont =
+  comp_expr_list stack_info env (List.rev argl) sz cont
 
-and comp_expr_list env exprl sz cont = match exprl with
+and comp_expr_list stack_info env exprl sz cont = match exprl with
     [] -> cont
-  | [exp] -> comp_expr env exp sz cont
+  | [exp] -> comp_expr stack_info env exp sz cont
   | exp :: rem ->
-      comp_expr env exp sz (Kpush :: comp_expr_list env rem (sz+1) cont)
+      comp_expr stack_info env exp sz
+        (Kpush :: comp_expr_list stack_info env rem (sz+1) cont)
 
-and comp_exit_args  env argl sz pos cont =
-   comp_expr_list_assign env (List.rev argl) sz pos cont
+and comp_exit_args stack_info env argl sz pos cont =
+   comp_expr_list_assign stack_info env (List.rev argl) sz pos cont
 
-and comp_expr_list_assign env exprl sz pos cont = match exprl with
+and comp_expr_list_assign stack_info env exprl sz pos cont = match exprl with
   | [] -> cont
   | exp :: rem ->
-      comp_expr env exp sz
-        (Kassign (sz-pos)::comp_expr_list_assign env rem sz (pos-1) cont)
+      comp_expr stack_info env exp sz
+        (Kassign (sz-pos)
+         ::comp_expr_list_assign stack_info env rem sz (pos-1) cont)
 
 (* Compile an if-then-else test. *)
 
-and comp_binary_test env cond ifso ifnot sz cont =
+and comp_binary_test stack_info env cond ifso ifnot sz cont =
   let cont_cond =
     if ifnot = Lconst const_unit then begin
       let (lbl_end, cont1) = label_code cont in
-      Kstrictbranchifnot lbl_end :: comp_expr env ifso sz cont1
+      Kstrictbranchifnot lbl_end :: comp_expr stack_info env ifso sz cont1
     end else
-    match code_as_jump ifso sz with
+    match code_as_jump stack_info ifso sz with
     | Some label ->
-      let cont = comp_expr env ifnot sz cont in
+      let cont = comp_expr stack_info env ifnot sz cont in
       Kbranchif label :: cont
-    | _ ->
-        match code_as_jump ifnot sz with
+    | None ->
+        match code_as_jump stack_info ifnot sz with
         | Some label ->
-            let cont = comp_expr env ifso sz cont in
+            let cont = comp_expr stack_info env ifso sz cont in
             Kbranchifnot label :: cont
-        | _ ->
+        | None ->
             let (branch_end, cont1) = make_branch cont in
-            let (lbl_not, cont2) = label_code(comp_expr env ifnot sz cont1) in
+            let (lbl_not, cont2) =
+              label_code(comp_expr stack_info env ifnot sz cont1) in
             Kbranchifnot lbl_not ::
-            comp_expr env ifso sz (branch_end :: cont2) in
+            comp_expr stack_info env ifso sz (branch_end :: cont2) in
 
-  comp_expr env cond sz cont_cond
+  comp_expr stack_info env cond sz cont_cond
 
 (**** Compilation of a code block (with tracking of stack usage) ****)
 
 let comp_block env exp sz cont =
-  max_stack_used := 0;
-  let code = comp_expr env exp sz cont in
-  let used_safe = !max_stack_used + Config.stack_safety_margin in
+  let stack_info = create_stack_info () in
+  let code = comp_expr stack_info env exp sz cont in
+  let used_safe = !(stack_info.max_stack_used) + Config.stack_safety_margin in
   if used_safe > Config.stack_threshold then
     Kconst(Const_base(Const_int used_safe)) ::
     Kccall("caml_ensure_stack_capacity", 1) ::
@@ -1121,29 +1148,25 @@ let comp_remainder cont =
 
 (**** Compilation of a lambda phrase ****)
 
-let compile_implementation modulename expr =
-  Stack.clear functions_to_compile;
+let reset () =
   label_counter := 0;
-  sz_static_raises := [] ;
+  compunit_name := "";
+  Stack.clear functions_to_compile
+
+let compile_implementation modulename expr =
+  reset ();
   compunit_name := modulename;
+  Fun.protect ~finally:reset (fun () ->
   let init_code = comp_block empty_env expr 0 [] in
   if Stack.length functions_to_compile > 0 then begin
     let lbl_init = new_label() in
     Kbranch lbl_init :: comp_remainder (Klabel lbl_init :: init_code)
   end else
-    init_code
+    init_code)
 
 let compile_phrase expr =
-  Stack.clear functions_to_compile;
-  label_counter := 0;
-  sz_static_raises := [] ;
+  reset ();
+  Fun.protect ~finally:reset (fun () ->
   let init_code = comp_block empty_env expr 1 [Kreturn 1] in
   let fun_code = comp_remainder [] in
-  (init_code, fun_code)
-
-let reset () =
-  label_counter := 0;
-  sz_static_raises := [];
-  compunit_name := "";
-  Stack.clear functions_to_compile;
-  max_stack_used := 0
+  (init_code, fun_code))
diff --git a/bytecomp/bytegen.mli b/bytecomp/bytegen.mli
index 80c222bff4..64bad6e423 100644
--- a/bytecomp/bytegen.mli
+++ b/bytecomp/bytegen.mli
@@ -20,7 +20,6 @@ open Instruct
 
 val compile_implementation: string -> lambda -> instruction list
 val compile_phrase: lambda -> instruction list * instruction list
-val reset: unit -> unit
 
 val merge_events:
   Instruct.debug_event -> Instruct.debug_event -> Instruct.debug_event