More assertions for local alloc and letrec (#1081)

stedolan · web-flow · commit b0e15275421e · 2023-02-21T15:03:30.000Z
diff --git a/backend/cmmgen.ml b/backend/cmmgen.ml
@@ -148,18 +148,22 @@ let get_field env ptr n dbg =
   get_field_gen mut ptr n dbg
 
 type rhs_kind =
-  | RHS_block of int
-  | RHS_infix of { blocksize : int; offset : int }
-  | RHS_floatblock of int
+  | RHS_block of Lambda.alloc_mode * int
+  | RHS_infix of { blocksize : int; offset : int; blockmode: Lambda.alloc_mode }
+  | RHS_floatblock of Lambda.alloc_mode * int
   | RHS_nonrec
 ;;
 
 let rec expr_size env = function
   | Uvar id ->
       begin try V.find_same id env with Not_found -> RHS_nonrec end
   | Uclosure { functions ; not_scanned_slots ; scanned_slots } ->
-      RHS_block (fundecls_size functions + List.length not_scanned_slots
-        + List.length scanned_slots)
+      (* should all have the same mode *)
+      let fn_mode = (List.hd functions).mode in
+      List.iter (fun f -> assert (Lambda.eq_mode fn_mode f.mode)) functions;
+      RHS_block (fn_mode,
+                 fundecls_size functions + List.length not_scanned_slots
+                 + List.length scanned_slots)
   | Ulet(_str, _kind, id, exp, body) ->
       expr_size (V.add (VP.var id) (expr_size env exp) env) body
   | Uletrec(bindings, body) ->
@@ -169,24 +173,24 @@ let rec expr_size env = function
           bindings env
       in
       expr_size env body
-  | Uprim(Pmakeblock _, args, _) ->
-      RHS_block (List.length args)
-  | Uprim(Pmakearray((Paddrarray | Pintarray), _, _), args, _) ->
-      RHS_block (List.length args)
-  | Uprim(Pmakearray(Pfloatarray, _, _), args, _) ->
-      RHS_floatblock (List.length args)
-  | Uprim(Pmakearray(Pgenarray, _, _), _, _) ->
+  | Uprim(Pmakeblock (_, _, _, mode), args, _) ->
+      RHS_block (mode, List.length args)
+  | Uprim(Pmakearray((Paddrarray | Pintarray), _, mode), args, _) ->
+      RHS_block (mode, List.length args)
+  | Uprim(Pmakearray(Pfloatarray, _, mode), args, _) ->
+      RHS_floatblock (mode, List.length args)
+  | Uprim(Pmakearray(Pgenarray, _, _mode), _, _) ->
      (* Pgenarray is excluded from recursive bindings by the
         check in Translcore.check_recursive_lambda *)
      RHS_nonrec
   | Uprim (Pduprecord ((Record_regular | Record_inlined _), sz), _, _) ->
-      RHS_block sz
+      RHS_block (Lambda.alloc_heap, sz)
   | Uprim (Pduprecord (Record_unboxed _, _), _, _) ->
       assert false
   | Uprim (Pduprecord (Record_extension _, sz), _, _) ->
-      RHS_block (sz + 1)
+      RHS_block (Lambda.alloc_heap, sz + 1)
   | Uprim (Pduprecord (Record_float, sz), _, _) ->
-      RHS_floatblock sz
+      RHS_floatblock (Lambda.alloc_heap, sz)
   | Uprim (Pccall { prim_name; _ }, closure::_, _)
         when prim_name = "caml_check_value_is_closure" ->
       (* Used for "-clambda-checks". *)
@@ -195,7 +199,8 @@ let rec expr_size env = function
       expr_size env exp'
   | Uoffset (exp, offset) ->
       (match expr_size env exp with
-      | RHS_block blocksize -> RHS_infix { blocksize; offset }
+      | RHS_block (blockmode, blocksize) ->
+         RHS_infix { blocksize; offset; blockmode }
       | RHS_nonrec -> RHS_nonrec
       | _ -> assert false)
   | Uregion exp ->
@@ -1437,14 +1442,19 @@ and transl_letrec env bindings cont =
         args, dbg) in
   let rec init_blocks = function
     | [] -> fill_nonrec bsz
-    | (id, _exp, RHS_block sz) :: rem ->
+    | (_, _,
+       (RHS_block (Alloc_local, _) |
+        RHS_infix {blockmode=Alloc_local; _} |
+        RHS_floatblock (Alloc_local, _))) :: _ ->
+      Misc.fatal_error "Invalid stack allocation found"
+    | (id, _exp, RHS_block (Alloc_heap, sz)) :: rem ->
         Clet(id, op_alloc "caml_alloc_dummy" [int_const dbg sz],
           init_blocks rem)
