zero_alloc: cleanup attributes in Lambda (#2723)

gretay-js · web-flow · commit fa9bd8595b7d · 2024-06-28T17:01:02.000+01:00
* Remove [Ignore_assert_all] and [opt] from Lambda onwards

* format

* Fix tests

* fix tests (no local allocation mode)

* Zero_alloc_attribute.t = Lambda.zero_alloc_attribute
diff --git a/middle_end/flambda2/from_lambda/closure_conversion.ml b/middle_end/flambda2/from_lambda/closure_conversion.ml
@@ -1880,8 +1880,7 @@ let make_unboxed_function_wrapper acc function_slot ~unarized_params:params
         (Poll_attribute.from_lambda (Function_decl.poll_attribute decl))
       ~zero_alloc_attribute:
         (Zero_alloc_attribute.from_lambda
-           (Function_decl.zero_alloc_attribute decl)
-           (Debuginfo.Scoped_location.to_location (Function_decl.loc decl)))
+           (Function_decl.zero_alloc_attribute decl))
       ~is_a_functor:(Function_decl.is_a_functor decl)
       ~is_opaque:false ~recursive ~newer_version_of:None ~cost_metrics
       ~inlining_arguments:(Inlining_arguments.create ~round:0)
@@ -2249,8 +2248,7 @@ let close_one_function acc ~code_id ~external_env ~by_function_slot
         (Poll_attribute.from_lambda (Function_decl.poll_attribute decl))
       ~zero_alloc_attribute:
         (Zero_alloc_attribute.from_lambda
-           (Function_decl.zero_alloc_attribute decl)
-           (Debuginfo.Scoped_location.to_location (Function_decl.loc decl)))
+           (Function_decl.zero_alloc_attribute decl))
       ~is_a_functor:(Function_decl.is_a_functor decl)
       ~is_opaque:(Function_decl.is_opaque decl)
       ~recursive ~newer_version_of:None ~cost_metrics
@@ -2379,7 +2377,6 @@ let close_functions acc external_env ~current_region function_declarations =
         let zero_alloc_attribute =
           Zero_alloc_attribute.from_lambda
             (Function_decl.zero_alloc_attribute decl)
-            (Debuginfo.Scoped_location.to_location (Function_decl.loc decl))
         in
         let cost_metrics = Cost_metrics.zero in
         let dbg = Debuginfo.from_location (Function_decl.loc decl) in
diff --git a/middle_end/flambda2/parser/fexpr_to_flambda.ml b/middle_end/flambda2/parser/fexpr_to_flambda.ml
@@ -925,7 +925,7 @@ let rec expr env (e : Fexpr.expr) : Flambda.Expr.t =
             ~first_complex_local_param:(Flambda_arity.num_params params_arity)
             ~result_arity ~result_types:Unknown ~result_mode
             ~contains_no_escaping_local_allocs:false ~stub:false ~inline
-            ~zero_alloc_attribute:Default_check
+            ~zero_alloc_attribute:Default_zero_alloc
               (* CR gyorsh: should [check] be set properly? *)
             ~is_a_functor:false ~is_opaque:false ~recursive
             ~cost_metrics (* CR poechsel: grab inlining arguments from fexpr. *)
