Zero alloc: propagate assume on applications of externals

ocaml/lambda/translcore.ml

@@ -410,8 +410,12 @@ and transl_exp0 ~in_new_scope ~scopes sort e =
         else Rc_normal
       in
       let lam =
+        let loc =
+          map_scopes (update_assume_zero_alloc ~assume_zero_alloc)
+            (of_location ~scopes e.exp_loc)
+        in
         Translprim.transl_primitive_application
-          (of_location ~scopes e.exp_loc) p e.exp_env prim_type
+          loc p e.exp_env prim_type
           ~poly_mode:pmode ~poly_sort:psort
           path prim_exp args (List.map fst arg_exps) position
       in

tests/backend/checkmach/dune

@@ -17,3 +17,24 @@
 (rule
  (alias  runtest)
  (action (copy test_all_opt.ml test_all_opt3.ml)))
+
+;; Tests whose outputs differ depending on stack_allocation configuration flag.
+;; This condition is not expressible in "enable_if" clause
+;; because dune does not support %{config:stack_allocation} yet.
+;; This test cannot be promoted automatically.
+
+(rule
+ (enabled_if (= %{context_name} "main"))
+ (targets test_assume_stub.output)
+ (deps
+        test_assume_stub.heap_allocation.output
+        test_assume_stub.stack_allocation.output)
+ (action
+   (with-accepted-exit-codes 0
+     (bash
+        "if %{bin:ocamlopt.opt} -config | grep -q \"stack_allocation: true\" ;
+        then
+           cp test_assume_stub.stack_allocation.output test_assume_stub.output
+        else
+           cp test_assume_stub.heap_allocation.output test_assume_stub.output
+        fi"))))

tests/backend/checkmach/dune.inc

@@ -978,3 +978,22 @@
  (enabled_if (= %{context_name} "main"))
  (deps test_assume_error.output test_assume_error.output.corrected)
  (action (diff test_assume_error.output test_assume_error.output.corrected)))
+
+(rule
+ (enabled_if (= %{context_name} "main"))
+ (targets test_assume_stub.output.corrected)
+ (deps (:ml test_assume_stub.ml) filter.sh)
+ (action
+   (with-outputs-to test_assume_stub.output.corrected
+    (pipe-outputs
+    (with-accepted-exit-codes 2
+     (run %{bin:ocamlopt.opt} %{ml} -g -color never -error-style short -c
+          -zero-alloc-check default -checkmach-details-cutoff 20 -O3))
+    (run "./filter.sh")
+   ))))
+
+(rule
+ (alias   runtest)
+ (enabled_if (= %{context_name} "main"))
+ (deps test_assume_stub.output test_assume_stub.output.corrected)
+ (action (diff test_assume_stub.output test_assume_stub.output.corrected)))

tests/backend/checkmach/gen/gen_dune.ml

@@ -181,4 +181,6 @@ let () =
     ~extra_dep:None ~exit_code:2 "test_assume_inlining";
   print_test_expected_output ~cutoff:default_cutoff
     ~extra_dep:None ~exit_code:2 "test_assume_error";
+  print_test_expected_output ~cutoff:default_cutoff
+    ~extra_dep:None ~exit_code:2 "test_assume_stub";
   ()

tests/backend/checkmach/test_assume_stub.heap_allocation.output

@@ -0,0 +1,11 @@
+File "test_assume_stub.ml", line 20, characters 7-17:
+Error: Annotation check for zero_alloc failed on function Test_assume_stub.A3.overapplied (camlTest_assume_stub.overapplied_HIDE_STAMP)
+
+File "test_assume_stub.ml", line 20, characters 37-74:
+Error: called function may allocate (indirect tailcall)
+
+File "test_assume_stub.ml", line 34, characters 7-17:
+Error: Annotation check for zero_alloc failed on function Test_assume_stub.A4.baz (camlTest_assume_stub.baz_HIDE_STAMP)
+
+File "test_assume_stub.ml", line 35, characters 4-58:
+Error: called function may allocate (indirect tailcall)

tests/backend/checkmach/test_assume_stub.ml

@@ -0,0 +1,44 @@
+module A1 = struct
+  external stub : int -> int = "test_stub"
+  let[@inline always] wrap x = (stub[@zero_alloc assume]) x
+
+  let[@zero_alloc] foo x = wrap x
+end
+
+module A2 = struct
+  external stub : int -> int = "test_stub"
+  let[@inline always][@zero_alloc assume] wrap x = stub x
+
+  let[@zero_alloc] foo x = wrap x
+end
+
+module A3 = struct
+  external revapply : 'a -> ('a -> 'b) -> 'b = "%revapply"
+
+  let[@inline never] id x = x
+
+  let[@zero_alloc] overapplied g x = (revapply[@zero_alloc assume]) g id x
+end
+
+module A4 = struct
+  external opaque_identity : 'a -> 'a = "%opaque"
+
+  let[@zero_alloc] foo _y =
+    (opaque_identity[@zero_alloc assume]) (fun x -> x)
+
+  (* overapply, still succeeds *)
+  let[@zero_alloc] bar y =
+    (opaque_identity[@zero_alloc assume]) (fun x -> x) y
+
+  (* apply, fails as expected *)
+  let[@zero_alloc] baz y =
+    ((opaque_identity[@zero_alloc assume]) (fun x -> x)) y
+end
+
+module A5 = struct
+  external get : 'a array -> int -> 'a = "%array_safe_get"
+
+  let[@zero_alloc] foo a i =
+    (get [@zero_alloc assume]) a (i+1)
+
+end

tests/backend/checkmach/test_assume_stub.stack_allocation.output

@@ -0,0 +1,5 @@
+File "test_assume_stub.ml", line 34, characters 7-17:
+Error: Annotation check for zero_alloc failed on function Test_assume_stub.A4.baz (camlTest_assume_stub.baz_HIDE_STAMP)
+
+File "test_assume_stub.ml", line 35, characters 4-58:
+Error: called function may allocate (indirect tailcall)