Parse zero_alloc attributes into the typed tree (ocaml-flambda#2400)

Check for [@zero_alloc] earlier, record it in the typed tree

backend/checkmach.ml

@@ -335,17 +335,17 @@ end = struct
      Assume_info due to cyclic dependencies. The witnesses in Assume_info are
      always empty and the translation is trivial. Is there a better way to avoid
      duplicating [Zero_alloc_utils]? *)
-  let transl w (v : Assume_info.V.t) : V.t =
+  let transl w (v : Zero_alloc_utils.Assume_info.V.t) : V.t =
     match v with Top _ -> Top w | Safe -> Safe | Bot -> Bot
 
-  let transl w (v : Assume_info.Value.t) : Value.t =
+  let transl w (v : Zero_alloc_utils.Assume_info.Value.t) : Value.t =
     { nor = transl w v.nor; exn = transl w v.exn; div = transl w v.div }
 
   let assume_value dbg ~can_raise w =
     (* [loc] can be obtained by [Debuginfo.to_location dbg], For now just return
        [Location.none] because it is not used. *)
     let a = Debuginfo.assume_zero_alloc dbg in
-    match Assume_info.get_value a with
+    match Zero_alloc_utils.Assume_info.get_value a with
     | None -> None
     | Some v ->
       let v = transl w v in

chamelon/compat.jst.ml

@@ -17,15 +17,19 @@ let mkTexp_ident ?id:(ident_kind, uu = (Id_value, shared_many_use))
   Texp_ident (path, longident, vd, ident_kind, uu)
 
 type nonrec apply_arg = apply_arg
-type texp_apply_identifier = apply_position * Locality.l
+
+type texp_apply_identifier =
+  apply_position * Locality.l * Zero_alloc_utils.Assume_info.t
 
 let mkTexp_apply
-    ?id:(pos, mode = (Default, Locality.disallow_right Locality.legacy))
-    (exp, args) =
+    ?id:(pos, mode, za =
+        ( Default,
+          Locality.disallow_right Locality.legacy,
+          Zero_alloc_utils.Assume_info.none )) (exp, args) =
   let args =
     List.map (fun (label, x) -> (Typetexp.transl_label label None, x)) args
   in
-  Texp_apply (exp, args, pos, mode)
+  Texp_apply (exp, args, pos, mode, za)
 
 type texp_tuple_identifier = string option list * Alloc.r
 
@@ -86,6 +90,7 @@ type texp_function_identifier = {
   ret_sort : Jkind.sort;
   region : bool;
   ret_mode : Alloc.l;
+  zero_alloc : Builtin_attributes.check_attribute;
 }
 
 let texp_function_cases_identifier_defaults =
@@ -110,6 +115,7 @@ let texp_function_defaults =
     ret_sort = Jkind.Sort.value;
     ret_mode = Alloc.disallow_right Alloc.legacy;
     region = false;
+    zero_alloc = Builtin_attributes.Default_check;
   }
 
 let mkTexp_function ?(id = texp_function_defaults)
@@ -163,6 +169,7 @@ let mkTexp_function ?(id = texp_function_defaults)
       region = id.region;
       ret_sort = id.ret_sort;
       ret_mode = id.ret_mode;
+      zero_alloc = id.zero_alloc;
     }
 
 type texp_sequence_identifier = Jkind.sort
@@ -206,15 +213,16 @@ let view_texp (e : expression_desc) =
   match e with
   | Texp_ident (path, longident, vd, ident_kind, uu) ->
       Texp_ident (path, longident, vd, (ident_kind, uu))
-  | Texp_apply (exp, args, pos, mode) ->
+  | Texp_apply (exp, args, pos, mode, za) ->
       let args = List.map (fun (label, x) -> (untype_label label, x)) args in
-      Texp_apply (exp, args, (pos, mode))
+      Texp_apply (exp, args, (pos, mode, za))
   | Texp_construct (name, desc, args, mode) ->
       Texp_construct (name, desc, args, mode)
   | Texp_tuple (args, mode) ->
       let labels, args = List.split args in
       Texp_tuple (args, (labels, mode))
