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

Author: gretay-js

backend/checkmach.ml

@@ -1162,7 +1162,13 @@ end
 module Spec_zero_alloc : Spec = struct
   let property = Cmm.Zero_alloc
 
-  let enabled () = !Clflags.zero_alloc_check
+  let enabled () =
+    (* Checkmach no longer distinguishes between opt and default checks. *)
+    match !Clflags.zero_alloc_check with
+    | No_check -> false
+    | Check_default -> true
+    | Check_all -> true
+    | Check_opt_only -> true
 
   (* Compact the mapping from function name to Value.t to reduce size of Checks
      in cmx and memory consumption Compilenv. Different components have

backend/cmm_helpers.ml

@@ -4563,8 +4563,10 @@ let transl_attrib : Lambda.check_attribute -> Cmm.codegen_option list = function
           never_returns_normally;
           loc
         } ]
-  | Check { property; strict; loc } ->
-    [Check { property = transl_property property; strict; loc }]
+  | Check { property; strict; loc; opt } ->
+    if Lambda.is_check_enabled ~opt property
+    then [Check { property = transl_property property; strict; loc }]
+    else []
 
 let kind_of_layout (layout : Lambda.layout) =
   match layout with

driver/flambda_backend_args.ml

@@ -90,8 +90,10 @@ let mk_heap_reduction_threshold f =
 ;;
 
 let mk_zero_alloc_check f =
-  "-zero-alloc-check", Arg.Unit f, " Check that annoted functions do not allocate \
-                                    and do not have indirect calls"
+  let annotations = Clflags.Annotations.(List.map to_string all) in
+  "-zero-alloc-check", Arg.Symbol (annotations, f),
+  " Check that annotated functions do not allocate \
+   and do not have indirect calls. "^Clflags.Annotations.doc
 
 let mk_dcheckmach f =
   "-dcheckmach", Arg.Unit f, " (undocumented)"
@@ -586,7 +588,7 @@ module type Flambda_backend_options = sig
   val dno_asm_comments : unit -> unit
 
   val heap_reduction_threshold : int -> unit
-  val zero_alloc_check : unit -> unit
+  val zero_alloc_check : string -> unit
   val dcheckmach : unit -> unit
   val checkmach_details_cutoff : int -> unit
 
@@ -837,7 +839,12 @@ module Flambda_backend_options_impl = struct
   let heap_reduction_threshold x =
     Flambda_backend_flags.heap_reduction_threshold := x
 
-  let zero_alloc_check = set' Clflags.zero_alloc_check
+  let zero_alloc_check s =
+    match Clflags.Annotations.of_string s with
+    | None -> () (* this should not occur as we use Arg.Symbol *)
+    | Some a ->
+      Clflags.zero_alloc_check := a
+
   let dcheckmach = set' Flambda_backend_flags.dump_checkmach
   let checkmach_details_cutoff n =
     let c : Flambda_backend_flags.checkmach_details_cutoff =
@@ -1096,7 +1103,13 @@ module Extra_params = struct
        set_int_option' Flambda_backend_flags.reorder_blocks_random
     | "basic-block-sections" -> set' Flambda_backend_flags.basic_block_sections
     | "heap-reduction-threshold" -> set_int' Flambda_backend_flags.heap_reduction_threshold
-    | "zero-alloc-check" -> set' Clflags.zero_alloc_check
+    | "zero-alloc-check" ->
+      (match Clflags.Annotations.of_string v with
+       | Some a -> Clflags.zero_alloc_check := a; true
+       | None ->
+         raise
+           (Arg.Bad
+              (Printf.sprintf "Unexpected value %s for %s" v name)))
     | "dump-checkmach" -> set' Flambda_backend_flags.dump_checkmach
     | "checkmach-details-cutoff" ->
       begin match Compenv.check_int ppf name v with

driver/flambda_backend_args.mli

@@ -44,7 +44,7 @@ module type Flambda_backend_options = sig
   val dno_asm_comments : unit -> unit
 
   val heap_reduction_threshold : int -> unit
-  val zero_alloc_check : unit -> unit
+  val zero_alloc_check : string -> unit
   val dcheckmach : unit -> unit
   val checkmach_details_cutoff : int -> unit
 

middle_end/flambda2/terms/check_attribute.ml

@@ -58,8 +58,10 @@ let from_lambda : Lambda.check_attribute -> t = function
         never_returns_normally;
         loc
       }
-  | Check { property; strict; loc } ->
-    Check { property = Property.from_lambda property; strict; loc }
+  | Check { property; strict; opt; loc } ->
+    if Lambda.is_check_enabled ~opt property
+    then Check { property = Property.from_lambda property; strict; loc }
+    else Default_check
 
 let equal x y =
   match x, y with

ocaml/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

ocaml/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

ocaml/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 =

ocaml/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]. *)

ocaml/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
       | _ ->