Zero alloc: assume error means never raises

backend/checkmach.ml

@@ -266,14 +266,16 @@ end = struct
     { strict : bool;  (** strict or relaxed? *)
       assume : bool;
       never_returns_normally : bool;
+      never_raises : bool;
       loc : Location.t
           (** Source location of the annotation, used for error messages. *)
     }
 
   let get_loc t = t.loc
 
-  let expected_value { strict; never_returns_normally; _ } =
-    Value.of_annotation ~strict ~never_returns_normally
+  let expected_value
+      { strict; never_returns_normally; assume = _; loc = _; never_raises } =
+    Value.of_annotation ~strict ~never_returns_normally ~never_raises
 
   let valid t v =
     (* Use Value.lessequal but ignore witnesses. *)
@@ -295,10 +297,23 @@ end = struct
         (fun (c : Cmm.codegen_option) ->
           match c with
           | Check { property; strict; loc } when property = spec ->
-            Some { strict; assume = false; never_returns_normally = false; loc }
-          | Assume { property; strict; never_returns_normally; loc }
+            Some
+              { strict;
+                assume = false;
+                never_returns_normally = false;
+                never_raises = false;
+                loc
+              }
+          | Assume
+              { property; strict; never_returns_normally; never_raises; loc }
             when property = spec ->
-            Some { strict; assume = true; never_returns_normally; loc }
+            Some
+              { strict;
+                assume = true;
+                never_returns_normally;
+                never_raises;
+                loc
+              }
           | Check _ | Assume _ | Reduce_code_size | No_CSE
           | Use_linscan_regalloc ->
             None)

backend/cmm.ml

@@ -319,6 +319,7 @@ type codegen_option =
   | No_CSE
   | Use_linscan_regalloc
   | Assume of { property: property; strict: bool; never_returns_normally: bool;
+                never_raises: bool;
                 loc: Location.t }
   | Check of { property: property; strict: bool; loc : Location.t; }
 

backend/cmm.mli

@@ -321,6 +321,7 @@ type codegen_option =
   | No_CSE
   | Use_linscan_regalloc
   | Assume of { property: property; strict: bool; never_returns_normally: bool;
+                never_raises: bool;
                 loc: Location.t }
   | Check of { property: property; strict: bool; loc: Location.t }
 

middle_end/flambda2/terms/check_attribute.ml

@@ -26,6 +26,7 @@ type t =
       { property : Property.t;
         strict : bool;
         never_returns_normally : bool;
+        never_raises : bool;
         loc : Location.t
       }
   | Check of
@@ -37,10 +38,12 @@ type t =
 let print ppf t =
   match t with
   | Default_check -> ()
-  | Assume { property; strict; never_returns_normally; loc = _ } ->
-    Format.fprintf ppf "@[assume_%a%s%s@]" Property.print property
+  | Assume { property; strict; never_returns_normally; never_raises; loc = _ }
+    ->
+    Format.fprintf ppf "@[assume_%a%s%s%s@]" Property.print property
       (if strict then "_strict" else "")
       (if never_returns_normally then "_never_returns_normally" else "")
+      (if never_raises then "_never_raises" else "")
   | Check { property; strict; loc = _ } ->
     Format.fprintf ppf "@[assert_%a%s@]" Property.print property
       (if strict then "_strict" else "")
@@ -54,11 +57,14 @@ let from_lambda : Lambda.check_attribute -> Location.t -> t =
     then Check { property = Zero_alloc; strict = false; loc }
     else Default_check
   | Ignore_assert_all Zero_alloc -> Default_check
-  | Assume { property; strict; never_returns_normally; loc; arity = _ } ->
+  | Assume
+      { property; strict; never_returns_normally; never_raises; loc; arity = _ }
+    ->
     Assume
       { property = Property.from_lambda property;
         strict;
         never_returns_normally;
+        never_raises;
         loc
       }
   | Check { property; strict; opt; loc; arity = _ } ->
@@ -73,11 +79,21 @@ let equal x y =
       Check { property = p2; strict = s2; loc = loc2 } ) ->
     Property.equal p1 p2 && Bool.equal s1 s2 && Location.compare loc1 loc2 = 0
   | ( Assume
-        { property = p1; strict = s1; never_returns_normally = n1; loc = loc1 },
+        { property = p1;
+          strict = s1;
+          never_returns_normally = n1;
+          never_raises = r1;
+          loc = loc1
+        },
       Assume
-        { property = p2; strict = s2; never_returns_normally = n2; loc = loc2 }
-    ) ->
+        { property = p2;
+          strict = s2;
+          never_returns_normally = n2;
+          never_raises = r2;
+          loc = loc2
+        } ) ->
     Property.equal p1 p2 && Bool.equal s1 s2 && Bool.equal n1 n2