-  | Texp_function { params; body; alloc_mode; region; ret_sort; ret_mode } ->
+  | Texp_function
+      { params; body; alloc_mode; region; ret_sort; ret_mode; zero_alloc } ->
       let params =
         List.map
           (fun param ->
@@ -259,7 +267,8 @@ let view_texp (e : expression_desc) =
               }
       in
       Texp_function
-        ({ params; body }, { alloc_mode; region; ret_sort; ret_mode })
+        ( { params; body },
+          { alloc_mode; region; ret_sort; ret_mode; zero_alloc } )
   | Texp_sequence (e1, sort, e2) -> Texp_sequence (e1, e2, sort)
   | Texp_match (e, sort, cases, partial) -> Texp_match (e, cases, partial, sort)
   | _ -> O e

middle_end/flambda2/terms/check_attribute.ml

@@ -50,7 +50,7 @@ let from_lambda : Lambda.check_attribute -> Location.t -> t =
   match a with
   | Default_check ->
     if !Clflags.zero_alloc_check_assert_all
-       && Lambda.is_check_enabled ~opt:false Zero_alloc
+       && Builtin_attributes.is_check_enabled ~opt:false Zero_alloc
     then Check { property = Zero_alloc; strict = false; loc }
     else Default_check
   | Ignore_assert_all Zero_alloc -> Default_check
@@ -62,7 +62,7 @@ let from_lambda : Lambda.check_attribute -> Location.t -> t =
         loc
       }
   | Check { property; strict; opt; loc } ->
-    if Lambda.is_check_enabled ~opt property
+    if Builtin_attributes.is_check_enabled ~opt property
     then Check { property = Property.from_lambda property; strict; loc }
     else Default_check
 

ocaml/.depend

@@ -400,6 +400,7 @@ parsing/attr_helper.cmi : \
     parsing/location.cmi \
     parsing/asttypes.cmi
 parsing/builtin_attributes.cmo : \
+    utils/zero_alloc_utils.cmi \
     utils/warnings.cmi \
     parsing/parsetree.cmi \
     utils/misc.cmi \
@@ -412,6 +413,7 @@ parsing/builtin_attributes.cmo : \
     parsing/ast_helper.cmi \
     parsing/builtin_attributes.cmi
 parsing/builtin_attributes.cmx : \
+    utils/zero_alloc_utils.cmx \
     utils/warnings.cmx \
     parsing/parsetree.cmi \
     utils/misc.cmx \
@@ -424,6 +426,7 @@ parsing/builtin_attributes.cmx : \
     parsing/ast_helper.cmx \
     parsing/builtin_attributes.cmi
 parsing/builtin_attributes.cmi : \
+    utils/zero_alloc_utils.cmi \
     parsing/parsetree.cmi \
     utils/misc.cmi \
     parsing/location.cmi
@@ -1452,6 +1455,7 @@ typing/printtyp.cmi : \
     typing/env.cmi \
     parsing/asttypes.cmi
 typing/printtyped.cmo : \
+    utils/zero_alloc_utils.cmi \
     typing/types.cmi \
     typing/typedtree.cmi \
     parsing/printast.cmi \
@@ -1468,6 +1472,7 @@ typing/printtyped.cmo : \
     parsing/asttypes.cmi \
     typing/printtyped.cmi
 typing/printtyped.cmx : \
+    utils/zero_alloc_utils.cmx \
     typing/types.cmx \
     typing/typedtree.cmx \
     parsing/printast.cmx \
@@ -1731,6 +1736,7 @@ typing/typeclass.cmi : \
     typing/ctype.cmi \
     parsing/asttypes.cmi
 typing/typecore.cmo : \
+    utils/zero_alloc_utils.cmi \
     utils/warnings.cmi \
     typing/untypeast.cmi \
     typing/uniqueness_analysis.cmi \
@@ -1773,6 +1779,7 @@ typing/typecore.cmo : \
     parsing/ast_helper.cmi \
     typing/typecore.cmi
 typing/typecore.cmx : \
+    utils/zero_alloc_utils.cmx \
     utils/warnings.cmx \
     typing/untypeast.cmx \
     typing/uniqueness_analysis.cmx \
