flambda-backend: Add support for using zero_alloc information from a signature (#2506)

Author: ccasin

typing/typecore.ml

@@ -5056,6 +5056,44 @@ let add_check_attribute expr attributes =
     end
   | _ -> expr
 
+let zero_alloc_of_application ~num_args attrs funct =
+  let zero_alloc =
+    Builtin_attributes.get_property_attribute ~in_signature:false
+      ~default_arity:num_args attrs Zero_alloc
+  in
+  let zero_alloc =
+    match zero_alloc with
+    | Assume _ | Ignore_assert_all _ | Check _ ->
+      (* The user wrote a zero_alloc attribute on the application - keep it.
+         (Note that `ignore` and `check` aren't really allowed here, and will be
+         rejected by the call to `Builtin_attributes.assume_zero_alloc` below.)
+       *)
+      zero_alloc
+    | Default_check ->
+      (* We assume the call is zero_alloc if the function is known to be
+         zero_alloc. If the function is zero_alloc opt, then we need to be sure
+         that the opt checks were run to license this assumption. We judge
+         whether the opt checks were run based on the argument to the
+         [-zero-alloc-check] command line flag. *)
+      let use_opt =
+        match !Clflags.zero_alloc_check with
+        | Check_default | No_check -> false
+        | Check_all | Check_opt_only -> true
+      in
+      match funct.exp_desc with
+      | Texp_ident (_, _, { val_zero_alloc = (Check c); _ }, _, _)
+        when c.arity = num_args && (use_opt || not c.opt) ->
+        Builtin_attributes.Assume {
+          property = Zero_alloc;
+          strict = c.strict;
+          never_returns_normally = false;
+          arity = c.arity;
+          loc = c.loc
+        }
+      | _ -> Builtin_attributes.Default_check
+  in
+  Builtin_attributes.assume_zero_alloc ~is_check_allowed:false zero_alloc
+
 let rec type_exp ?recarg env expected_mode sexp =
   (* We now delegate everything to type_expect *)
   type_expect ?recarg env expected_mode sexp
@@ -5384,12 +5422,8 @@ and type_expect_
         type_application env loc expected_mode pm funct funct_mode sargs rt
       in
       let assume_zero_alloc =
-        let default_arity = List.length args in
-        let zero_alloc =
-          Builtin_attributes.get_property_attribute ~in_signature:false
-            ~default_arity sfunct.pexp_attributes Zero_alloc
-        in
-        Builtin_attributes.assume_zero_alloc ~is_check_allowed:false zero_alloc
+        zero_alloc_of_application ~num_args:(List.length args)
+          sfunct.pexp_attributes funct
       in
 
       rue {