Zero alloc annotation: assume a function never returns normally

backend/checkmach.ml

@@ -403,13 +403,16 @@ end = struct
   type t =
     { strict : bool;  (** strict or relaxed? *)
       assume : bool;
+      never_returns_normally : bool;
       loc : Location.t
           (** Source location of the annotation, used for error messages. *)
     }
 
   let get_loc t = t.loc
 
-  let expected_value t = if t.strict then Value.safe else Value.relaxed
+  let expected_value t =
+    let res = if t.strict then Value.safe else Value.relaxed in
+    if t.never_returns_normally then { res with nor = V.Bot } else res
 
   let is_assume t = t.assume
 
@@ -419,12 +422,15 @@ end = struct
       List.filter_map
         (fun (c : Cmm.codegen_option) ->
           match c with
-          | Check { property; strict; assume; loc } when property = spec ->
-            Some { strict; assume; loc }
+          | Check { property; strict; loc } when property = spec ->
+            Some { strict; assume = false; never_returns_normally = false; loc }
+          | Assume { property; strict; never_returns_normally; loc }
+            when property = spec ->
+            Some { strict; assume = true; never_returns_normally; loc }
           | Ignore_assert_all property when property = spec ->
             ignore_assert_all := true;
             None
-          | Ignore_assert_all _ | Check _ | Reduce_code_size | No_CSE
+          | Ignore_assert_all _ | Check _ | Assume _ | Reduce_code_size | No_CSE
           | Use_linscan_regalloc ->
             None)
         codegen_options
@@ -433,7 +439,12 @@ end = struct
     | [] ->
       if !Clflags.zero_alloc_check_assert_all && not !ignore_assert_all
       then
-        Some { strict = false; assume = false; loc = Debuginfo.to_location dbg }
+        Some
+          { strict = false;
+            assume = false;
+            never_returns_normally = false;
+            loc = Debuginfo.to_location dbg
+          }
       else None
     | [p] -> Some p
     | _ :: _ ->

backend/cmm.ml

@@ -302,8 +302,9 @@ type codegen_option =
   | No_CSE
   | Use_linscan_regalloc
   | Ignore_assert_all of property