@@ -1994,6 +2001,7 @@ typing/typedecl_variance.cmi : \
     typing/env.cmi \
     parsing/asttypes.cmi
 typing/typedtree.cmo : \
+    utils/zero_alloc_utils.cmi \
     typing/types.cmi \
     typing/shape.cmi \
     typing/primitive.cmi \
@@ -2007,9 +2015,11 @@ typing/typedtree.cmo : \
     parsing/jane_asttypes.cmi \
     typing/ident.cmi \
     typing/env.cmi \
+    parsing/builtin_attributes.cmi \
     parsing/asttypes.cmi \
     typing/typedtree.cmi
 typing/typedtree.cmx : \
+    utils/zero_alloc_utils.cmx \
     typing/types.cmx \
     typing/shape.cmx \
     typing/primitive.cmx \
@@ -2023,9 +2033,11 @@ typing/typedtree.cmx : \
     parsing/jane_asttypes.cmx \
     typing/ident.cmx \
     typing/env.cmx \
+    parsing/builtin_attributes.cmx \
     parsing/asttypes.cmi \
     typing/typedtree.cmi
 typing/typedtree.cmi : \
+    utils/zero_alloc_utils.cmi \
     typing/types.cmi \
     typing/shape.cmi \
     typing/primitive.cmi \
@@ -2039,6 +2051,7 @@ typing/typedtree.cmi : \
     parsing/jane_asttypes.cmi \
     typing/ident.cmi \
     typing/env.cmi \
+    parsing/builtin_attributes.cmi \
     parsing/asttypes.cmi
 typing/typemod.cmo : \
     utils/warnings.cmi \
@@ -3958,36 +3971,28 @@ middle_end/variable.cmi : \
     utils/identifiable.cmi \
     typing/ident.cmi \
     utils/compilation_unit.cmi
-lambda/assume_info.cmo : \
-    utils/zero_alloc_utils.cmi \
-    lambda/assume_info.cmi
-lambda/assume_info.cmx : \
-    utils/zero_alloc_utils.cmx \
-    lambda/assume_info.cmi
-lambda/assume_info.cmi : \
-    utils/zero_alloc_utils.cmi
 lambda/debuginfo.cmo : \
+    utils/zero_alloc_utils.cmi \
     parsing/location.cmi \
     utils/int_replace_polymorphic_compare.cmi \
     typing/ident.cmi \
     utils/compilation_unit.cmi \
     parsing/asttypes.cmi \
-    lambda/assume_info.cmi \
     lambda/debuginfo.cmi
 lambda/debuginfo.cmx : \
+    utils/zero_alloc_utils.cmx \
     parsing/location.cmx \
     utils/int_replace_polymorphic_compare.cmx \
     typing/ident.cmx \
     utils/compilation_unit.cmx \
     parsing/asttypes.cmi \
-    lambda/assume_info.cmx \
     lambda/debuginfo.cmi
 lambda/debuginfo.cmi : \
+    utils/zero_alloc_utils.cmi \
     parsing/location.cmi \
     typing/ident.cmi \
     utils/compilation_unit.cmi \
-    parsing/asttypes.cmi \
-    lambda/assume_info.cmi
+    parsing/asttypes.cmi
 lambda/lambda.cmo : \
     typing/types.cmi \
     typing/typedtree.cmi \
@@ -4003,6 +4008,7 @@ lambda/lambda.cmo : \
     utils/config.cmi \
     utils/compilation_unit.cmi \
     utils/clflags.cmi \
+    parsing/builtin_attributes.cmi \
     parsing/asttypes.cmi \
     lambda/lambda.cmi
 lambda/lambda.cmx : \
@@ -4020,6 +4026,7 @@ lambda/lambda.cmx : \
     utils/config.cmx \
     utils/compilation_unit.cmx \
     utils/clflags.cmx \
+    parsing/builtin_attributes.cmx \
     parsing/asttypes.cmi \
     lambda/lambda.cmi
 lambda/lambda.cmi : \
@@ -4034,6 +4041,7 @@ lambda/lambda.cmi : \
     typing/env.cmi \
     lambda/debuginfo.cmi \
     utils/compilation_unit.cmi \