+    && Bool.equal r1 r2
     && Location.compare loc1 loc2 = 0
   | (Default_check | Check _ | Assume _), _ -> false
 

middle_end/flambda2/terms/check_attribute.mli

@@ -20,6 +20,7 @@ type t =
       { property : Property.t;
         strict : bool;
         never_returns_normally : bool;
+        never_raises : bool;
         loc : Location.t
       }
   | Check of

middle_end/flambda2/to_cmm/to_cmm_set_of_closures.ml

@@ -327,11 +327,12 @@ let transl_property : Check_attribute.Property.t -> Cmm.property = function
 let transl_check_attrib : Check_attribute.t -> Cmm.codegen_option list =
   function
   | Default_check -> []
-  | Assume { property; strict; never_returns_normally; loc } ->
+  | Assume { property; strict; never_returns_normally; never_raises; loc } ->
     [ Assume
         { property = transl_property property;
           strict;
           never_returns_normally;
+          never_raises;
           loc
         } ]
   | Check { property; strict; loc } ->

ocaml/lambda/lambda.ml

@@ -640,6 +640,7 @@ type check_attribute = Builtin_attributes.check_attribute =
   | Assume of { property: property;
                 strict: bool;
                 never_returns_normally: bool;
+                never_raises: bool;
                 arity: int;
                 loc: Location.t;
               }

ocaml/lambda/lambda.mli

@@ -521,6 +521,7 @@ type check_attribute = Builtin_attributes.check_attribute =
   | Assume of { property: property;
                 strict: bool;
                 never_returns_normally: bool;
+                never_raises: bool;
                 arity: int;
                 loc: Location.t;
               }

ocaml/parsing/builtin_attributes.ml

@@ -691,6 +691,7 @@ type check_attribute =
   | Assume of { property: property;
                 strict: bool;
                 never_returns_normally: bool;
+                never_raises: bool;
                 arity: int;
                 loc: Location.t;
               }
@@ -825,6 +826,7 @@ let zero_alloc_lookup_table =
     (["assume"],
      fun arity loc ->
        Assume { property; strict = false; never_returns_normally = false;
+                never_raises = false;
                 arity; loc; });
     (["strict"],
      fun arity loc ->
@@ -838,18 +840,22 @@ let zero_alloc_lookup_table =
     (["assume"; "strict"],
      fun arity loc ->
        Assume { property; strict = true; never_returns_normally = false;
+                never_raises = false;
                 arity; loc; });
     (["assume"; "never_returns_normally"],
      fun arity loc ->
        Assume { property; strict = false; never_returns_normally = true;
+                never_raises = false;
                 arity; loc; });
     (["assume"; "never_returns_normally"; "strict"],
      fun arity loc ->
        Assume { property; strict = true; never_returns_normally = true;
+                never_raises = false;
                 arity; loc; });
     (["assume"; "error"],
      fun arity loc ->
-       Assume { property; strict = false; never_returns_normally = true;
+       Assume { property; strict = true; never_returns_normally = true;
+                never_raises = true;
                 arity; loc; });
     (["ignore"], fun _ _ -> Ignore_assert_all property)
   ]
@@ -926,8 +932,8 @@ let assume_zero_alloc ~is_check_allowed check : Zero_alloc_utils.Assume_info.t =
   match check with
   | Default_check -> Zero_alloc_utils.Assume_info.none
   | Ignore_assert_all Zero_alloc -> Zero_alloc_utils.Assume_info.none
-  | Assume { property=Zero_alloc; strict; never_returns_normally; } ->
-    Zero_alloc_utils.Assume_info.create ~strict ~never_returns_normally
+  | Assume { property=Zero_alloc; strict; never_returns_normally; never_raises; } ->
+    Zero_alloc_utils.Assume_info.create ~strict ~never_returns_normally ~never_raises
   | Check { property=Zero_alloc; loc; _ } ->
     if not is_check_allowed then begin
       let name = "zero_alloc" in

ocaml/parsing/builtin_attributes.mli

@@ -257,6 +257,11 @@ type check_attribute =
   | Assume of { property: property;
                 strict: bool;
                 never_returns_normally: bool;
+                never_raises: bool;
+                (* [never_raises=true] the function never returns
+                   via an exception. The function (directly or transitively)
+                   may raise exceptions that do not escape, i.e.,
+                   handled before the function returns. *)
                 arity: int;
                 loc: Location.t;
               }

ocaml/typing/includecore.ml

@@ -117,8 +117,10 @@ let zero_alloc za1 za2 =
     | Default_check | Ignore_assert_all _ -> ZA.Assume_info.Value.top ()
     | Check { strict; _ } ->
       ZA.Assume_info.Value.of_annotation ~strict ~never_returns_normally:false
-    | Assume { strict; never_returns_normally } ->
+        ~never_raises:false
+    | Assume { strict; never_returns_normally; never_raises; } ->
       ZA.Assume_info.Value.of_annotation ~strict ~never_returns_normally
+        ~never_raises
   in
   let v1 = abstract_value za1 in
   let v2 = abstract_value za2 in

ocaml/typing/typecore.ml

@@ -5087,6 +5087,7 @@ let zero_alloc_of_application ~num_args attrs funct =
           property = Zero_alloc;
           strict = c.strict;
           never_returns_normally = false;
+          never_raises = false;
           arity = c.arity;
           loc = c.loc
         }