-  | Check of { property: property; strict: bool; assume: bool;
-               loc : Location.t; }
+  | Assume of { property: property; strict: bool; never_returns_normally: bool;
+                loc: Location.t }
+  | Check of { property: property; strict: bool; loc : Location.t; }
 
 type fundecl =
   { fun_name: symbol;

backend/cmm.mli

@@ -317,7 +317,9 @@ type codegen_option =
   | No_CSE
   | Use_linscan_regalloc
   | Ignore_assert_all of property
-  | Check of { property: property; strict: bool; assume: bool; loc: Location.t }
+  | Assume of { property: property; strict: bool; never_returns_normally: bool;
+                loc: Location.t }
+  | Check of { property: property; strict: bool; loc: Location.t }
 
 type fundecl =
   { fun_name: symbol;

backend/cmm_helpers.ml

@@ -4512,8 +4512,15 @@ let transl_property : Lambda.property -> Cmm.property = function
 let transl_attrib : Lambda.check_attribute -> Cmm.codegen_option list = function
   | Default_check -> []
   | Ignore_assert_all p -> [Ignore_assert_all (transl_property p)]
-  | Check { property; strict; assume; loc } ->
-    [Check { property = transl_property property; strict; assume; loc }]
+  | Assume { property; strict; never_returns_normally; loc } ->
+    [ Assume
+        { property = transl_property property;
+          strict;
+          never_returns_normally;
+          loc
+        } ]
+  | Check { property; strict; loc } ->
+    [Check { property = transl_property property; strict; loc }]
 
 let kind_of_layout (layout : Lambda.layout) =
   match layout with

backend/printcmm.ml

@@ -390,11 +390,15 @@ let codegen_option = function
   | Use_linscan_regalloc -> "linscan"
   | Ignore_assert_all property ->
     Printf.sprintf "ignore %s" (property_to_string property)
-  | Check { property; strict; assume; loc = _ } ->
-    Printf.sprintf "%s %s%s"
-      (if assume then "assume" else "assert")
+  | Assume { property; strict; never_returns_normally; loc = _ } ->
+    Printf.sprintf "assume_%s%s%s"
       (property_to_string property)
-      (if strict then " strict" else "")
+      (if strict then "_strict" else "")
+      (if strict then "_never_returns_normally" else "")
+  | Check { property; strict; loc = _ } ->
+    Printf.sprintf "assert_%s%s"
+      (property_to_string property)
+      (if strict then "_strict" else "")
 
 let print_codegen_options ppf l =
   List.iter (fun c -> fprintf ppf " %s" (codegen_option c)) l

middle_end/flambda2/simplify/simplify_set_of_closures.ml

@@ -493,6 +493,7 @@ let simplify_function context ~outer_dacc function_slot code_id
       match Code_metadata.check code_metadata with
       | Default_check -> !Clflags.zero_alloc_check_assert_all
       | Ignore_assert_all Zero_alloc -> false
+      | Assume { property = Zero_alloc; _ } -> false
       | Check { property = Zero_alloc; _ } -> true
     in
     if never_delete then Code_id.Set.singleton code_id else Code_id.Set.empty

middle_end/flambda2/terms/check_attribute.ml

@@ -23,10 +23,15 @@ end
 type t =
   | Default_check
   | Ignore_assert_all of Property.t
+  | Assume of
+      { property : Property.t;
+        strict : bool;
+        never_returns_normally : bool;
+        loc : Location.t
+      }
   | Check of
       { property : Property.t;
         strict : bool;
-        assume : bool;
         loc : Location.t
       }
 
@@ -35,27 +40,43 @@ let print ppf t =
   | Default_check -> ()
   | Ignore_assert_all property ->
     Format.fprintf ppf "@[ignore %a@]" Property.print property
-  | Check { property; strict; assume; loc = _ } ->
-    Format.fprintf ppf "@[%s%s %a@]"
-      (if assume then "assume" else "assert")
-      (if strict then " strict" else "")
-      Property.print property
+  | Assume { property; strict; never_returns_normally; loc = _ } ->
+    Format.fprintf ppf "@[assume_%a%s%s@]" Property.print property
+      (if strict then "_strict" else "")
+      (if never_returns_normally then "_never_returns_normally" else "")
+  | Check { property; strict; loc = _ } ->
+    Format.fprintf ppf "@[assert_%a%s@]" Property.print property
+      (if strict then "_strict" else "")
 
 let from_lambda : Lambda.check_attribute -> t = function
   | Default_check -> Default_check
   | Ignore_assert_all p -> Ignore_assert_all (Property.from_lambda p)
-  | Check { property; strict; assume; loc } ->
-    Check { property = Property.from_lambda property; strict; assume; loc }
+  | Assume { property; strict; never_returns_normally; loc } ->
+    Assume
+      { property = Property.from_lambda property;
+        strict;
+        never_returns_normally;
+        loc
+      }
+  | Check { property; strict; loc } ->
+    Check { property = Property.from_lambda property; strict; loc }
 
 let equal x y =
   match x, y with
   | Default_check, Default_check -> true
   | Ignore_assert_all p1, Ignore_assert_all p2 -> Property.equal p1 p2
-  | ( Check { property = p1; strict = s1; assume = a1; loc = loc1 },
-      Check { property = p2; strict = s2; assume = a2; loc = loc2 } ) ->
-    Property.equal p1 p2 && Bool.equal s1 s2 && Bool.equal a1 a2 && loc1 = loc2
-  | (Default_check | Ignore_assert_all _ | Check _), _ -> false
+  | ( Check { property = p1; strict = s1; loc = loc1 },
+      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 },
+      Assume
+        { property = p2; strict = s2; never_returns_normally = n2; loc = loc2 }
+    ) ->
+    Property.equal p1 p2 && Bool.equal s1 s2 && Bool.equal n1 n2
+    && Location.compare loc1 loc2 = 0
+  | (Default_check | Ignore_assert_all _ | Check _ | Assume _), _ -> false
 
 let is_default : t -> bool = function
   | Default_check -> true
-  | Ignore_assert_all _ | Check _ -> false
+  | Ignore_assert_all _ | Check _ | Assume _ -> false

middle_end/flambda2/terms/check_attribute.mli

@@ -17,10 +17,15 @@ end
 type t =
   | Default_check
   | Ignore_assert_all of Property.t
+  | Assume of
+      { property : Property.t;
+        strict : bool;
+        never_returns_normally : bool;
+        loc : Location.t
+      }
   | Check of
       { property : Property.t;
         strict : bool;
-        assume : bool;
         loc : Location.t
       }
 

middle_end/flambda2/to_cmm/to_cmm_set_of_closures.ml

@@ -324,8 +324,15 @@ let transl_check_attrib : Check_attribute.t -> Cmm.codegen_option list =
   function
   | Default_check -> []
   | Ignore_assert_all p -> [Ignore_assert_all (transl_property p)]
-  | Check { property; strict; assume; loc } ->
-    [Check { property = transl_property property; strict; assume; loc }]
+  | Assume { property; strict; never_returns_normally; loc } ->
+    [ Assume
+        { property = transl_property property;
+          strict;
+          never_returns_normally;
+          loc
+        } ]
+  | Check { property; strict; loc } ->
+    [Check { property = transl_property property; strict; loc }]
 
 (* Translation of the bodies of functions. *)
 

middle_end/printclambda.ml

@@ -105,7 +105,7 @@ and one_fun ppf f =
     in
     iter f.params f.arity.params_layout
   in
-  fprintf ppf "(fun@ %s%s%a@ %d@ @[<2>%t@]@ @[<2>%a@])"
+  fprintf ppf "(fun@ %s%s@ %a@ %d@ @[<2>%t@]@ @[<2>%a@])"
     f.label (layout f.arity.return_layout) Printlambda.check_attribute f.check
     (List.length f.arity.params_layout) idents lam f.body
 

ocaml/lambda/lambda.ml

@@ -496,9 +496,13 @@ type check_attribute =
   | Ignore_assert_all of property
   | Check of { property: property;
                strict: bool;
-               assume: bool;
                loc: Location.t;
              }
+  | Assume of { property: property;
+                strict: bool;
+                loc: Location.t;
+                never_returns_normally: bool;
+              }
 
 type loop_attribute =
   | Always_loop (* [@loop] or [@loop always] *)

ocaml/lambda/lambda.mli

@@ -391,11 +391,13 @@ type check_attribute =
                   then the property holds (but property violations on
                   exceptional returns or divering loops are ignored).
                   This definition may not be applicable to new properties. *)