-    | (id, _exp, RHS_infix { blocksize; offset}) :: rem ->
+    | (id, _exp, RHS_infix { blocksize; offset; blockmode=Alloc_heap }) :: rem ->
         Clet(id, op_alloc "caml_alloc_dummy_infix"
              [int_const dbg blocksize; int_const dbg offset],
              init_blocks rem)
-    | (id, _exp, RHS_floatblock sz) :: rem ->
+    | (id, _exp, RHS_floatblock (Alloc_heap, sz)) :: rem ->
         Clet(id, op_alloc "caml_alloc_dummy_float" [int_const dbg sz],
           init_blocks rem)
     | (id, _exp, RHS_nonrec) :: rem ->
diff --git a/middle_end/flambda2/from_lambda/dissect_letrec.ml b/middle_end/flambda2/from_lambda/dissect_letrec.ml
@@ -187,6 +187,12 @@ let is_simple (lam : Lambda.lambda) =
   match lam with Lvar _ | Lconst _ -> true | _ -> false
   [@@ocaml.warning "-fragile-match"]
 
+let assert_not_local ~lam : Lambda.alloc_mode -> unit = function
+  | Alloc_heap -> ()
+  | Alloc_local ->
+    Misc.fatal_errorf "Invalid stack allocation found in %a" Printlambda.lambda
+      lam
+
 let dead_code lam letrec =
   (* Some cases generate code without effects, and bound to nothing. We use this
      function to insert it as [Lsequence] in [effects], for documentation. It
@@ -203,6 +209,7 @@ let rec prepare_letrec (recursive_set : Ident.Set.t)
       (lam : Lambda.lambda) (letrec : letrec) =
   match lam with
   | Lfunction funct -> (
+    assert_not_local ~lam funct.mode;
     match current_let with
     | Some current_let when Ident.Set.mem current_let.ident recursive_set ->
       { letrec with functions = (current_let.ident, funct) :: letrec.functions }
@@ -221,6 +228,10 @@ let rec prepare_letrec (recursive_set : Ident.Set.t)
     | None -> dead_code lam letrec)
   | Lprim (((Pmakeblock _ | Pmakearray _ | Pduprecord _) as prim), args, dbg)
     when not (List.for_all is_simple args) ->
+    (match prim with
+    | Pmakeblock (_, _, _, mode) | Pmakearray (_, _, mode) ->
+      assert_not_local ~lam mode
+    | _ -> ());
     (* If there are some non-trivial expressions as arguments, we first extract
        the arguments (to let-bound variables) before deconstructing. Arguments
        could contain side effects and other blocks declarations. *)
@@ -246,15 +257,16 @@ let rec prepare_letrec (recursive_set : Ident.Set.t)
         defs
     in
     prepare_letrec recursive_set current_let lam letrec
-  | Lprim (Pmakeblock _, args, _)
-  | Lprim (Pmakearray ((Paddrarray | Pintarray), _, _), args, _) -> (
+  | Lprim (Pmakeblock (_, _, _, mode), args, _)
+  | Lprim (Pmakearray ((Paddrarray | Pintarray), _, mode), args, _) -> (
+    assert_not_local ~lam mode;
     match current_let with
     | Some cl -> build_block cl (List.length args) (Normal 0) lam letrec
-    | None ->
-      dead_code lam letrec
-      (* We know that [args] are all "simple" at this point, so no effects *))
-  | Lprim (Pmakearray (Pfloatarray, _, _), args, _)
-  | Lprim (Pmakefloatblock _, args, _) -> (
+    | None -> dead_code lam letrec
+    (* We know that [args] are all "simple" at this point, so no effects *))
+  | Lprim (Pmakearray (Pfloatarray, _, mode), args, _)
+  | Lprim (Pmakefloatblock (_, mode), args, _) -> (
+    assert_not_local ~lam mode;
     match current_let with
     | Some cl -> build_block cl (List.length args) Boxed_float lam letrec
     | None -> dead_code lam letrec)
diff --git a/ocaml/asmcomp/cmmgen.ml b/ocaml/asmcomp/cmmgen.ml
@@ -122,18 +122,22 @@ let get_field env ptr n dbg =
   get_field_gen mut ptr n dbg
 
 type rhs_kind =
-  | RHS_block of int
-  | RHS_infix of { blocksize : int; offset : int }
-  | RHS_floatblock of int
+  | RHS_block of Lambda.alloc_mode * int
+  | RHS_infix of { blocksize : int; offset : int; blockmode: Lambda.alloc_mode }
+  | RHS_floatblock of Lambda.alloc_mode * int
   | RHS_nonrec
 ;;
 
 let rec expr_size env = function
   | Uvar id ->
       begin try V.find_same id env with Not_found -> RHS_nonrec end
   | Uclosure { functions ; not_scanned_slots ; scanned_slots } ->
-      RHS_block (fundecls_size functions + List.length not_scanned_slots
-        + List.length scanned_slots)
+      (* should all have the same mode *)
+      let fn_mode = (List.hd functions).mode in
+      List.iter (fun f -> assert (Lambda.eq_mode fn_mode f.mode)) functions;
+      RHS_block (fn_mode,
+                 fundecls_size functions + List.length not_scanned_slots
+                 + List.length scanned_slots)
   | Ulet(_str, _kind, id, exp, body) ->
       expr_size (V.add (VP.var id) (expr_size env exp) env) body
   | Uletrec(bindings, body) ->
@@ -143,24 +147,24 @@ let rec expr_size env = function
           bindings env
       in
       expr_size env body
-  | Uprim(Pmakeblock _, args, _) ->
-      RHS_block (List.length args)
-  | Uprim(Pmakearray((Paddrarray | Pintarray), _, _), args, _) ->
-      RHS_block (List.length args)
-  | Uprim(Pmakearray(Pfloatarray, _, _), args, _) ->
-      RHS_floatblock (List.length args)
-  | Uprim(Pmakearray(Pgenarray, _, _), _, _) ->
+  | Uprim(Pmakeblock (_, _, _, mode), args, _) ->
+      RHS_block (mode, List.length args)
+  | Uprim(Pmakearray((Paddrarray | Pintarray), _, mode), args, _) ->
+      RHS_block (mode, List.length args)
+  | Uprim(Pmakearray(Pfloatarray, _, mode), args, _) ->
+      RHS_floatblock (mode, List.length args)
+  | Uprim(Pmakearray(Pgenarray, _, _mode), _, _) ->
      (* Pgenarray is excluded from recursive bindings by the
         check in Translcore.check_recursive_lambda *)
      RHS_nonrec
   | Uprim (Pduprecord ((Record_regular | Record_inlined _), sz), _, _) ->
-      RHS_block sz
+      RHS_block (Lambda.alloc_heap, sz)
   | Uprim (Pduprecord (Record_unboxed _, _), _, _) ->
       assert false
   | Uprim (Pduprecord (Record_extension _, sz), _, _) ->
-      RHS_block (sz + 1)
+      RHS_block (Lambda.alloc_heap, sz + 1)
   | Uprim (Pduprecord (Record_float, sz), _, _) ->
-      RHS_floatblock sz
+      RHS_floatblock (Lambda.alloc_heap, sz)
   | Uprim (Pccall { prim_name; _ }, closure::_, _)
         when prim_name = "caml_check_value_is_closure" ->
       (* Used for "-clambda-checks". *)
@@ -169,7 +173,8 @@ let rec expr_size env = function
       expr_size env exp'
   | Uoffset (exp, offset) ->
       (match expr_size env exp with
-      | RHS_block blocksize -> RHS_infix { blocksize; offset }
+      | RHS_block (blockmode, blocksize) ->
+         RHS_infix { blocksize; offset; blockmode }
       | RHS_nonrec -> RHS_nonrec
       | _ -> assert false)
   | Uregion exp ->
@@ -1388,14 +1393,19 @@ and transl_letrec env bindings cont =
     Cop(Cextcall(prim, typ_val, [], true), args, dbg) in
   let rec init_blocks = function
     | [] -> fill_nonrec bsz
-    | (id, _exp, RHS_block sz) :: rem ->
+    | (_, _,
+       (RHS_block (Alloc_local, _) |
+        RHS_infix {blockmode=Alloc_local; _} |
+        RHS_floatblock (Alloc_local, _))) :: _ ->
+      Misc.fatal_error "Invalid stack allocation found"
+    | (id, _exp, RHS_block (Alloc_heap, sz)) :: rem ->
         Clet(id, op_alloc "caml_alloc_dummy" [int_const dbg sz],
           init_blocks rem)
-    | (id, _exp, RHS_infix { blocksize; offset}) :: rem ->
+    | (id, _exp, RHS_infix { blocksize; offset; blockmode=Alloc_heap }) :: rem ->
         Clet(id, op_alloc "caml_alloc_dummy_infix"
              [int_const dbg blocksize; int_const dbg offset],
              init_blocks rem)
-    | (id, _exp, RHS_floatblock sz) :: rem ->
+    | (id, _exp, RHS_floatblock (Alloc_heap, sz)) :: rem ->
         Clet(id, op_alloc "caml_alloc_dummy_float" [int_const dbg sz],
           init_blocks rem)
     | (id, _exp, RHS_nonrec) :: rem ->
