Static check for noalloc: ignore allocations post-dominated by a raise

backend/checkmach.ml

@@ -89,6 +89,14 @@ module type Spec = sig
   val check_specific : Arch.specific_operation -> bool
 
   val annotation : Cmm.property
+
+  (** [ignore_postdominated_by_raise = true] Do not perform checks on paths that
+      must lead to a raise, as indicated by a [raise] with a backtrace, i.e.,
+      paths postdominated by [raise_notrace] are not ignored.
+
+      In other words, when a function returns normally, it is guaranteed to pass
+      the check, but not upon exceptional return. *)
+  val ignore_postdominated_by_raise : bool
 end
 (* CR-someday gyorsh: We may also want annotations on call sites, not only on
    functions. *)
@@ -274,7 +282,18 @@ end = struct
   type t =
     { ppf : Format.formatter;
       fun_name : string;
-      mutable unresolved_dependencies : bool
+      mutable unresolved_dependencies : bool;
+      mutable raise_after_handlers : int list
+          (** all Iexit indexes for which the corresponding Icatch handler leads
+              to a [raise] with backtrace. Conservative: recursive catch
+              handlers are never added. *)
+    }
+
+  let create ppf fun_name =
+    { ppf;
+      fun_name;
+      unresolved_dependencies = false;
+      raise_after_handlers = []
     }
 
   let report t ~msg ~desc dbg =
@@ -368,7 +387,22 @@ end = struct
     | Iswitch (_index, cases) ->
       let raise_after = check_instr_exn t i.next raise_after in
       Array.for_all (fun case -> check_instr_exn t case raise_after) cases
-    | Icatch (_rec, _ts, handlers, body) ->
+    | Icatch (Nonrecursive, _ts, handlers, body) ->
+      let raise_after = check_instr_exn t i.next raise_after in
+      (* order is important: first analyze the handlers to get accurate
+         information for Iexit that may appear in the body. *)
+      let raise_after_all_handlers =
+        List.for_all
+          (fun (n, _ts, handler) ->
+            let res = check_instr_exn t handler raise_after in
+            (* no duplicates in this list because handler indexes are unique *)
+            if res then t.raise_after_handlers <- n :: t.raise_after_handlers;
+            res)
+          handlers
+      in
+      let raise_after_body = check_instr_exn t body raise_after in
+      raise_after_body && raise_after_all_handlers
+    | Icatch (Recursive, _ts, handlers, body) ->
       (* conservative *)
       let raise_after = check_instr_exn t i.next raise_after in
       let _ = check_instr_exn t body raise_after in
@@ -382,9 +416,16 @@ end = struct
       let _ = check_instr_exn t body false in
       let _ = check_instr_exn t handler false in
       false
-    | Iraise _ -> false
+    | Iraise Raise_notrace ->
+      (* The intended use of [S.ignore_postdominated_by_raise] is to ignore
+         allocation on error paths, whereas [raise_notrace] is typically used
+         for control flow, not for indicating an error. *)
+      false
+    | Iraise (Raise_reraise | Raise_regular) -> S.ignore_postdominated_by_raise
     | Ireturn _ -> false
-    | Iexit _ -> false
+    | Iexit (n, _) ->
+      (* conservative *)
+      List.exists (Int.equal n) t.raise_after_handlers
 
   let debug t expected =
     match Unit_info.get_value unit_info t.fun_name with