-               assume: bool;
-               (* [assume=true] assume without checking that the
-                  property holds *)
                loc: Location.t;
              }
+  | Assume of { property: property;
+                strict: bool;
+                loc: Location.t;
+                never_returns_normally: bool;
+              }
 
 type loop_attribute =
   | Always_loop (* [@loop] or [@loop always] *)

ocaml/lambda/printlambda.ml

@@ -634,11 +634,15 @@ let check_attribute ppf check =
   | Default_check -> ()
   | Ignore_assert_all p ->
     fprintf ppf "ignore assert all %s@ " (check_property p)
-  | Check {property=p; assume; strict; loc = _} ->
-    fprintf ppf "%s %s%s@ "
-      (if assume then "assume" else "assert")
+  | Assume {property=p; strict; never_returns_normally; loc = _} ->
+    fprintf ppf "assume_%s%s%s@ "
       (check_property p)
-      (if strict then " strict" else "")
+      (if strict then "_strict" else "")
+      (if never_returns_normally then "_never_returns_normally" else "")
+  | Check {property=p; strict; loc = _} ->
+    fprintf ppf "assert_%s%s@ "
+      (check_property p)
+      (if strict then "_strict" else "")
 
 let function_attribute ppf t =
   if t.is_a_functor then

ocaml/lambda/translattribute.ml

