Compiler: check stack compatibility when jumping

CHANGES.md

@@ -16,7 +16,7 @@
 * Compiler: fix global flow analysis (#1494)
 * Compiler: fix js parser/printer wrt async functions (#1515)
 * Compiler: fix free variables pass wrt parameters' default value (#1521)
-
+* Compiler: fix internal invariant (continuation)
 # 5.4.0 (2023-07-06) - Lille
 
 ## Bug fixes

TODO.txt

@@ -7,9 +7,6 @@ Compiler optimizations
 - constant hoisting (including functions, out of loops and functions)
 - inline also partially applied functions
 
-- we should check stack compatibility when parsing:
-  when jumping somewhere, the stack should keep the same shape
-
 - cross-function optimizations
 
 - deadcode elimination inside blocks

compiler/lib/code.ml

@@ -687,7 +687,7 @@ let invariant { blocks; start; _ } =
     let defs = Var.ISet.empty () in
     let check_cont (cont, args) =
       let b = Addr.Map.find cont blocks in
-      assert (List.length args >= List.length b.params)
+      assert (List.length args = List.length b.params)
     in
     let define x =
       if check_defs

compiler/lib/deadcode.ml

@@ -130,7 +130,7 @@ let rec filter_args st pl al =
   match pl, al with
   | x :: pl, y :: al ->
       if st.live.(Var.idx x) > 0 then y :: filter_args st pl al else filter_args st pl al
-  | [], _ -> []
+  | [], [] -> []
   | _ -> assert false
 
 let filter_cont blocks st (pc, args) =
@@ -184,7 +184,8 @@ let rec add_arg_dep defs params args =
   | x :: params, y :: args ->
       add_def defs x (Var y);
       add_arg_dep defs params args
-  | _ -> ()
+  | [], [] -> ()
+  | _ -> assert false
 
 let add_cont_dep blocks defs (pc, args) =
   match try Some (Addr.Map.find pc blocks) with Not_found -> None with

compiler/lib/effects.ml

@@ -926,6 +926,8 @@ let remove_empty_blocks ~live_vars (p : Code.program) : Code.program =
 let f (p, live_vars) =
   let t = Timer.make () in
   let p = remove_empty_blocks ~live_vars p in
+  (* [remove_empty_blocks] can affect [Deadcode.variable_uses] *)
+  let p, live_vars = Deadcode.f p in
   let flow_info = Global_flow.f ~fast:false p in
   let cps_needed = Partial_cps_analysis.f p flow_info in
   let p, cps_needed = rewrite_toplevel ~cps_needed p in

compiler/lib/flow.ml

@@ -81,7 +81,8 @@ let rec arg_deps vars deps defs params args =
       add_dep deps x y;
       add_assign_def vars defs x y;
       arg_deps vars deps defs params args
-  | _ -> ()
+  | [], [] -> ()
+  | _ -> assert false
 
 let cont_deps blocks vars deps defs (pc, args) =
   let block = Addr.Map.find pc blocks in

compiler/lib/global_flow.ml

@@ -139,7 +139,8 @@ let rec arg_deps st ?ignore params args =
       | Some y' when Var.equal y y' -> ()
       | _ -> add_assign_def st x y);
       arg_deps st params args
-  | _ -> ()
+  | [], [] -> ()
+  | _ -> assert false
 
 let cont_deps blocks st ?ignore (pc, args) =
   let block = Addr.Map.find pc blocks in

compiler/lib/parse_bytecode.ml

@@ -332,7 +332,8 @@ end = struct
     | v1 :: r1, v2 :: r2 ->
         Var.propagate_name v1 v2;
         propagate r1 r2
-    | _ -> ()
+    | [], [] -> ()
+    | _ -> assert false
 
   let fold t f acc =
     Int_table.fold (fun k { event; _ } acc -> f k event acc) t.events_by_pc acc
@@ -646,8 +647,6 @@ module State = struct
 
   let set_accu st x loc = { st with accu = Var (x, loc) }
 
-  let clear_accu st = { st with accu = Dummy }
-
   let peek n st = elt_to_var (List.nth st.stack n)
 
   let grab n st = List.map (list_start n st.stack) ~f:elt_to_var, pop n st
@@ -811,7 +810,7 @@ let get_global state instrs i loc =
         g.vars.(i) <- Some x;
         x, state, instrs)
 
