Fix #5127 Memory corruption with the VM

The bytecode interpreter ensures that the stack space available at some
points is above a static threshold. However, arbitrary large stack space
can be needed between two check points, leading to segmentation faults
in some cases.

We track the use of stack space at compilation time and add
an instruction to ensure greater stack capacity when required. This is
inspired from OCaml's PR#339 and PR#7168.

Patch written with Benjamin Grégoire.

9 files changed, 196 insertions, 103 deletions

diff --git a/kernel/byterun/coq_fix_code.c b/kernel/byterun/coq_fix_code.c
index 29e33d349..d5feafbf9 100644
--- a/kernel/byterun/coq_fix_code.c
+++ b/kernel/byterun/coq_fix_code.c
@@ -57,7 +57,7 @@ void init_arity () {
     arity[MAKEBLOCK1]=arity[MAKEBLOCK2]=arity[MAKEBLOCK3]=arity[MAKEBLOCK4]=
     arity[MAKEACCU]=arity[CONSTINT]=arity[PUSHCONSTINT]=arity[GRABREC]=
     arity[PUSHFIELDS]=arity[GETFIELD]=arity[SETFIELD]=
-    arity[BRANCH]=arity[ISCONST]= 1;
+    arity[BRANCH]=arity[ISCONST]=arity[ENSURESTACKCAPACITY]=1;
   /* instruction with two operands */
   arity[APPTERM]=arity[MAKEBLOCK]=arity[CLOSURE]=
   arity[ARECONST]=arity[PROJ]=2;
@@ -79,7 +79,7 @@ void * coq_stat_alloc (asize_t sz)
 
 value coq_makeaccu (value i) {
   code_t q;
-  code_t res = coq_stat_alloc(8);
+  code_t res = coq_stat_alloc(2 * sizeof(opcode_t));
   q = res;
   *q++ = VALINSTR(MAKEACCU);
   *q = (opcode_t)Int_val(i);
@@ -91,13 +91,13 @@ value coq_pushpop (value i) {
   int n;
   n = Int_val(i);
   if (n == 0) {
-    res = coq_stat_alloc(4);
+    res = coq_stat_alloc(sizeof(opcode_t));
     *res = VALINSTR(STOP);
     return (value)res;
   }
   else {
     code_t q;
-    res = coq_stat_alloc(12);
+    res = coq_stat_alloc(3 * sizeof(opcode_t));
     q = res;
     *q++ = VALINSTR(POP);
     *q++ = (opcode_t)n;
diff --git a/kernel/byterun/coq_instruct.h b/kernel/byterun/coq_instruct.h
index 8c5ab0ecb..dc3a88818 100644
--- a/kernel/byterun/coq_instruct.h
+++ b/kernel/byterun/coq_instruct.h
@@ -37,6 +37,7 @@ enum instructions {
   GETFIELD0, GETFIELD1, GETFIELD, 
   SETFIELD0, SETFIELD1, SETFIELD,
   PROJ,
+  ENSURESTACKCAPACITY,
   CONST0, CONST1, CONST2, CONST3, CONSTINT,
   PUSHCONST0, PUSHCONST1, PUSHCONST2, PUSHCONST3, PUSHCONSTINT,
   ACCUMULATE,
diff --git a/kernel/byterun/coq_interp.c b/kernel/byterun/coq_interp.c
index d634b726b..35abd011b 100644
--- a/kernel/byterun/coq_interp.c
+++ b/kernel/byterun/coq_interp.c
@@ -84,6 +84,14 @@ sp is a local copy of the global variable extern_sp. */
 #   define print_lint(i)
 #endif
 
+#define CHECK_STACK(num_args) {                                                \
+if (sp - num_args < coq_stack_threshold) {                                     \
+  coq_sp = sp;                                                                 \
+  realloc_coq_stack(num_args + Coq_stack_threshold / sizeof(value));           \
+  sp = coq_sp;                                                                 \
+ }                                                                             \
+}
+
 /* GC interface */
 #define Setup_for_gc { sp -= 2; sp[0] = accu; sp[1] = coq_env; coq_sp = sp; }
 #define Restore_after_gc { accu = sp[0]; coq_env = sp[1]; sp += 2; }
@@ -206,6 +214,9 @@ value coq_interprete
   sp = coq_sp;
   pc = coq_pc;
   accu = coq_accu;
+
+  CHECK_STACK(0);
+
 #ifdef THREADED_CODE
   goto *(void *)(coq_jumptbl_base + *pc++); /* Jump to the first instruction */
 #else
@@ -362,7 +373,7 @@ value coq_interprete
 	coq_extra_args = *pc - 1;
 	pc = Code_val(accu);
 	coq_env = accu;
-	goto check_stacks;
+	goto check_stack;
       }
       Instruct(APPLY1) {
 	value arg1 = sp[0];
@@ -379,7 +390,7 @@ value coq_interprete
 	pc = Code_val(accu);
 	coq_env = accu;
 	coq_extra_args = 0;
-	goto check_stacks;
+	goto check_stack;
       }
       Instruct(APPLY2) {
 	value arg1 = sp[0];
@@ -394,7 +405,7 @@ value coq_interprete
 	pc = Code_val(accu);
 	coq_env = accu;
 	coq_extra_args = 1;
-	goto check_stacks;
+	goto check_stack;
       }
       Instruct(APPLY3) {
 	value arg1 = sp[0];
@@ -411,17 +422,13 @@ value coq_interprete
 	pc = Code_val(accu);
 	coq_env = accu;
 	coq_extra_args = 2;
-	goto check_stacks;
+	goto check_stack;
       }
       /* Stack checks */
       
-    check_stacks:
-      print_instr("check_stacks");
-      if (sp < coq_stack_threshold) {
-	coq_sp = sp;
-	realloc_coq_stack(Coq_stack_threshold);
-	sp = coq_sp;
-      }
+    check_stack:
+      print_instr("check_stack");
+      CHECK_STACK(0);
       /* We also check for signals */
       if (caml_signals_are_pending) {
 	/* If there's a Ctrl-C, we reset the vm */
@@ -430,6 +437,16 @@ value coq_interprete
       }
       Next;
 
+      Instruct(ENSURESTACKCAPACITY) {
+        print_instr("ENSURESTACKCAPACITY");
+        int size = *pc++;
+        /* CHECK_STACK may trigger here a useless allocation because of the
+        threshold, but check_stack: often does it anyway, so we prefer to
+        factorize the code. */
+        CHECK_STACK(size);
+        Next;
+      }
+
       Instruct(APPTERM) {
 	int nargs = *pc++;
 	int slotsize = *pc;
@@ -444,7 +461,7 @@ value coq_interprete
 	pc = Code_val(accu);
 	coq_env = accu;
 	coq_extra_args += nargs - 1;
-	goto check_stacks;
+	goto check_stack;
       }
       Instruct(APPTERM1) {
 	value arg1 = sp[0];
@@ -453,7 +470,7 @@ value coq_interprete
 	sp[0] = arg1;
 	pc = Code_val(accu);
 	coq_env = accu;
-	goto check_stacks;
+	goto check_stack;
       }
       Instruct(APPTERM2) {
 	value arg1 = sp[0];
@@ -466,7 +483,7 @@ value coq_interprete
 	print_lint(accu);
 	coq_env = accu;
 	coq_extra_args += 1;
-	goto check_stacks;
+	goto check_stack;
       }
       Instruct(APPTERM3) {
 	value arg1 = sp[0];
@@ -480,7 +497,7 @@ value coq_interprete
 	pc = Code_val(accu);
 	coq_env = accu;
 	coq_extra_args += 2;
-	goto check_stacks;
+	goto check_stack;
       }
       
       Instruct(RETURN) {
@@ -511,6 +528,7 @@ value coq_interprete
 	int num_args = Wosize_val(coq_env) - 2;
 	int i;
 	print_instr("RESTART");
+        CHECK_STACK(num_args);
 	sp -= num_args;
 	for (i = 0; i < num_args; i++) sp[i] = Field(coq_env, i + 2);
 	coq_env = Field(coq_env, 1);
@@ -870,29 +888,7 @@ value coq_interprete
        	sp++;
 	Next; 
       }
-  
-      /* *sp = accu;
-	* Netoyage des cofix *
-        size = Wosize_val(accu);
-        for (i = 2; i < size; i++) {
- 	  accu = Field(*sp, i);
-          if (IS_EVALUATED_COFIX(accu)) {
-            size_aux = Wosize_val(accu);
-            *--sp = accu;
-	    Alloc_small(accu, size_aux, Accu_tag);
-	    for(j = 0; j < size_aux; j++) Field(accu, j) = Field(*sp, j);
-	    *sp = accu;
-	    Alloc_small(accu, 1, ATOM_COFIX_TAG);
-	    Field(accu, 0) = Field(Field(*sp, 1), 0);
-	    caml_modify(&Field(*sp, 1), accu);
-	    accu = *sp; sp++;
-            caml_modify(&Field(*sp, i), accu);
-	  }
-	}
-	sp++;
-	Next; 
-	} */
-      
+       
       Instruct(SETFIELD){
 	print_instr("SETFIELD");
 	caml_modify(&Field(accu, *pc),*sp);
@@ -986,28 +982,31 @@ value coq_interprete
       }  
       Instruct(MAKESWITCHBLOCK) {
 	print_instr("MAKESWITCHBLOCK");
-	*--sp = accu;
-	accu = Field(accu,1);
+	*--sp = accu; // Save matched block on stack
+	accu = Field(accu,1); // Save atom to accu register
 	switch (Tag_val(accu)) {
-	case ATOM_COFIX_TAG: 
+	case ATOM_COFIX_TAG: // We are forcing a cofix
 	  {
 	    mlsize_t i, nargs;
 	    print_instr("COFIX_TAG");
 	    sp-=2;
 	    pc++;
+            // Push the return address
 	    sp[0] = (value) (pc + *pc);
 	    sp[1] = coq_env;
-	    coq_env = Field(accu,0);
-	    accu = sp[2];
-	    sp[2] = Val_long(coq_extra_args);
-	    nargs = Wosize_val(accu) - 2;
+	    coq_env = Field(accu,0); // Pointer to suspension
+	    accu = sp[2]; // Save accumulator to accu register
+	    sp[2] = Val_long(coq_extra_args); // Push number of args for return
+	    nargs = Wosize_val(accu) - 2; // Number of args = size of accumulator - 1 (accumulator code) - 1 (atom)
+            // Push arguments to stack
+            CHECK_STACK(nargs+1);
 	    sp -= nargs;
-	    for (i = 0; i < nargs; i++) sp[i] = Field(accu, i + 2);
-	    *--sp = accu;
+	    for (i = 0; i < nargs; i++) sp[i] = Field(accu, i + 2); 
+	    *--sp = accu; // Last argument is the pointer to the suspension
 	    print_lint(nargs);
 	    coq_extra_args = nargs;
-	    pc = Code_val(coq_env);
-	    goto check_stacks;
+	    pc = Code_val(coq_env); // Trigger evaluation
+	    goto check_stack;
 	  }
 	case ATOM_COFIXEVALUATED_TAG: 
 	  {
@@ -1464,26 +1463,32 @@ value coq_push_val(value v) {
 
 value coq_push_arguments(value args) {
   int nargs,i;
+  value * sp = coq_sp;
   nargs = Wosize_val(args) - 2;
+  CHECK_STACK(nargs);
   coq_sp -= nargs;
   print_instr("push_args");print_int(nargs);
   for(i = 0; i < nargs; i++) coq_sp[i] = Field(args, i+2);
   return Val_unit;
 }
 
-value coq_push_vstack(value stk) {
+value coq_push_vstack(value stk, value max_stack_size) {
   int len,i;
+  value * sp = coq_sp;
   len = Wosize_val(stk);
+  CHECK_STACK(len);
   coq_sp -= len;
   print_instr("push_vstack");print_int(len);
    for(i = 0; i < len; i++) coq_sp[i] = Field(stk,i);
+  sp = coq_sp;
+  CHECK_STACK(uint32_of_value(max_stack_size));
   return Val_unit;
 }
 
 value  coq_interprete_ml(value tcode, value a, value e, value ea) {
   print_instr("coq_interprete");
   return coq_interprete((code_t)tcode, a, e, Long_val(ea));
-   print_instr("end coq_interprete");
+  print_instr("end coq_interprete");
 }
 
 value coq_eval_tcode (value tcode, value e) {
diff --git a/kernel/byterun/coq_memory.c b/kernel/byterun/coq_memory.c
index c9bcdc32f..45cfae509 100644
--- a/kernel/byterun/coq_memory.c
+++ b/kernel/byterun/coq_memory.c
@@ -130,6 +130,7 @@ value init_coq_vm(value unit) /* ML */
   return Val_unit;;
 }
 
+/* [required_space] is a size in words */
 void realloc_coq_stack(asize_t required_space)
 {
   asize_t size;
diff --git a/kernel/cbytecodes.ml b/kernel/cbytecodes.ml
index f9cf2691e..40f968f36 100644
--- a/kernel/cbytecodes.ml
+++ b/kernel/cbytecodes.ml
@@ -43,7 +43,7 @@ type structured_constant =
 type reloc_table = (tag * int) array
 
 type annot_switch =
-   {ci : case_info; rtbl : reloc_table; tailcall : bool}
+   {ci : case_info; rtbl : reloc_table; tailcall : bool; max_stack_size : int}
 
 module Label =
   struct
@@ -87,6 +87,7 @@ type instruction =
   | Ksequence of bytecodes * bytecodes
   | Kproj of int * Constant.t  (* index of the projected argument,
                                            name of projection *)
+  | Kensurestackcapacity of int
 (* spiwack: instructions concerning integers *)
   | Kbranch of Label.t                  (* jump to label *)
   | Kaddint31                           (* adds the int31 in the accu
@@ -249,6 +250,8 @@ let rec pp_instr i =
 
   | Kproj(n,p) -> str "proj " ++ int n ++ str " " ++ Constant.print p
 
+  | Kensurestackcapacity size -> str "growstack " ++ int size
+
   | Kaddint31 -> str "addint31"
   | Kaddcint31 -> str "addcint31"
   | Kaddcarrycint31 -> str "addcarrycint31"
diff --git a/kernel/cbytecodes.mli b/kernel/cbytecodes.mli
index 6fa0841af..5a854e717 100644
--- a/kernel/cbytecodes.mli
+++ b/kernel/cbytecodes.mli
@@ -39,7 +39,7 @@ val pp_struct_const : structured_constant -> Pp.std_ppcmds
 type reloc_table = (tag * int) array
 
 type annot_switch =
-   {ci : case_info; rtbl : reloc_table; tailcall : bool}
+   {ci : case_info; rtbl : reloc_table; tailcall : bool; max_stack_size : int}
 
 module Label :
   sig
@@ -84,6 +84,7 @@ type instruction =
   | Ksequence of bytecodes * bytecodes
   | Kproj of int * Constant.t  (** index of the projected argument,
                                             name of projection *)
+  | Kensurestackcapacity of int
 
 (** spiwack: instructions concerning integers *)
   | Kbranch of Label.t                  (** jump to label, is it needed ? *)
diff --git a/kernel/cbytegen.ml b/kernel/cbytegen.ml
index 77eac9ee9..1d538e294 100644
--- a/kernel/cbytegen.ml
+++ b/kernel/cbytegen.ml
@@ -91,6 +91,11 @@ open Pre_env
 (* In Cfxe_t accumulators, we need to store [fcofixi] for testing         *)
 (* conversion of cofixpoints (which is intentional).                      *)
 
+module Config = struct
+  let stack_threshold = 256 (* see byterun/coq_memory.h *)
+  let stack_safety_margin = 15
+end
+
 type argument = ArgConstr of Constr.t | ArgUniv of Univ.Level.t
 
 let empty_fv = { size= 0;  fv_rev = [] }
@@ -107,6 +112,26 @@ let empty_comp_env ?(univs=0) ()=
     in_env = ref empty_fv
   }
 
+(* Maximal stack size reached during the current function body. Used to
+   reallocate the stack if we lack space. *)
+let max_stack_size = ref 0
+
+let set_max_stack_size stack_size =
+  if stack_size > !max_stack_size then
+    max_stack_size := stack_size
+
+let ensure_stack_capacity f x =
+  let old = !max_stack_size in
+  max_stack_size := 0;
+  let code = f x in
+  let used_safe =
+    !max_stack_size + Config.stack_safety_margin
+  in
+  max_stack_size := old;
+  if used_safe > Config.stack_threshold then
+    Kensurestackcapacity used_safe :: code
+  else code
+
 (*i Creation functions for comp_env *)
 
 let rec add_param n sz l =
@@ -375,14 +400,28 @@ let const_bn tag args =
   else
     Const_bn(last_variant_tag, Array.append [|Const_b0 (tag - last_variant_tag) |] args)
 
-  
-let code_makeblock arity tag cont = 
+(*
+If [tag] hits the OCaml limitation for non constant constructors, we switch to
+another representation for the remaining constructors:
+[last_variant_tag|tag - last_variant_tag|args]
+
+We subtract last_variant_tag for efficiency of match interpretation.
+ *)
+
+let nest_block tag arity cont =
+  Kconst (Const_b0 (tag - last_variant_tag)) ::
+    Kmakeblock(arity+1, last_variant_tag) :: cont
+
+let code_makeblock ~stack_size ~arity ~tag cont = 
   if tag < last_variant_tag then
     Kmakeblock(arity, tag) :: cont
-  else
-    Kpush :: Kconst (Const_b0 (tag - last_variant_tag)) ::
-    Kmakeblock(arity+1, last_variant_tag) :: cont
+  else begin
+    set_max_stack_size (stack_size + 1);
+    Kpush :: nest_block tag arity cont
+  end
 
+(* [code_construct] compiles an abstracted constructor dropping parameters and
+   updates [fun_code] *)
 (* Inv : nparam + arity > 0 *)
 let code_construct tag nparams arity cont =
   let f_cont =
@@ -391,11 +430,11 @@ let code_construct tag nparams arity cont =
 	[Kconst (Const_b0 tag); Kreturn 0]
        else if tag < last_variant_tag then
          [Kacc 0; Kpop 1; Kmakeblock(arity, tag); Kreturn 0]
-       else 
-         [Kconst (Const_b0 (tag - last_variant_tag)); 
-          Kmakeblock(arity+1, last_variant_tag); Kreturn 0])
+       else
+         nest_block tag arity [Kreturn 0])
     in
     let lbl = Label.create() in
+    (* No need to grow the stack here, as the function does not push stuff. *)
     fun_code := [Ksequence (add_grab (nparams+arity) lbl f_cont,!fun_code)];
     Kclosure(lbl,0) :: cont
 
@@ -540,11 +579,12 @@ let compile_fv_elem reloc fv sz cont =
 let rec compile_fv reloc l sz cont =
   match l with
   | [] -> cont
-  | [fvn] -> compile_fv_elem reloc fvn sz cont
+  | [fvn] -> set_max_stack_size (sz + 1); compile_fv_elem reloc fvn sz cont
   | fvn :: tl ->
       compile_fv_elem reloc fvn sz
 	(Kpush :: compile_fv reloc tl (sz + 1) cont)
 
+
 (* Compiling constants *)
 
 let rec get_alias env kn =
@@ -559,6 +599,7 @@ let rec get_alias env kn =
 
 (* sz is the size of the local stack *)
 let rec compile_constr reloc c sz cont =
+  set_max_stack_size sz;
   match kind_of_term c with
   | Meta _ -> invalid_arg "Cbytegen.compile_constr : Meta"
   | Evar _ -> invalid_arg "Cbytegen.compile_constr : Evar"
@@ -607,6 +648,7 @@ let rec compile_constr reloc c sz cont =
 	compile_str_cst reloc (Bstrconst (Const_sorts (Type uglob))) sz cont
       else
 	let compile_get_univ reloc idx sz cont =
+          set_max_stack_size sz;
 	  compile_fv_elem reloc (FVuniv_var idx) sz cont
 	in
         comp_app compile_str_cst compile_get_univ reloc
@@ -626,7 +668,8 @@ let rec compile_constr reloc c sz cont =
       let r_fun = comp_env_fun arity in
       let lbl_fun = Label.create() in
       let cont_fun =
-	compile_constr r_fun body arity [Kreturn arity] in
+        ensure_stack_capacity (compile_constr r_fun body arity) [Kreturn arity]
+      in
       fun_code := [Ksequence(add_grab arity lbl_fun cont_fun,!fun_code)];
       let fv = fv r_fun in
       compile_fv reloc fv.fv_rev sz (Kclosure(lbl_fun,fv.size) :: cont)
@@ -646,9 +689,10 @@ let rec compile_constr reloc c sz cont =
       (* Compilation des types *)
       let env_type = comp_env_fix_type rfv in
       for i = 0 to ndef - 1 do
-	let lbl,fcode =
-	  label_code
-	    (compile_constr env_type type_bodies.(i) 0 [Kstop]) in
+        let fcode =
+          ensure_stack_capacity (compile_constr env_type type_bodies.(i) 0) [Kstop]
+        in
+	let lbl,fcode = label_code fcode in
 	lbl_types.(i) <- lbl;
 	fun_code := [Ksequence(fcode,!fun_code)]
       done;
@@ -658,7 +702,8 @@ let rec compile_constr reloc c sz cont =
 	let arity = List.length params in
 	let env_body = comp_env_fix ndef i arity rfv in
 	let cont1 =
-	  compile_constr env_body body arity [Kreturn arity] in
+	  ensure_stack_capacity (compile_constr env_body body arity) [Kreturn arity]
+        in
 	let lbl = Label.create () in
 	lbl_bodies.(i) <- lbl;
 	let fcode =  add_grabrec rec_args.(i) arity lbl cont1 in
@@ -676,9 +721,10 @@ let rec compile_constr reloc c sz cont =
       let rfv = ref empty_fv in
       let env_type = comp_env_cofix_type ndef rfv in
       for i = 0 to ndef - 1 do
-	let lbl,fcode =
-	  label_code
-	    (compile_constr env_type type_bodies.(i) 0 [Kstop]) in
+        let fcode =
+          ensure_stack_capacity (compile_constr env_type type_bodies.(i) 0) [Kstop]
+        in
+	let lbl,fcode = label_code fcode in
 	lbl_types.(i) <- lbl;
 	fun_code := [Ksequence(fcode,!fun_code)]
       done;
@@ -688,14 +734,17 @@ let rec compile_constr reloc c sz cont =
 	let arity = List.length params in
 	let env_body = comp_env_cofix ndef arity rfv in
 	let lbl = Label.create () in
-	let cont1 =
-	  compile_constr env_body body (arity+1) (cont_cofix arity) in
-	let cont2 =
-	  add_grab (arity+1) lbl cont1 in
+        let comp arity =
+          (* 4 stack slots are needed to update the cofix when forced *)
+          set_max_stack_size (arity + 4);
+          compile_constr env_body body (arity+1) (cont_cofix arity)
+        in
+	let cont = ensure_stack_capacity comp arity in
 	lbl_bodies.(i) <- lbl;
-	fun_code := [Ksequence(cont2,!fun_code)];
+	fun_code := [Ksequence(add_grab (arity+1) lbl cont,!fun_code)];
       done;
       let fv = !rfv in
+      set_max_stack_size (sz + fv.size + ndef + 2);
       compile_fv reloc fv.fv_rev sz
 	(Kclosurecofix(fv.size, init, lbl_types, lbl_bodies) :: cont)
 
@@ -713,9 +762,11 @@ let rec compile_constr reloc c sz cont =
       let lbl_eblocks = Array.make neblock Label.no in 
       let branch1,cont = make_branch cont in
       (* Compiling return type *)
-      let lbl_typ,fcode =
-	label_code (compile_constr reloc t sz [Kpop sz; Kstop])
-      in fun_code := [Ksequence(fcode,!fun_code)];
+      let fcode =
+        ensure_stack_capacity (compile_constr reloc t sz) [Kpop sz; Kstop]
+      in
+      let lbl_typ,fcode = label_code fcode in
+      fun_code := [Ksequence(fcode,!fun_code)];
       (* Compiling branches *)
       let lbl_sw = Label.create () in
       let sz_b,branch,is_tailcall =
@@ -725,14 +776,9 @@ let rec compile_constr reloc c sz cont =
 	  sz, branch1, true
 	| _ -> sz+3, Kjump, false
       in
-      let annot = {ci = ci; rtbl = tbl; tailcall = is_tailcall} in
-     (* Compiling branch for accumulators *)
-      let lbl_accu, code_accu =
-	label_code(Kmakeswitchblock(lbl_typ,lbl_sw,annot,sz) :: branch::cont)
-      in
-      lbl_blocks.(0) <- lbl_accu;
-      let c = ref code_accu in
-      (* perform the extra match if needed (to many block constructors) *)
+
+      let c = ref cont in
+      (* Perform the extra match if needed (too many block constructors) *)
       if neblock <> 0 then begin
         let lbl_b, code_b = 
           label_code (
@@ -762,14 +808,34 @@ let rec compile_constr reloc c sz cont =
 	      compile_constr reloc branchs.(i) (sz_b+arity)
 		(Kappterm(arity,sz_appterm) :: !c) in
           let code_b = 
-            if tag < last_variant_tag then Kpushfields arity :: code_b
-            else Kacc 0::Kpop 1::Kpushfields(arity+1)::Kpop 1::code_b in
+            if tag < last_variant_tag then begin
+                set_max_stack_size (sz_b + arity);
+                Kpushfields arity :: code_b
+              end
+            else begin
+                set_max_stack_size (sz_b + arity + 1);
+                Kacc 0::Kpop 1::Kpushfields(arity+1)::Kpop 1::code_b
+              end
+          in
           let lbl_b,code_b = label_code code_b in
           if tag < last_variant_tag then lbl_blocks.(tag) <- lbl_b
           else lbl_eblocks.(tag - last_variant_tag) <- lbl_b;
 	  c := code_b
       done;
-      c := Klabel lbl_sw :: Kswitch(lbl_consts,lbl_blocks) :: !c;
+
+      let annot =
+        {ci = ci; rtbl = tbl; tailcall = is_tailcall;
+         max_stack_size = !max_stack_size - sz}
+      in
+
+     (* Compiling branch for accumulators *)
+      let lbl_accu, code_accu =
+        set_max_stack_size (sz+3);
+	label_code(Kmakeswitchblock(lbl_typ,lbl_sw,annot,sz) :: branch :: !c)
+      in
+      lbl_blocks.(0) <- lbl_accu;
+
+      c := Klabel lbl_sw :: Kswitch(lbl_consts,lbl_blocks) :: code_accu;
       let code_sw =
  	match branch1 with
         (* spiwack : branch1 can't be a lbl anymore it's a Branch instead
@@ -786,12 +852,14 @@ let rec compile_constr reloc c sz cont =
 	code_sw)
 
 and compile_str_cst reloc sc sz cont =
+  set_max_stack_size sz;
   match sc with
   | Bconstr c -> compile_constr reloc c sz cont
   | Bstrconst sc -> Kconst sc :: cont
   | Bmakeblock(tag,args) ->
-      let nargs = Array.length args in
-      comp_args compile_str_cst reloc args sz (code_makeblock nargs tag cont)
+      let arity = Array.length args in
+      let cont = code_makeblock ~stack_size:(sz+arity-1) ~arity ~tag cont in
+      comp_args compile_str_cst reloc args sz cont
   | Bconstruct_app(tag,nparams,arity,args) ->
       if Int.equal (Array.length args) 0 then
         code_construct tag nparams arity cont
@@ -805,6 +873,7 @@ and compile_str_cst reloc sc sz cont =
 (* spiwack : compilation of constants with their arguments.
    Makes a special treatment with 31-bit integer addition *)
 and compile_get_global reloc (kn,u) sz cont =
+  set_max_stack_size sz;
   let kn = get_alias !global_env kn in
   if Univ.Instance.is_empty u then
     Kgetglobal kn :: cont
@@ -813,11 +882,13 @@ and compile_get_global reloc (kn,u) sz cont =
       compile_universe reloc () (Univ.Instance.to_array u) sz cont
 
 and compile_universe reloc uni sz cont =
+  set_max_stack_size sz;
   match Univ.Level.var_index uni with
   | None -> Kconst (Const_univ_level uni) :: cont
   | Some idx -> pos_universe_var idx reloc sz :: cont
 
 and compile_const reloc kn u args sz cont =
+  set_max_stack_size sz;
   let nargs = Array.length args in
   (* spiwack: checks if there is a specific way to compile the constant
               if there is not, Not_found is raised, and the function
@@ -879,7 +950,7 @@ let compile fail_on_error ?universes:(universes=0) env c =
     let reloc, init_code =
       if Int.equal universes 0 then
         let reloc = empty_comp_env () in
-        reloc, compile_constr reloc c 0 cont
+        reloc, ensure_stack_capacity (compile_constr reloc c 0) cont
       else
         (* We are going to generate a lambda, but merge the universe closure
          * with the function closure if it exists.
@@ -896,11 +967,16 @@ let compile fail_on_error ?universes:(universes=0) env c =
         let r_fun = comp_env_fun ~univs:universes arity in
         let lbl_fun = Label.create () in
         let cont_fun =
-          compile_constr r_fun body full_arity [Kreturn full_arity]
+          ensure_stack_capacity (compile_constr r_fun body full_arity)
+                         [Kreturn full_arity]
         in
         fun_code := [Ksequence(add_grab full_arity lbl_fun cont_fun,!fun_code)];
         let fv = fv r_fun in
-        reloc, compile_fv reloc fv.fv_rev 0 (Kclosure(lbl_fun,fv.size) :: cont)
+        let init_code =
+          ensure_stack_capacity (compile_fv reloc fv.fv_rev 0)
+                         (Kclosure(lbl_fun,fv.size) :: cont)
+        in
+        reloc, init_code
     in
     let fv = List.rev (!(reloc.in_env).fv_rev) in
     (if !Flags.dump_bytecode then
diff --git a/kernel/cemitcodes.ml b/kernel/cemitcodes.ml
index 57e32684a..ccea67e41 100644
--- a/kernel/cemitcodes.ml
+++ b/kernel/cemitcodes.ml
@@ -226,6 +226,7 @@ let emit_instr = function
       else (out opSETFIELD;out_int n)
   | Ksequence _ -> invalid_arg "Cemitcodes.emit_instr"
   | Kproj (n,p) -> out opPROJ; out_int n; slot_for_const (Const_proj p)
+  | Kensurestackcapacity size -> out opENSURESTACKCAPACITY; out_int size
   (* spiwack *)
   | Kbranch lbl -> out opBRANCH; out_label lbl
   | Kaddint31 -> out opADDINT31
diff --git a/kernel/vm.ml b/kernel/vm.ml
index 702987643..781a05884 100644
--- a/kernel/vm.ml
+++ b/kernel/vm.ml
@@ -170,7 +170,7 @@ type whd =
 external push_ra : tcode -> unit = "coq_push_ra"
 external push_val : values -> unit = "coq_push_val"
 external push_arguments : arguments -> unit = "coq_push_arguments"
-external push_vstack : vstack -> unit = "coq_push_vstack"
+external push_vstack : vstack -> int -> unit = "coq_push_vstack"
 
 
 (* interpreteur *)
@@ -206,7 +206,9 @@ let apply_varray vf varray =
   else
     begin
       push_ra stop;
-      push_vstack varray;
+      (* The fun code of [vf] will make sure we have enough stack, so we put 0
+         here. *)
+      push_vstack varray 0;
       interprete (fun_code vf) vf (Obj.magic vf) (n - 1)
     end
 
@@ -560,7 +562,9 @@ let check_switch sw1 sw2 = sw1.sw_annot.rtbl = sw2.sw_annot.rtbl
 let case_info sw = sw.sw_annot.ci
 
 let type_of_switch sw =
-  push_vstack sw.sw_stk;
+  (* The fun code of types will make sure we have enough stack, so we put 0
+  here. *)
+  push_vstack sw.sw_stk 0;
   interprete sw.sw_type_code crazy_val sw.sw_env 0
 
 let branch_arg k (tag,arity) =
@@ -580,9 +584,10 @@ let branch_arg k (tag,arity) =
 let apply_switch sw arg =
   let tc = sw.sw_annot.tailcall in
   if tc then
-    (push_ra stop;push_vstack sw.sw_stk)
+    (push_ra stop;push_vstack sw.sw_stk sw.sw_annot.max_stack_size)
   else
-    (push_vstack sw.sw_stk; push_ra (popstop_code (Array.length sw.sw_stk)));
+    (push_vstack sw.sw_stk sw.sw_annot.max_stack_size;
+     push_ra (popstop_code (Array.length sw.sw_stk)));
   interprete sw.sw_code arg sw.sw_env 0
 
 let branch_of_switch k sw =