@@ -251,11 +251,17 @@ let parse_property_attribute attr property =
   | Some {Parsetree.attr_name = {txt; loc}; attr_payload = payload}->
       parse_ids_payload txt loc
         ~default:Default_check
-        ~empty:(Check { property; strict = false; assume = false; loc; } )
+        ~empty:(Check { property; strict = false; loc; } )
         [
-          ["assume"], Check { property; strict = false; assume = true; loc; };
-          ["strict"], Check { property; strict = true; assume = false; loc; };
-          ["assume"; "strict"], Check { property; strict = true; assume = true; loc; };
+          ["assume"],
+          Assume { property; strict = false; never_returns_normally = false; loc; };
+          ["strict"], Check { property; strict = true; loc; };
+          ["assume"; "strict"],
+          Assume { property; strict = true; never_returns_normally = false; loc; };
+          ["assume"; "never_returns_normally"],
+          Assume { property; strict = false; never_returns_normally = true; loc; };
+          ["assume"; "strict"; "never_returns_normally"],
+          Assume { property; strict = true; never_returns_normally = true; loc; };
           ["ignore"], Ignore_assert_all property
         ]
         payload
@@ -303,15 +309,15 @@ let get_property_attribute l p =
   (match attr, res with
    | None, Default_check -> ()
    | _, Default_check -> ()
-   | None, (Check _ | Ignore_assert_all _ ) -> assert false
+   | None, (Check _ | Assume _ | Ignore_assert_all _) -> assert false
    | Some _, Ignore_assert_all _ -> ()
-   | Some attr, Check { assume; _ } ->
+   | Some _, Assume _ -> ()
+   | Some attr, Check _ ->
      if !Clflags.zero_alloc_check && !Clflags.native_code then
        (* The warning for unchecked functions will not trigger if the check is requested
           through the [@@@zero_alloc all] top-level annotation rather than through the
           function annotation [@zero_alloc]. *)
-       if not assume then
-         Builtin_attributes.register_property attr.attr_name);
+       Builtin_attributes.register_property attr.attr_name);
    res
 
 let get_poll_attribute l =
@@ -414,7 +420,8 @@ let assume_zero_alloc attributes =
   match parse_property_attribute attr p with
   | Default_check -> false
   | Ignore_assert_all _ -> false
-  | Check { property = Zero_alloc; assume } -> assume
+  | Assume { property = Zero_alloc; _ } -> true
+  | Check { property = Zero_alloc; _ } -> false
 
 let assume_zero_alloc attributes =
   (* This function is used for "look-ahead" to find attributes
@@ -428,9 +435,12 @@ let add_check_attribute expr loc attributes =
     | Zero_alloc -> "zero_alloc"
   in
   let to_string = function
-    | Check { property; strict; assume; loc = _} ->
-      Printf.sprintf "%s %s%s"
-        (if assume then "assume" else "assert")
+    | Check { property; strict; loc = _} ->
+      Printf.sprintf "assert %s%s"
+        (to_string property)
+        (if strict then " strict" else "")
+    | Assume { property; strict; loc = _} ->
+      Printf.sprintf "assume %s%s"
         (to_string property)
         (if strict then " strict" else "")
     | Ignore_assert_all property ->
@@ -441,11 +451,13 @@ let add_check_attribute expr loc attributes =
   | Lfunction({ attr = { stub = false } as attr; } as funct) ->
     begin match get_property_attribute attributes Zero_alloc with
     | Default_check -> expr
-    | (Ignore_assert_all p | Check { property = p; _ }) as check ->
+    | (Ignore_assert_all p | Check { property = p; _ } | Assume { property = p; _ })
+      as check ->
       begin match attr.check with
       | Default_check -> ()
       | Ignore_assert_all p'
-      | Check { property = p'; strict = _; assume = _; loc = _; } ->
+      | Assume { property = p'; strict = _; loc = _; }
+      | Check { property = p'; strict = _; loc = _; } ->
         if p = p' then
           Location.prerr_warning loc
             (Warnings.Duplicated_attribute (to_string check));

tests/backend/checkmach/dune.inc

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