-let tagged_blocks = ref Addr.Set.empty
+let tagged_blocks = ref Addr.Map.empty
 
 let compiled_blocks = ref Addr.Map.empty
 
@@ -854,27 +853,44 @@ let ( ||| ) x y =
   | _ -> x
 
 let rec compile_block blocks debug_data code pc state =
-  if not (Addr.Set.mem pc !tagged_blocks)
-  then (
-    let limit = Blocks.next blocks pc in
-    assert (limit > pc);
-    if debug_parser () then Format.eprintf "Compiling from %d to %d@." pc (limit - 1);
-    let state = State.start_block pc state in
-    tagged_blocks := Addr.Set.add pc !tagged_blocks;
-    let instr, last, state' =
-      compile { blocks; code; limit; debug = debug_data } pc state []
-    in
-    assert (not (Addr.Map.mem pc !compiled_blocks));
-    compiled_blocks := Addr.Map.add pc (state, List.rev instr, last) !compiled_blocks;
-    match fst last with
-    | Branch (pc', _) | Poptrap (pc', _) ->
-        compile_block blocks debug_data code pc' state'
-    | Cond (_, (pc1, _), (pc2, _)) ->
-        compile_block blocks debug_data code pc1 state';
-        compile_block blocks debug_data code pc2 state'
-    | Switch (_, _) -> ()
-    | Pushtrap _ -> ()
-    | Raise _ | Return _ | Stop -> ())
+  match Addr.Map.find_opt pc !tagged_blocks with
+  | Some old_state -> (
+      let rec check (xs : State.elt list) (ys : State.elt list) =
+        match xs, ys with
+        | Var _ :: xs, Var _ :: ys -> check xs ys
+        | Dummy :: xs, Dummy :: ys -> check xs ys
+        | [], [] -> ()
+        | Var _ :: _, Dummy :: _ -> assert false
+        | Dummy :: _, Var _ :: _ -> assert false
+        | _ :: _, [] -> assert false
+        | [], _ :: _ -> assert false
+      in
+      check old_state.State.stack state.State.stack;
+      match old_state.State.accu, state.State.accu with
+      | Dummy, Dummy -> ()
+      | Var _, Var _ -> ()
+      | Var _, Dummy -> assert false
+      | Dummy, Var _ -> assert false)
+  | None -> (
+      let limit = Blocks.next blocks pc in
+      assert (limit > pc);
+      if debug_parser () then Format.eprintf "Compiling from %d to %d@." pc (limit - 1);
+      let state = State.start_block pc state in
+      tagged_blocks := Addr.Map.add pc state !tagged_blocks;
+      let instr, last, state' =
+        compile { blocks; code; limit; debug = debug_data } pc state []
+      in
+      assert (not (Addr.Map.mem pc !compiled_blocks));
+      compiled_blocks := Addr.Map.add pc (state, List.rev instr, last) !compiled_blocks;
+      match fst last with
+      | Branch (pc', _) -> compile_block blocks debug_data code pc' state'
+      | Cond (_, (pc1, _), (pc2, _)) ->
+          compile_block blocks debug_data code pc1 state';
+          compile_block blocks debug_data code pc2 state'
+      | Poptrap (_, _) -> ()
+      | Switch (_, _) -> ()
+      | Pushtrap _ -> ()
+      | Raise _ | Return _ | Stop -> ())
 
 and compile infos pc state instrs =
   if debug_parser () then State.print state;
@@ -1208,7 +1224,9 @@ and compile infos pc state instrs =
         let state' = State.start_function state env 0 in
         let params, state' = State.make_stack nparams state' loc in
         if debug_parser () then Format.printf ") {@.";
-        let state' = State.clear_accu state' in
+        (* We can't use [Dummy] here as it breaks the stack shape
+           invariant. *)
+        let _dummy, state' = State.fresh_var state' No in
         compile_block infos.blocks infos.debug code addr state';
         if debug_parser () then Format.printf "}@.";
         let args = State.stack_vars state' in
@@ -1265,7 +1283,9 @@ and compile infos pc state instrs =
               let state' = State.start_function state env offset in
               let params, state' = State.make_stack nparams state' loc in
               if debug_parser () then Format.printf ") {@.";
-              let state' = State.clear_accu state' in
+              (* We can't use [Dummy] here as it breaks the stack
+                 shape invariant. *)
+              let _dummy, state' = State.fresh_var state' No in
               compile_block infos.blocks infos.debug code addr state';
               if debug_parser () then Format.printf "}@.";
               let args = State.stack_vars state' in
@@ -1677,7 +1697,7 @@ and compile infos pc state instrs =
             let isint_branch = pc + 1 in
             let isblock_branch = pc + 2 in
             let () =
-              tagged_blocks := Addr.Set.add isint_branch !tagged_blocks;
+              tagged_blocks := Addr.Map.add isint_branch state !tagged_blocks;
               let i_state = State.start_block isint_branch state in
               let i_args = State.stack_vars i_state in
               compiled_blocks :=
@@ -1687,7 +1707,7 @@ and compile infos pc state instrs =
                   !compiled_blocks
             in
             let () =
-              tagged_blocks := Addr.Set.add isblock_branch !tagged_blocks;
+              tagged_blocks := Addr.Map.add isblock_branch state !tagged_blocks;
               let x_tag = Var.fresh () in
               let b_state = State.start_block isblock_branch state in
               let b_args = State.stack_vars b_state in
@@ -1723,7 +1743,7 @@ and compile infos pc state instrs =
         let handler_addr = pc + 1 + gets code (pc + 1) in
         let x, handler_state = State.fresh_var handler_ctx_state loc in
 
-        tagged_blocks := Addr.Set.add interm_addr !tagged_blocks;
+        tagged_blocks := Addr.Map.add interm_addr state !tagged_blocks;
         compiled_blocks :=
           Addr.Map.add
             interm_addr
@@ -2492,7 +2512,7 @@ let parse_bytecode code globals debug_data =
   in
   pushpop := Addr.Map.empty;
   compiled_blocks := Addr.Map.empty;
-  tagged_blocks := Addr.Set.empty;
+  tagged_blocks := Addr.Map.empty;
   p
 
 (* HACK - override module *)

compiler/lib/phisimpl.ml

@@ -42,7 +42,8 @@ let rec arg_deps vars deps defs params args =
       add_dep deps x y;
       add_def vars defs x y;
       arg_deps vars deps defs params args
-  | _ -> ()
+  | [], [] -> ()
+  | _ -> assert false
 
 let cont_deps blocks vars deps defs (pc, args) =
   let block = Addr.Map.find pc blocks in

compiler/tests-compiler/effects_continuations.ml

@@ -103,43 +103,43 @@ let%expect_test "test-compiler/lib-effects/test1.ml" =
     {|
 
     function exceptions(s, cont){
-     try{var _B_ = runtime.caml_int_of_string(s), n = _B_;}
-     catch(_F_){
-      var _u_ = caml_wrap_exception(_F_);
+     try{var _z_ = runtime.caml_int_of_string(s), n = _z_;}
+     catch(_D_){
+      var _u_ = caml_wrap_exception(_D_);
       if(_u_[1] !== Stdlib[7]){
        var raise$1 = caml_pop_trap();
        return raise$1(caml_maybe_attach_backtrace(_u_, 0));
       }
-      var n = 0, _v_ = 0;
+      var n = 0;
      }
      try{
       if(caml_string_equal(s, cst$0))
        throw caml_maybe_attach_backtrace(Stdlib[8], 1);
-      var _A_ = 7, m = _A_;
+      var _y_ = 7, m = _y_;
      }
-     catch(_E_){
-      var _w_ = caml_wrap_exception(_E_);
-      if(_w_ !== Stdlib[8]){
+     catch(_C_){
+      var _v_ = caml_wrap_exception(_C_);
+      if(_v_ !== Stdlib[8]){
        var raise$0 = caml_pop_trap();
-       return raise$0(caml_maybe_attach_backtrace(_w_, 0));
+       return raise$0(caml_maybe_attach_backtrace(_v_, 0));
       }
-      var m = 0, _x_ = 0;
+      var m = 0;
      }
      runtime.caml_push_trap
-      (function(_D_){
-        if(_D_ === Stdlib[8]) return cont(0);
+      (function(_B_){
+        if(_B_ === Stdlib[8]) return cont(0);
         var raise = caml_pop_trap();
-        return raise(caml_maybe_attach_backtrace(_D_, 0));
+        return raise(caml_maybe_attach_backtrace(_B_, 0));
        });
      if(caml_string_equal(s, cst)){
-      var _y_ = Stdlib[8], raise = caml_pop_trap();
-      return raise(caml_maybe_attach_backtrace(_y_, 1));
+      var _w_ = Stdlib[8], raise = caml_pop_trap();
+      return raise(caml_maybe_attach_backtrace(_w_, 1));
      }
-     var _z_ = Stdlib[79];
+     var _x_ = Stdlib[79];
      return caml_cps_call2
-             (_z_,
+             (_x_,
               cst_toto,
-              function(_C_){caml_pop_trap(); return cont([0, [0, _C_, n, m]]);});
+              function(_A_){caml_pop_trap(); return cont([0, [0, _A_, n, m]]);});
     }
     //end
     function cond1(b, cont){

compiler/tests-compiler/effects_exceptions.ml

@@ -57,43 +57,43 @@ let%expect_test "test-compiler/lib-effects/test1.ml" =
     {|
 
     function exceptions(s, cont){
-     try{var _p_ = runtime.caml_int_of_string(s), n = _p_;}
-     catch(_t_){
-      var _i_ = caml_wrap_exception(_t_);
+     try{var _n_ = runtime.caml_int_of_string(s), n = _n_;}
+     catch(_r_){
+      var _i_ = caml_wrap_exception(_r_);
       if(_i_[1] !== Stdlib[7]){
        var raise$1 = caml_pop_trap();
        return raise$1(caml_maybe_attach_backtrace(_i_, 0));
       }
-      var n = 0, _j_ = 0;
+      var n = 0;
      }
      try{
       if(caml_string_equal(s, cst$0))
        throw caml_maybe_attach_backtrace(Stdlib[8], 1);
-      var _o_ = 7, m = _o_;
+      var _m_ = 7, m = _m_;
      }
-     catch(_s_){
-      var _k_ = caml_wrap_exception(_s_);
-      if(_k_ !== Stdlib[8]){
+     catch(_q_){
+      var _j_ = caml_wrap_exception(_q_);
+      if(_j_ !== Stdlib[8]){
        var raise$0 = caml_pop_trap();
-       return raise$0(caml_maybe_attach_backtrace(_k_, 0));
+       return raise$0(caml_maybe_attach_backtrace(_j_, 0));
       }
-      var m = 0, _l_ = 0;
+      var m = 0;
      }
      caml_push_trap
-      (function(_r_){
-        if(_r_ === Stdlib[8]) return cont(0);
+      (function(_p_){
+        if(_p_ === Stdlib[8]) return cont(0);
         var raise = caml_pop_trap();
-        return raise(caml_maybe_attach_backtrace(_r_, 0));
+        return raise(caml_maybe_attach_backtrace(_p_, 0));
        });
      if(caml_string_equal(s, cst)){
-      var _m_ = Stdlib[8], raise = caml_pop_trap();
-      return raise(caml_maybe_attach_backtrace(_m_, 1));
+      var _k_ = Stdlib[8], raise = caml_pop_trap();
+      return raise(caml_maybe_attach_backtrace(_k_, 1));
      }
-     var _n_ = Stdlib[79];
+     var _l_ = Stdlib[79];
      return caml_cps_call2
-             (_n_,
+             (_l_,
               cst_toto,
-              function(_q_){caml_pop_trap(); return cont([0, [0, _q_, n, m]]);});
+              function(_o_){caml_pop_trap(); return cont([0, [0, _o_, n, m]]);});
     }
     //end |}];
   print_fun_decl code (Some "handler_is_loop");

compiler/tests-compiler/inlining.ml

@@ -24,11 +24,19 @@ let%expect_test "inline recursive function" =
     let rec f () = f ()
     and g () = f ()
   |} in
+  print_fun_decl program (Some "f$0");
   print_fun_decl program (Some "f");
   print_fun_decl program (Some "g");
   [%expect
     {|
-    function f(param){for(;;) ;}
+    function f$0(counter, param){
+     var _a_ = 0;
+     if(counter >= 50) return caml_trampoline_return(f$0, [0, _a_]);
+     var counter$0 = counter + 1 | 0;
+     return f$0(counter$0, _a_);
+    }
+    //end
+    function f(param){return caml_trampoline(f$0(0, param));}
     //end
     function g(param){return f(0);}
     //end |}]