flambda-backend: Infer needed zero_alloc checks from signatures (#2742)

Author: ccasin

compilerlibs/Makefile.compilerlibs

@@ -89,6 +89,7 @@ TYPING = \
   typing/jkind_types.cmo \
   typing/primitive.cmo \
   typing/shape.cmo \
+  typing/zero_alloc.cmo \
   typing/types.cmo \
   typing/jkind.cmo \
   typing/btype.cmo \

dune

@@ -88,7 +88,8 @@
    jane_syntax ; manual update: mli only files asttypes parsetree
 
    ;; TYPING
-   ident path jkind primitive shape shape_reduce types btype oprint subst predef datarepr
+   ident path jkind primitive shape shape_reduce zero_alloc types btype oprint subst
+   predef datarepr
    global_module cmi_format persistent_env env errortrace mode jkind_types jkind_intf
    typedtree printtyped ctype printtyp includeclass mtype envaux includecore
    tast_iterator tast_mapper signature_group cmt_format cms_format untypeast
@@ -291,6 +292,7 @@
     (jkind.mli as compiler-libs/jkind.mli)
     (jkind_types.mli as compiler-libs/jkind_types.mli)
     (primitive.mli as compiler-libs/primitive.mli)
+    (zero_alloc.mli as compiler-libs/zero_alloc.mli)
     (types.mli as compiler-libs/types.mli)
     (btype.mli as compiler-libs/btype.mli)
     (binutils.mli as compiler-libs/binutils.mli)

lambda/translcore.ml

@@ -350,6 +350,38 @@ let can_apply_primitive p pmode pos args =
     end
   end
 
+let zero_alloc_of_application
+      ~num_args (annotation : Zero_alloc.assume option) funct =
+  match annotation, funct.exp_desc with
+  | Some assume, _ ->
+    (* The user wrote a zero_alloc attribute on the application - keep it. *)
+    Builtin_attributes.assume_zero_alloc assume
+  | None, Texp_ident (_, _, { val_zero_alloc; _ }, _, _) ->
+    (* 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
+    begin match Zero_alloc.get val_zero_alloc with
+    | Check c when c.arity = num_args && (use_opt || not c.opt) ->
+      let assume : Zero_alloc.assume =
+        { strict = c.strict;
+          never_returns_normally = false;
+          never_raises = false;
+          arity = c.arity;
+          loc = c.loc }
+      in
+      Builtin_attributes.assume_zero_alloc assume
+    | Check _ | Default_zero_alloc | Ignore_assert_all | Assume _ ->
+      Zero_alloc_utils.Assume_info.none
+    end
+  | None, _ -> Zero_alloc_utils.Assume_info.none
+
 let rec transl_exp ~scopes sort e =
   transl_exp1 ~scopes ~in_new_scope:false sort e
 
@@ -387,7 +419,7 @@ and transl_exp0 ~in_new_scope ~scopes sort e =
   | Texp_apply({ exp_desc = Texp_ident(path, _, {val_kind = Val_prim p},
                                        Id_prim (pmode, psort), _);
                  exp_type = prim_type; } as funct,
-               oargs, pos, ap_mode, assume_zero_alloc)
+               oargs, pos, ap_mode, zero_alloc)
     when can_apply_primitive p pmode pos oargs ->
       let rec cut_args prim_repr oargs =
         match prim_repr, oargs with
@@ -409,6 +441,11 @@ and transl_exp0 ~in_new_scope ~scopes sort e =
         if extra_args = [] then transl_apply_position pos
         else Rc_normal
       in
+      let assume_zero_alloc =
+        match zero_alloc with
+        | None -> Zero_alloc_utils.Assume_info.none
+        | Some assume -> Builtin_attributes.assume_zero_alloc assume
+      in
       let lam =
         let loc =
           map_scopes (update_assume_zero_alloc ~assume_zero_alloc)
@@ -433,14 +470,17 @@ and transl_exp0 ~in_new_scope ~scopes sort e =
              ~position ~mode
              ~result_layout lam extra_args (of_location ~scopes e.exp_loc))
       end
-  | Texp_apply(funct, oargs, position, ap_mode, assume_zero_alloc)
+  | Texp_apply(funct, oargs, position, ap_mode, zero_alloc)
     ->
       let tailcall = Translattribute.get_tailcall_attribute funct in
       let inlined = Translattribute.get_inlined_attribute funct in
       let specialised = Translattribute.get_specialised_attribute funct in
       let result_layout = layout_exp sort e in
       let position = transl_apply_position position in
       let mode = transl_locality_mode_l ap_mode in