diff --git a/middle_end/flambda2/simplify/simplify_apply_expr.ml b/middle_end/flambda2/simplify/simplify_apply_expr.ml
@@ -631,7 +631,7 @@ let simplify_direct_partial_application ~simplify_expr dacc apply
                 (Code_metadata.contains_no_escaping_local_allocs
                    callee's_code_metadata)
               ~stub:true ~inline:Default_inline ~poll_attribute:Default
-              ~zero_alloc_attribute:Zero_alloc_attribute.Default_check
+              ~zero_alloc_attribute:Zero_alloc_attribute.Default_zero_alloc
               ~is_a_functor:false ~is_opaque:false ~recursive
               ~cost_metrics:cost_metrics_of_body
               ~inlining_arguments:(DE.inlining_arguments (DA.denv dacc))
diff --git a/middle_end/flambda2/simplify/simplify_set_of_closures.ml b/middle_end/flambda2/simplify/simplify_set_of_closures.ml
@@ -510,7 +510,7 @@ let simplify_function context ~outer_dacc function_slot code_id
     let code_metadata = Code_or_metadata.code_metadata code_or_metadata in
     let never_delete =
       match Code_metadata.zero_alloc_attribute code_metadata with
-      | Default_check -> !Clflags.zero_alloc_check_assert_all
+      | Default_zero_alloc -> false
       | Assume _ -> false
       | Check _ -> true
     in
diff --git a/middle_end/flambda2/terms/zero_alloc_attribute.ml b/middle_end/flambda2/terms/zero_alloc_attribute.ml
@@ -10,22 +10,22 @@
 (*                                                                        *)
 (**************************************************************************)
 
-type t =
-  | Default_check
-  | Assume of
+type t = Lambda.zero_alloc_attribute =
+  | Default_zero_alloc
+  | Check of
       { strict : bool;
-        never_returns_normally : bool;
-        never_raises : bool;
         loc : Location.t
       }
-  | Check of
+  | Assume of
       { strict : bool;
+        never_returns_normally : bool;
+        never_raises : bool;
         loc : Location.t
       }
 
 let print ppf t =
   match t with
-  | Default_check -> ()
+  | Default_zero_alloc -> ()
   | Assume { strict; never_returns_normally; never_raises; loc = _ } ->
     Format.fprintf ppf "@[assume_zero_alloc%s%s%s@]"
       (if strict then "_strict" else "")
@@ -35,25 +35,11 @@ let print ppf t =
     Format.fprintf ppf "@[assert_zero_alloc%s@]"
       (if strict then "_strict" else "")
 
-let from_lambda : Lambda.zero_alloc_attribute -> Location.t -> t =
- fun a loc ->
-  match a with
-  | Default_zero_alloc ->
-    if !Clflags.zero_alloc_check_assert_all
-       && Builtin_attributes.is_zero_alloc_check_enabled ~opt:false
-    then Check { strict = false; loc }
-    else Default_check
-  | Ignore_assert_all -> Default_check
-  | Assume { strict; never_returns_normally; never_raises; loc; arity = _ } ->
-    Assume { strict; never_returns_normally; never_raises; loc }
-  | Check { strict; opt; loc; arity = _ } ->
-    if Builtin_attributes.is_zero_alloc_check_enabled ~opt
-    then Check { strict; loc }
-    else Default_check
+let from_lambda : Lambda.zero_alloc_attribute -> t = Fun.id
 
 let equal x y =
   match x, y with
-  | Default_check, Default_check -> true
+  | Default_zero_alloc, Default_zero_alloc -> true
   | Check { strict = s1; loc = loc1 }, Check { strict = s2; loc = loc2 } ->
     Bool.equal s1 s2 && Location.compare loc1 loc2 = 0
   | ( Assume
@@ -70,8 +56,8 @@ let equal x y =
         } ) ->
     Bool.equal s1 s2 && Bool.equal n1 n2 && Bool.equal r1 r2
     && Location.compare loc1 loc2 = 0
-  | (Default_check | Check _ | Assume _), _ -> false
+  | (Default_zero_alloc | Check _ | Assume _), _ -> false
 
 let is_default : t -> bool = function
-  | Default_check -> true
+  | Default_zero_alloc -> true
   | Check _ | Assume _ -> false
diff --git a/middle_end/flambda2/terms/zero_alloc_attribute.mli b/middle_end/flambda2/terms/zero_alloc_attribute.mli
@@ -10,16 +10,16 @@
 (*                                                                        *)
 (**************************************************************************)
 (** [@zero_alloc ...] annotations on function declaration (not call sites) *)
-type t =
-  | Default_check
-  | Assume of
+type t = Lambda.zero_alloc_attribute =
+  | Default_zero_alloc
+  | Check of
       { strict : bool;
-        never_returns_normally : bool;
-        never_raises : bool;
         loc : Location.t
       }
-  | Check of
+  | Assume of
       { strict : bool;
+        never_returns_normally : bool;
+        never_raises : bool;
         loc : Location.t
       }
 
@@ -29,4 +29,4 @@ val equal : t -> t -> bool
 
 val is_default : t -> bool
 
-val from_lambda : Lambda.zero_alloc_attribute -> Location.t -> t
+val from_lambda : Lambda.zero_alloc_attribute -> t
diff --git a/middle_end/flambda2/to_cmm/to_cmm_set_of_closures.ml b/middle_end/flambda2/to_cmm/to_cmm_set_of_closures.ml
@@ -351,7 +351,7 @@ end)
 
 let transl_check_attrib : Zero_alloc_attribute.t -> Cmm.codegen_option list =
   function
