flambda-backend: Zero alloc: add payload "opt" and "-zero-alloc-check {default|all|none|opt}" flag (#1936)

Author: gretay-js

lambda/lambda.ml

@@ -509,6 +509,7 @@ type check_attribute =
   | Ignore_assert_all of property
   | Check of { property: property;
                strict: bool;
+               opt: bool;
                loc: Location.t;
              }
   | Assume of { property: property;
@@ -1662,3 +1663,12 @@ let array_set_kind mode = function
   | Paddrarray -> Paddrarray_set mode
   | Pintarray -> Pintarray_set
   | Pfloatarray -> Pfloatarray_set
+
+let is_check_enabled ~opt property =
+  match property with
+  | Zero_alloc ->
+    match !Clflags.zero_alloc_check with
+    | No_check -> false
+    | Check_all -> true
+    | Check_default -> not opt
+    | Check_opt_only -> opt

lambda/lambda.mli

@@ -401,6 +401,7 @@ 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. *)
+               opt: bool;
                loc: Location.t;
              }
   | Assume of { property: property;
@@ -770,3 +771,4 @@ val array_ref_kind : alloc_mode -> array_kind -> array_ref_kind
 
 (** The mode will be discarded if unnecessary for the given [array_kind] *)
 val array_set_kind : modify_mode -> array_kind -> array_set_kind
+val is_check_enabled : opt:bool -> property -> bool

lambda/printlambda.ml

@@ -687,9 +687,9 @@ let check_attribute ppf check =
       (check_property p)
       (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)
+  | Check {property=p; strict; loc = _; opt} ->
+    fprintf ppf "assert_%s%s%s@ "
+      (check_property p) (if opt then "_opt" else "")
       (if strict then "_strict" else "")
 
 let function_attribute ppf t =

lambda/translattribute.ml

@@ -251,11 +251,13 @@ 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; loc; } )
+        ~empty:(Check { property; strict = false; opt = false; loc; } )
         [
           ["assume"],
           Assume { property; strict = false; never_returns_normally = false; loc; };
-          ["strict"], Check { property; strict = true; loc; };
+          ["strict"], Check { property; strict = true; opt = false; loc; };
+          ["opt"], Check { property; strict = false; opt = true; loc; };
+          ["opt"; "strict"; ], Check { property; strict = true; opt = true; loc; };
           ["assume"; "strict"],
           Assume { property; strict = true; never_returns_normally = false; loc; };
           ["assume"; "never_returns_normally"],
@@ -312,8 +314,8 @@ let get_property_attribute l p =
    | None, (Check _ | Assume _ | Ignore_assert_all _) -> assert false
    | Some _, Ignore_assert_all _ -> ()
    | Some _, Assume _ -> ()
-   | Some attr, Check _ ->
-     if !Clflags.zero_alloc_check && !Clflags.native_code then
+   | Some attr, Check { opt; _ } ->
+     if Lambda.is_check_enabled ~opt p && !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]. *)

parsing/builtin_attributes.ml

@@ -586,7 +586,7 @@ let parse_attribute_with_ident_payload attr ~name ~f =
 let zero_alloc_attribute (attr : Parsetree.attribute)  =
   parse_attribute_with_ident_payload attr
     ~name:"zero_alloc" ~f:(function
-      | "check" -> Clflags.zero_alloc_check := true
+      | "check" -> Clflags.zero_alloc_check := Clflags.Annotations.Check_default
       | "all" ->
         Clflags.zero_alloc_check_assert_all := true
       | _ ->

utils/clflags.ml

@@ -635,5 +635,37 @@ let create_usage_msg program =
 let print_arguments program =
   Arg.usage !arg_spec (create_usage_msg program)
 
-let zero_alloc_check = ref false            (* -zero-alloc-check *)
+module Annotations = struct
+  type t = Check_default | Check_all | Check_opt_only | No_check
+
+  let all = [ Check_default; Check_all; Check_opt_only; No_check ]
+
+  let to_string = function
+    | Check_default -> "default"
+    | Check_all -> "all"
+    | Check_opt_only -> "opt"
+    | No_check -> "none"
+
+  let equal t1 t2 =
+    match t1, t2 with
+    | Check_default, Check_default -> true
+    | Check_all, Check_all -> true
+    | No_check, No_check -> true
+    | Check_opt_only, Check_opt_only -> true
+    | (Check_default | Check_all | Check_opt_only | No_check), _ -> false
+
+  let of_string v =
+    let f t =
+      if String.equal (to_string t) v then Some t else None
+    in
+    List.find_map f all
+
+  let doc =
+    "\n\    The argument specifies which annotations to check: \n\
+     \      \"opt\" means attributes with \"opt\" payload and is intended for debugging;\n\
+     \      \"default\" means attributes without \"opt\" payload; \n\
+     \      \"all\" covers both \"opt\" and \"default\" and is intended for optimized builds."
+end
+
+let zero_alloc_check = ref Annotations.No_check         (* -zero-alloc-check *)
 let zero_alloc_check_assert_all = ref false (* -zero-alloc-check-assert-all *)

utils/clflags.mli

@@ -286,6 +286,15 @@ val print_arguments : string -> unit
 (* [reset_arguments ()] clear all declared arguments *)
 val reset_arguments : unit -> unit
 
-val zero_alloc_check : bool ref
+(* [Annotations] specifies which zero_alloc attributes to check. *)
+module Annotations : sig
+  type t = Check_default | Check_all | Check_opt_only | No_check
+  val all : t list
+  val to_string : t -> string
+  val of_string : string -> t option
+  val equal : t -> t -> bool
+  val doc : string
+end
+val zero_alloc_check : Annotations.t ref
 val zero_alloc_check_assert_all : bool ref