+      let assume_zero_alloc =
+        zero_alloc_of_application ~num_args:(List.length oargs) zero_alloc funct
+      in
       event_after ~scopes e
         (transl_apply ~scopes ~tailcall ~inlined ~specialised
            ~assume_zero_alloc
@@ -1579,8 +1619,13 @@ and transl_function ~in_new_scope ~scopes e params body
       ~zero_alloc =
   let attrs = e.exp_attributes in
   let mode = transl_alloc_mode_r alloc_mode in
+  let zero_alloc = Zero_alloc.get zero_alloc in
   let assume_zero_alloc =
-    Builtin_attributes.assume_zero_alloc ~is_check_allowed:true zero_alloc
+    match zero_alloc with
+    | Default_zero_alloc | Check _ | Ignore_assert_all ->
+      Zero_alloc_utils.Assume_info.none
+    | Assume assume ->
+      Builtin_attributes.assume_zero_alloc assume
   in
   let scopes =
     if in_new_scope then

otherlibs/dynlink/Makefile

@@ -121,6 +121,7 @@ COMPILERLIBS_SOURCES=\
   typing/shape.ml \
   typing/jkind_types.ml \
   typing/primitive.ml \
+  typing/zero_alloc.ml \
   typing/types.ml \
   typing/jkind.ml \
   typing/typedtree.ml \

otherlibs/dynlink/dune

@@ -94,6 +94,7 @@
     jkind_intf
     jkind_types
     primitive
+    zero_alloc
     types
     jkind
     value_rec_types
@@ -198,6 +199,7 @@
 (copy_files ../../typing/solver.ml)
 (copy_files ../../typing/shape_reduce.ml)
 (copy_files ../../typing/mode.ml)
+(copy_files ../../typing/zero_alloc.ml)
 (copy_files ../../typing/types.ml)
 (copy_files ../../typing/btype.ml)
 (copy_files ../../typing/subst.ml)
@@ -267,6 +269,7 @@
 (copy_files ../../typing/solver.mli)
 (copy_files ../../typing/shape_reduce.mli)
 (copy_files ../../typing/mode.mli)
+(copy_files ../../typing/zero_alloc.mli)
 (copy_files ../../typing/types.mli)
 (copy_files ../../typing/btype.mli)
 (copy_files ../../typing/subst.mli)
@@ -379,6 +382,7 @@
       .dynlink_compilerlibs.objs/byte/dynlink_compilerlibs__Mode.cmo
       .dynlink_compilerlibs.objs/byte/dynlink_compilerlibs__Jkind_intf.cmo
       .dynlink_compilerlibs.objs/byte/dynlink_compilerlibs__Jkind_types.cmo
+      .dynlink_compilerlibs.objs/byte/dynlink_compilerlibs__Zero_alloc.cmo
       .dynlink_compilerlibs.objs/byte/dynlink_compilerlibs__Types.cmo
       .dynlink_compilerlibs.objs/byte/dynlink_compilerlibs__Attr_helper.cmo
       .dynlink_compilerlibs.objs/byte/dynlink_compilerlibs__Primitive.cmo
@@ -462,6 +466,7 @@
       .dynlink_compilerlibs.objs/native/dynlink_compilerlibs__Mode.cmx
       .dynlink_compilerlibs.objs/native/dynlink_compilerlibs__Jkind_intf.cmx
       .dynlink_compilerlibs.objs/native/dynlink_compilerlibs__Jkind_types.cmx
+      .dynlink_compilerlibs.objs/native/dynlink_compilerlibs__Zero_alloc.cmx
       .dynlink_compilerlibs.objs/native/dynlink_compilerlibs__Types.cmx
       .dynlink_compilerlibs.objs/native/dynlink_compilerlibs__Attr_helper.cmx
       .dynlink_compilerlibs.objs/native/dynlink_compilerlibs__Primitive.cmx

parsing/builtin_attributes.ml

@@ -681,20 +681,26 @@ let error_message_attr l =
     | _ -> None in
   List.find_map inner l
 
+type zero_alloc_check =
+  { strict: bool;
+    opt: bool;
+    arity: int;
+    loc: Location.t;
+  }
+
+type zero_alloc_assume =
+  { strict: bool;
+    never_returns_normally: bool;
+    never_raises: bool;
+    arity: int;
+    loc: Location.t;
+  }
+
 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;
-              }
+  | Check of zero_alloc_check
+  | Assume of zero_alloc_assume
 
 let is_zero_alloc_check_enabled ~opt =
   match !Clflags.zero_alloc_check with