+    parsing/builtin_attributes.cmi \
     parsing/asttypes.cmi
 lambda/matching.cmo : \
     typing/types.cmi \
@@ -4249,36 +4257,25 @@ lambda/translattribute.cmo : \
     utils/warnings.cmi \
     typing/typedtree.cmi \
     parsing/parsetree.cmi \
-    utils/misc.cmi \
-    parsing/longident.cmi \
     parsing/location.cmi \
     lambda/lambda.cmi \
     utils/config.cmi \
-    utils/clflags.cmi \
     parsing/builtin_attributes.cmi \
-    parsing/asttypes.cmi \
-    lambda/assume_info.cmi \
     lambda/translattribute.cmi
 lambda/translattribute.cmx : \
     utils/warnings.cmx \
     typing/typedtree.cmx \
     parsing/parsetree.cmi \
-    utils/misc.cmx \
-    parsing/longident.cmx \
     parsing/location.cmx \
     lambda/lambda.cmx \
     utils/config.cmx \
-    utils/clflags.cmx \
     parsing/builtin_attributes.cmx \
-    parsing/asttypes.cmi \
-    lambda/assume_info.cmx \
     lambda/translattribute.cmi
 lambda/translattribute.cmi : \
     typing/typedtree.cmi \
     parsing/parsetree.cmi \
     parsing/location.cmi \
-    lambda/lambda.cmi \
-    lambda/assume_info.cmi
+    lambda/lambda.cmi
 lambda/translclass.cmo : \
     typing/types.cmi \
     typing/typeopt.cmi \
@@ -4323,6 +4320,7 @@ lambda/translclass.cmi : \
     lambda/debuginfo.cmi \
     parsing/asttypes.cmi
 lambda/translcore.cmo : \
+    utils/zero_alloc_utils.cmi \
     typing/types.cmi \
     typing/typeopt.cmi \
     typing/typedtree.cmi \
@@ -4352,11 +4350,12 @@ lambda/translcore.cmo : \
     typing/ctype.cmi \
     utils/config.cmi \
     utils/clflags.cmi \
+    parsing/builtin_attributes.cmi \
     typing/btype.cmi \
     parsing/asttypes.cmi \
-    lambda/assume_info.cmi \
     lambda/translcore.cmi
 lambda/translcore.cmx : \
+    utils/zero_alloc_utils.cmx \
     typing/types.cmx \
     typing/typeopt.cmx \
     typing/typedtree.cmx \
@@ -4386,9 +4385,9 @@ lambda/translcore.cmx : \
     typing/ctype.cmx \
     utils/config.cmx \
     utils/clflags.cmx \
+    parsing/builtin_attributes.cmx \
     typing/btype.cmx \
     parsing/asttypes.cmi \
-    lambda/assume_info.cmx \
     lambda/translcore.cmi
 lambda/translcore.cmi : \
     typing/types.cmi \

ocaml/compilerlibs/Makefile.compilerlibs

@@ -125,7 +125,6 @@ TYPING = \
   typing/typedecl_properties.cmo \
   typing/typedecl_variance.cmo \
   typing/typedecl_separability.cmo \
-  lambda/assume_info.cmo \
   lambda/debuginfo.cmo lambda/lambda.cmo \
   typing/typedecl.cmo \
   typing/typeopt.cmo \

ocaml/dune

@@ -99,7 +99,7 @@
    annot outcometree
 
    ;; lambda/
-   assume_info debuginfo lambda matching printlambda runtimedef runtimetags tmc simplif switch
+   debuginfo lambda matching printlambda runtimedef runtimetags tmc simplif switch
    translmode
    transl_comprehension_utils
    transl_array_comprehension transl_list_comprehension
@@ -330,7 +330,6 @@
     (annot.mli as compiler-libs/annot.mli)
     (outcometree.mli as compiler-libs/outcometree.mli)
     (debuginfo.mli as compiler-libs/debuginfo.mli)
-    (assume_info.mli as compiler-libs/assume_info.mli)
     (lambda.mli as compiler-libs/lambda.mli)
     (matching.mli as compiler-libs/matching.mli)
     (printlambda.mli as compiler-libs/printlambda.mli)