ocaml/typing/typemod.ml

@@ -2816,7 +2816,8 @@ and type_structure ?(toplevel = None) funct_body anchor env sstr =
                 | Default_check | Ignore_assert_all _ -> Default_check
                 | Check _ -> zero_alloc
                 | Assume { property; strict; arity; loc;
-                           never_returns_normally = _ } ->
+                           never_returns_normally = _;
+                           never_raises = _} ->
                   Check { strict; property; arity; loc; opt = false }
               in
               let (first_loc, _, _) = List.hd id_info in

ocaml/utils/zero_alloc_utils.ml

@@ -137,8 +137,9 @@ module Make (Witnesses : WS) = struct
 
     let relaxed w = { nor = V.Safe; exn = V.Top w; div = V.Top w }
 
-    let of_annotation ~strict ~never_returns_normally =
+    let of_annotation ~strict ~never_returns_normally ~never_raises =
       let res = if strict then safe else relaxed Witnesses.empty in
+      let res = if never_raises then { res with exn = V.Bot } else res in
       if never_returns_normally then { res with nor = V.Bot } else res
 
     let print ~witnesses ppf { nor; exn; div } =
@@ -211,8 +212,8 @@ module Assume_info = struct
 
   let none = No_assume
 
-  let create ~strict ~never_returns_normally =
-    Assume (Value.of_annotation ~strict ~never_returns_normally)
+  let create ~strict ~never_returns_normally ~never_raises =
+    Assume (Value.of_annotation ~strict ~never_returns_normally ~never_raises)
 
   let get_value t = match t with No_assume -> None | Assume v -> Some v
 

ocaml/utils/zero_alloc_utils.mli

@@ -83,7 +83,8 @@ module Make (Witnesses : WS) : sig
     val relaxed : Witnesses.t -> t
 
     (** Constructs a value from a user annotation. The witness will be empty. *)
-    val of_annotation : strict:bool -> never_returns_normally:bool -> t
+    val of_annotation :
+      strict:bool -> never_returns_normally:bool -> never_raises:bool -> t
 
     val print : witnesses:bool -> Format.formatter -> t -> unit
 
@@ -113,7 +114,8 @@ module Assume_info : sig
 
   val none : t
 
-  val create : strict:bool -> never_returns_normally:bool -> t
+  val create :
+    strict:bool -> never_returns_normally:bool -> never_raises:bool -> t
 
   val compare : t -> t -> int
 

tests/backend/checkmach/dune.inc

@@ -959,3 +959,22 @@
  (enabled_if (= %{context_name} "main"))
  (deps test_assume_inlining.output test_assume_inlining.output.corrected)
  (action (diff test_assume_inlining.output test_assume_inlining.output.corrected)))
+
+(rule
+ (enabled_if (= %{context_name} "main"))
+ (targets test_assume_error.output.corrected)
+ (deps (:ml test_assume_error.ml) filter.sh)
+ (action
+   (with-outputs-to test_assume_error.output.corrected
+    (pipe-outputs
+    (with-accepted-exit-codes 2
+     (run %{bin:ocamlopt.opt} %{ml} -g -color never -error-style short -c
+          -zero-alloc-check default -checkmach-details-cutoff 20 -O3))
+    (run "./filter.sh")
+   ))))
+
+(rule
+ (alias   runtest)
+ (enabled_if (= %{context_name} "main"))
+ (deps test_assume_error.output test_assume_error.output.corrected)
+ (action (diff test_assume_error.output test_assume_error.output.corrected)))

tests/backend/checkmach/gen/gen_dune.ml

@@ -179,4 +179,6 @@ let () =
     ~exit_code:2 "test_signatures_separate_b";
   print_test_expected_output ~cutoff:default_cutoff
     ~extra_dep:None ~exit_code:2 "test_assume_inlining";
+  print_test_expected_output ~cutoff:default_cutoff
+    ~extra_dep:None ~exit_code:2 "test_assume_error";
   ()