@@ -923,25 +929,26 @@ let get_zero_alloc_attribute ~in_signature ~default_arity l =
        register_zero_alloc_attribute attr.attr_name);
    res
 
-let assume_zero_alloc ~is_check_allowed check : Zero_alloc_utils.Assume_info.t =
-  match check with
-  | Default_zero_alloc -> Zero_alloc_utils.Assume_info.none
-  | Ignore_assert_all -> Zero_alloc_utils.Assume_info.none
-  | Assume { strict; never_returns_normally; never_raises; } ->
-    Zero_alloc_utils.Assume_info.create ~strict ~never_returns_normally ~never_raises
+let zero_alloc_attribute_only_assume_allowed za =
+  match za with
+  | Assume assume -> Some assume
+  | Default_zero_alloc | Ignore_assert_all -> None
   | Check { loc; _ } ->
-    if not is_check_allowed then begin
-      let name = "zero_alloc" in
-      let msg = "Only the following combinations are supported in this context: \
-                 'zero_alloc assume', \
-                 `zero_alloc assume strict`, \
-                 `zero_alloc assume error`,\
-                 `zero_alloc assume never_returns_normally`,\
-                 `zero_alloc assume never_returns_normally strict`."
-      in
-      Location.prerr_warning loc (Warnings.Attribute_payload (name, msg))
-    end;
-    Zero_alloc_utils.Assume_info.none
+    let name = "zero_alloc" in
+    let msg = "Only the following combinations are supported in this context: \
+               'zero_alloc assume', \
+               `zero_alloc assume strict`, \
+               `zero_alloc assume error`,\
+               `zero_alloc assume never_returns_normally`,\
+               `zero_alloc assume never_returns_normally strict`."
+    in
+    Location.prerr_warning loc (Warnings.Attribute_payload (name, msg));
+    None
+
+let assume_zero_alloc assume : Zero_alloc_utils.Assume_info.t =
+  match assume with
+  | { strict; never_returns_normally; never_raises; } ->
+    Zero_alloc_utils.Assume_info.create ~strict ~never_returns_normally ~never_raises
 
 type tracing_probe =
   { name : string;

parsing/builtin_attributes.mli

@@ -237,30 +237,36 @@ val parse_optional_id_payload :
   string -> Location.t -> empty:'a -> (string * 'a) list ->
   Parsetree.payload -> ('a,unit) Result.t
 
-(* Support for property attributes like zero_alloc *)
+(* Support for zero_alloc *)
+type zero_alloc_check =
+  { 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 diverging loops are ignored).
+       This definition may not be applicable to new properties. *)
+    opt: bool;
+    arity: int;
+    loc: Location.t;
+  }
+
+type zero_alloc_assume =
+  { strict: bool;
+    never_returns_normally: bool;
+    never_raises: bool;
+    (* [never_raises=true] the function never returns
+       via an exception. The function (directly or transitively)
+       may raise exceptions that do not escape, i.e.,
+       handled before the function returns. *)
+    arity: int;
+    loc: Location.t;
+  }
+
 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 diverging 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;
-                (* [never_raises=true] the function never returns
-                   via an exception. The function (directly or transitively)
-                   may raise exceptions that do not escape, i.e.,
-                   handled before the function returns. *)
-                arity: int;
-                loc: Location.t;
-              }
+  | Check of zero_alloc_check
+  | Assume of zero_alloc_assume
 
 val is_zero_alloc_check_enabled : opt:bool -> bool
 
@@ -271,8 +277,12 @@ val get_zero_alloc_attribute :
   in_signature:bool -> default_arity:int -> Parsetree.attributes ->
   zero_alloc_attribute
 
-val assume_zero_alloc :
-  is_check_allowed:bool -> zero_alloc_attribute -> Zero_alloc_utils.Assume_info.t
+(* This returns the [zero_alloc_assume] if the input is an assume.  Otherwise,
+   it returns None. If the input attribute is [Check], this issues a warning. *)
+val zero_alloc_attribute_only_assume_allowed :
+  zero_alloc_attribute -> zero_alloc_assume option
+
+val assume_zero_alloc : zero_alloc_assume -> Zero_alloc_utils.Assume_info.t
 
 type tracing_probe =
   { name : string;