-  | Default_check -> []
+  | Default_zero_alloc -> []
   | Assume { strict; never_returns_normally; never_raises; loc } ->
     [Assume_zero_alloc { strict; never_returns_normally; never_raises; loc }]
   | Check { strict; loc } -> [Check_zero_alloc { strict; loc }]
diff --git a/ocaml/lambda/lambda.ml b/ocaml/lambda/lambda.ml
@@ -634,18 +634,14 @@ type poll_attribute =
   | Error_poll (* [@poll error] *)
   | Default_poll (* no [@poll] attribute *)
 
-type zero_alloc_attribute = Builtin_attributes.zero_alloc_attribute =
+type zero_alloc_attribute =
   | Default_zero_alloc
-  | Ignore_assert_all
   | Check of { strict: bool;
-               opt: bool;
-               arity: int;
                loc: Location.t;
              }
   | Assume of { strict: bool;
                 never_returns_normally: bool;
                 never_raises: bool;
-                arity: int;
                 loc: Location.t;
               }
 
@@ -925,7 +921,7 @@ let default_function_attribute = {
 }
 
 let default_stub_attribute =
-  { default_function_attribute with stub = true; zero_alloc = Ignore_assert_all }
+  { default_function_attribute with stub = true; zero_alloc = Default_zero_alloc }
 
 let default_param_attribute = { unbox_param = false }
 
diff --git a/ocaml/lambda/lambda.mli b/ocaml/lambda/lambda.mli
@@ -510,23 +510,19 @@ type poll_attribute =
   | Error_poll (* [@poll error] *)
   | Default_poll (* no [@poll] attribute *)
 
-type zero_alloc_attribute = Builtin_attributes.zero_alloc_attribute =
+type zero_alloc_attribute =
   | Default_zero_alloc
-  | Ignore_assert_all
   | Check of { strict: bool;
                (* [strict=true] property holds on all paths.
                   [strict=false] if the function returns normally,
                   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;
-               arity: int;
                loc: Location.t;
              }
   | Assume of { strict: bool;
                 never_returns_normally: bool;
                 never_raises: bool;
-                arity: int;
                 loc: Location.t;
               }
 
diff --git a/ocaml/lambda/printlambda.ml b/ocaml/lambda/printlambda.ml
@@ -939,15 +939,12 @@ let name_of_primitive = function
 let zero_alloc_attribute ppf check =
   match check with
   | Default_zero_alloc -> ()
-  | Ignore_assert_all ->
-    fprintf ppf "ignore assert all zero_alloc@ "
   | Assume {strict; never_returns_normally; loc = _} ->
     fprintf ppf "assume_zero_alloc%s%s@ "
       (if strict then "_strict" else "")
       (if never_returns_normally then "_never_returns_normally" else "")
-  | Check {strict; loc = _; opt} ->
-    fprintf ppf "assert_zero_alloc%s%s@ "
-      (if opt then "_opt" else "")
+  | Check {strict; loc = _; } ->
+    fprintf ppf "assert_zero_alloc%s@ "
       (if strict then "_strict" else "")
 
 let function_attribute ppf t =
diff --git a/ocaml/lambda/translcore.ml b/ocaml/lambda/translcore.ml
@@ -1608,6 +1608,21 @@ and transl_function ~in_new_scope ~scopes e params body
            ~mode ~return_sort ~return_mode
            ~scopes e.exp_loc repr ~region params body)
   in
+  let zero_alloc : Lambda.zero_alloc_attribute =
+    match (zero_alloc : Builtin_attributes.zero_alloc_attribute) with
+    | Default_zero_alloc ->
+      if !Clflags.zero_alloc_check_assert_all &&
+         Builtin_attributes.is_zero_alloc_check_enabled ~opt:false
+      then Check { strict = false; loc = e.exp_loc }
+      else Default_zero_alloc
+    | Check { strict; opt; arity = _; loc } ->
+      if Builtin_attributes.is_zero_alloc_check_enabled ~opt
+      then Check { strict; loc }
+      else Default_zero_alloc
+    | Assume { strict; never_returns_normally; never_raises; loc; arity = _; } ->
+      Assume { strict; never_returns_normally; never_raises; loc }
+    | Ignore_assert_all -> Default_zero_alloc
+  in
   let attr =
     { function_attribute_disallowing_arity_fusion with zero_alloc }
   in