@@ -402,7 +443,7 @@ end = struct
     then
       Profile.record_call ~accumulate:true ("check " ^ S.name) (fun () ->
           let fun_name = f.fun_name in
-          let t = { ppf; fun_name; unresolved_dependencies = false } in
+          let t = create ppf fun_name in
           Unit_info.in_current_unit unit_info fun_name f.fun_dbg;
           if List.mem (Cmm.Assume S.annotation) f.fun_codegen_options
           then (
@@ -422,7 +463,7 @@ end = struct
 end
 
 module Spec_alloc : Spec = struct
-  let name = "noalloc"
+  let name = "noalloc_strict"
 
   let enabled () = !Flambda_backend_flags.alloc_check
 
@@ -438,6 +479,30 @@ module Spec_alloc : Spec = struct
   let check_specific s = not (Arch.operation_allocates s)
 
   let annotation = Cmm.Noalloc
+
+  let ignore_postdominated_by_raise = false
+end
+
+module Spec_alloc_exn : Spec = struct
+  let name = "noalloc"
+
+  let enabled () = !Flambda_backend_flags.alloc_check
+
+  let check_callee s (checks : Cmx_format.checks) =
+    String.Set.mem s checks.ui_noalloc_exn_functions
+    || String.Set.mem s checks.ui_noalloc_functions
+
+  let add_callee s (checks : Cmx_format.checks) =
+    if not (check_callee s checks)
+    then
+      checks.ui_noalloc_exn_functions
+        <- String.Set.add s checks.ui_noalloc_exn_functions
+
+  let check_specific s = not (Arch.operation_allocates s)
+
+  let annotation = Cmm.Noalloc_exn
+
+  let ignore_postdominated_by_raise = true
 end
 
 (***************************************************************************
@@ -449,17 +514,32 @@ end
  *****************************************************************************)
 module Check_alloc = Analysis (Spec_alloc)
 
+(* Same as [Check_alloc] except no restrictions on instructions post-dominated
+   by a raise. *)
+module Check_alloc_exn = Analysis (Spec_alloc_exn)
+
 let fundecl ppf_dump fd =
+  (* We can check both properties at the same time in one pass, or even as part
+     of another pass, such as dead code elimination. For simplicity, it's a
+     separate pass for each at the moment. *)
   Check_alloc.fundecl ppf_dump fd;
+  Check_alloc_exn.fundecl ppf_dump fd;
   fd
 
 let reset_unit_info () =
   Check_alloc.reset_unit_info ();
+  Check_alloc_exn.reset_unit_info ();
   ()
 
+(* If a function is noalloc, then it is also noalloc_exn but not vice versa: it
+   is possible the noalloc_exn check succeeds, but noalloc check fails.
+
+   We use this fact to represent the result of the check using disjoint sets. *)
 let record_unit_info ppf_dump =
   let checks = (Compilenv.current_unit_infos ()).ui_checks in
+  (* Order is important here. *)
   Check_alloc.record_unit_info ppf_dump checks;
+  Check_alloc_exn.record_unit_info ppf_dump checks;
   Compilenv.cache_checks checks
 
 (* Error report *)

backend/cmm.ml

@@ -246,6 +246,7 @@ type expression =
 
 type property =
   | Noalloc
+  | Noalloc_exn
 
 type codegen_option =
   | Reduce_code_size

backend/cmm.mli

@@ -255,6 +255,7 @@ type expression =
 
 type property =
   | Noalloc
+  | Noalloc_exn
 
 type codegen_option =
   | Reduce_code_size

backend/cmm_helpers.ml

@@ -4172,6 +4172,7 @@ let cmm_arith_size (e : Cmm.expression) =
 
 let transl_property : Lambda.property -> Cmm.property = function
   | Noalloc -> Noalloc
+  | Noalloc_exn -> Noalloc_exn
 
 let transl_attrib : Lambda.check_attribute -> Cmm.codegen_option list = function
   | Default_check -> []

backend/printcmm.ml

@@ -340,7 +340,8 @@ and sequence ppf = function
 and expression ppf e = fprintf ppf "%a" expr e
 
 let property : Cmm.property -> string = function
-    | Noalloc -> "noalloc"
+  | Noalloc -> "noalloc"
+  | Noalloc_exn -> "noalloc_exn"
 
 let codegen_option = function
   | Reduce_code_size -> "reduce_code_size"

file_formats/cmx_format.mli

@@ -50,6 +50,10 @@ type checks =
     (* CR gyorsh: refactor to use lists. *)
     mutable ui_noalloc_functions: Misc.Stdlib.String.Set.t;
     (* Functions without allocations and indirect calls *)
+
+    mutable ui_noalloc_exn_functions: Misc.Stdlib.String.Set.t;
+    (* Same as noalloc, except no restrictions on instructions
+       (allocations and indirect calls) post-dominated by a raise. *)
   }
 
 type unit_infos =

middle_end/compilenv.ml

@@ -90,16 +90,20 @@ end = struct
   let create () =
     {
       ui_noalloc_functions = String.Set.empty;
+      ui_noalloc_exn_functions = String.Set.empty;
     }
 
   let reset t =
-    t.ui_noalloc_functions <- String.Set.empty
+    t.ui_noalloc_functions <- String.Set.empty;
+    t.ui_noalloc_exn_functions <- String.Set.empty
 
   let merge src ~into:dst =
     if !Flambda_backend_flags.alloc_check
     then (
       dst.ui_noalloc_functions
-        <- String.Set.union dst.ui_noalloc_functions src.ui_noalloc_functions)
+      <- String.Set.union dst.ui_noalloc_functions src.ui_noalloc_functions;
+      dst.ui_noalloc_exn_functions
+      <- String.Set.union dst.ui_noalloc_exn_functions src.ui_noalloc_exn_functions)
 end
 
 let cached_checks : Cmx_format.checks = Checks.create ()

