More assertions for local alloc and letrec

Author: stedolan

backend/cmmgen.ml

@@ -154,10 +154,15 @@ type rhs_kind =
   | RHS_nonrec
 ;;
 
+let assert_not_local : Lambda.alloc_mode -> unit = function
+  | Alloc_heap -> ()
+  | Alloc_local -> Misc.fatal_error "Invalid stack allocation found"
+
 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 } ->
+      List.iter (fun f -> assert_not_local f.mode) functions;
       RHS_block (fundecls_size functions + List.length not_scanned_slots
         + List.length scanned_slots)
   | Ulet(_str, _kind, id, exp, body) ->
@@ -169,13 +174,17 @@ let rec expr_size env = function
           bindings env
       in
       expr_size env body
-  | Uprim(Pmakeblock _, args, _) ->
+  | Uprim(Pmakeblock (_, _, _, mode), args, _) ->
+      assert_not_local mode;
       RHS_block (List.length args)
-  | Uprim(Pmakearray((Paddrarray | Pintarray), _, _), args, _) ->
+  | Uprim(Pmakearray((Paddrarray | Pintarray), _, mode), args, _) ->
+      assert_not_local mode;
       RHS_block (List.length args)
-  | Uprim(Pmakearray(Pfloatarray, _, _), args, _) ->
+  | Uprim(Pmakearray(Pfloatarray, _, mode), args, _) ->
+      assert_not_local mode;
       RHS_floatblock (List.length args)
-  | Uprim(Pmakearray(Pgenarray, _, _), _, _) ->
+  | Uprim(Pmakearray(Pgenarray, _, mode), _, _) ->
+      assert_not_local mode;
      (* Pgenarray is excluded from recursive bindings by the
         check in Translcore.check_recursive_lambda *)
      RHS_nonrec

middle_end/flambda2/from_lambda/dissect_letrec.ml

@@ -187,6 +187,10 @@ let is_simple (lam : Lambda.lambda) =
   match lam with Lvar _ | Lconst _ -> true | _ -> false
   [@@ocaml.warning "-fragile-match"]
 
+let assert_not_local : Lambda.alloc_mode -> unit = function
+  | Alloc_heap -> ()
+  | Alloc_local -> Misc.fatal_error "Invalid stack allocation found"
+
 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 +207,7 @@ let rec prepare_letrec (recursive_set : Ident.Set.t)
       (lam : Lambda.lambda) (letrec : letrec) =
   match lam with
   | Lfunction funct -> (
+    assert_not_local 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 +226,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 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 +255,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 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 mode;
     match current_let with
     | Some cl -> build_block cl (List.length args) Boxed_float lam letrec
     | None -> dead_code lam letrec)