diff --git a/ocaml/lambda/translmod.ml b/ocaml/lambda/translmod.ml
@@ -162,7 +162,7 @@ and apply_coercion_result loc strict funct params args cc_res =
              ~return:Lambda.layout_module
              ~attr:{ default_function_attribute with
                         is_a_functor = true;
-                        zero_alloc = Ignore_assert_all;
+                        zero_alloc = Default_zero_alloc;
                         stub = true; }
              ~loc
              ~mode:alloc_heap
@@ -572,7 +572,7 @@ let rec compile_functor ~scopes mexp coercion root_path loc =
       loop = Never_loop;
       is_a_functor = true;
       is_opaque = false;
-      zero_alloc = Ignore_assert_all;
+      zero_alloc = Default_zero_alloc;
       stub = false;
       tmc_candidate = false;
       may_fuse_arity = true;
diff --git a/ocaml/testsuite/tests/functors/functors.compilers.reference b/ocaml/testsuite/tests/functors/functors.compilers.reference
@@ -1,34 +1,30 @@
 (setglobal Functors!
   (let
     (O =
-       (function {nlocal = 0} X is_a_functor always_inline
-         ignore assert all zero_alloc never_loop
+       (function {nlocal = 0} X is_a_functor always_inline never_loop
          (let
            (cow =
               (function {nlocal = 0} x[int] : int (apply (field_imm 0 X) x))
             sheep = (function {nlocal = 0} x[int] : int (+ 1 (apply cow x))))
            (makeblock 0 cow sheep)))
      F =
-       (function {nlocal = 0} X Y is_a_functor always_inline
-         ignore assert all zero_alloc never_loop
+       (function {nlocal = 0} X Y is_a_functor always_inline never_loop
          (let
            (cow =
               (function {nlocal = 0} x[int] : int
                 (apply (field_imm 0 Y) (apply (field_imm 0 X) x)))
             sheep = (function {nlocal = 0} x[int] : int (+ 1 (apply cow x))))
            (makeblock 0 cow sheep)))
      F1 =
-       (function {nlocal = 0} X Y is_a_functor always_inline
-         ignore assert all zero_alloc never_loop
+       (function {nlocal = 0} X Y is_a_functor always_inline never_loop
          (let
            (cow =
               (function {nlocal = 0} x[int] : int
                 (apply (field_imm 0 Y) (apply (field_imm 0 X) x)))
             sheep = (function {nlocal = 0} x[int] : int (+ 1 (apply cow x))))
            (makeblock 0 sheep)))
      F2 =
-       (function {nlocal = 0} X Y is_a_functor always_inline
-         ignore assert all zero_alloc never_loop
+       (function {nlocal = 0} X Y is_a_functor always_inline never_loop
          (let
            (X =a (makeblock 0 (field_imm 1 X))
             Y =a (makeblock 0 (field_imm 1 Y))
@@ -40,8 +36,7 @@
      M =
        (let
          (F =
-            (function {nlocal = 0} X Y is_a_functor always_inline
-              ignore assert all zero_alloc never_loop
+            (function {nlocal = 0} X Y is_a_functor always_inline never_loop
               (let
                 (cow =
                    (function {nlocal = 0} x[int] : int
@@ -51,7 +46,6 @@
                 (makeblock 0 cow sheep))))
          (makeblock 0
            (function {nlocal = 0} funarg funarg is_a_functor stub
-             ignore assert all zero_alloc
              (let
                (let =
                   (apply F (makeblock 0 (field_imm 1 funarg))
diff --git a/ocaml/testsuite/tests/translprim/array_spec.heap.flat.reference b/ocaml/testsuite/tests/translprim/array_spec.heap.flat.reference
diff --git a/ocaml/testsuite/tests/translprim/array_spec.stack.flat.reference b/ocaml/testsuite/tests/translprim/array_spec.stack.flat.reference
diff --git a/ocaml/testsuite/tests/translprim/comparison_table.heap.reference b/ocaml/testsuite/tests/translprim/comparison_table.heap.reference
diff --git a/ocaml/testsuite/tests/translprim/comparison_table.stack.reference b/ocaml/testsuite/tests/translprim/comparison_table.stack.reference
diff --git a/ocaml/testsuite/tests/translprim/module_coercion.compilers.flat.reference b/ocaml/testsuite/tests/translprim/module_coercion.compilers.flat.reference