middle_end/flambda2/terms/check_attribute.ml

@@ -11,13 +11,21 @@
 (**************************************************************************)
 
 module Property = struct
-  type t = Noalloc
+  type t =
+    | Noalloc
+    | Noalloc_exn
 
-  let to_string = function Noalloc -> "noalloc"
+  let to_string = function Noalloc -> "noalloc" | Noalloc_exn -> "noalloc_exn"
 
-  let equal x y = match x, y with Noalloc, Noalloc -> true
+  let equal x y =
+    match x, y with
+    | Noalloc, Noalloc -> true
+    | Noalloc_exn, Noalloc_exn -> true
+    | (Noalloc | Noalloc_exn), _ -> false
 
-  let from_lambda : Lambda.property -> t = function Noalloc -> Noalloc
+  let from_lambda : Lambda.property -> t = function
+    | Noalloc -> Noalloc
+    | Noalloc_exn -> Noalloc_exn
 end
 
 type t =

middle_end/flambda2/terms/check_attribute.mli

@@ -11,7 +11,9 @@
 (**************************************************************************)
 (** Annotations on function declaration (not call sites) *)
 module Property : sig
-  type t = Noalloc
+  type t =
+    | Noalloc
+    | Noalloc_exn
 end
 
 type t =

middle_end/flambda2/to_cmm/to_cmm_set_of_closures.ml

@@ -251,6 +251,7 @@ end)
 
 let transl_property : Check_attribute.Property.t -> Cmm.property = function
   | Noalloc -> Noalloc
+  | Noalloc_exn -> Noalloc_exn
 
 let transl_check_attrib : Check_attribute.t -> Cmm.codegen_option list =
   function

ocaml/lambda/lambda.ml

@@ -374,6 +374,7 @@ type local_attribute =
 
 type property =
   | Noalloc
+  | Noalloc_exn
 
 type check_attribute =
   | Default_check

ocaml/lambda/lambda.mli

@@ -287,6 +287,7 @@ type local_attribute =
 
 type property =
   | Noalloc
+  | Noalloc_exn
 
 type check_attribute =
   | Default_check

ocaml/lambda/printlambda.ml

@@ -553,6 +553,7 @@ let name_of_primitive = function
 let check_attribute ppf check =
   let check_property = function
     | Noalloc -> "noalloc"
+    | Noalloc_exn -> "noalloc"
   in
   match check with
   | Default_check -> ()

ocaml/lambda/translattribute.ml

@@ -40,7 +40,8 @@ let is_local_attribute = function
 
 let is_property_attribute p a =
   match p, a with
-  | Noalloc, {txt=("noalloc"|"ocaml.noalloc")} -> true
+  | Noalloc, {txt=("noalloc_strict"|"ocaml.noalloc_strict")} -> true
+  | Noalloc_exn, {txt=("noalloc"|"ocaml.noalloc")} -> true
   | _ -> false
 
 let find_attribute p attributes =
@@ -235,7 +236,7 @@ let get_check_attribute l =
     match get_property_attribute l p with
     | Default_check -> None
     | a -> Some a)
-    [Noalloc]
+    [Noalloc; Noalloc_exn]
 
 let check_local_inline loc attr =
   match attr.local, attr.inline with
@@ -300,7 +301,8 @@ let add_local_attribute expr loc attributes =
 
 let add_check_attribute expr loc attributes =
   let to_string = function
-    | Noalloc -> "noalloc"
+    | Noalloc -> "noalloc_strict"
+    | Noalloc_exn -> "noalloc"
   in
   let to_string = function
     | Assert p -> to_string p
@@ -312,7 +314,8 @@ let add_check_attribute expr loc attributes =
   | Lfunction({ attr = { stub = false } as attr } as funct), [check] ->
       begin match attr.check with
       | Default_check -> ()
-      | Assert Noalloc | Assume Noalloc ->
+      | Assert (Noalloc | Noalloc_exn)
+      | Assume (Noalloc | Noalloc_exn) ->
           Location.prerr_warning loc
             (Warnings.Duplicated_attribute (to_string check))
       end;