flambda-backend: Add a test for an interaction between omitted mli and overeager heap allocation of argument (#1816)

* Add tests

These tests show how the defaulting of mode variables
affect whether objects can be stack allocated in the
presence/absence of an mli.

* Add documentation

Author: ncik-roberts

jane/doc/local-reference.md

@@ -79,6 +79,35 @@ function types, see below) then inference will resolve them according to what
 appears in the `.mli`. If there is no `.mli` file, then inference will always
 choose `global` for anything that can be accessed from another file.
 
+Local annotations (or the lack thereof) in the mli don't affect inference
+within the ml. In the below example, the `~foo` parameter is inferred to
+be local internally to `A`, so `foo:(Some x)` can be constructed locally.
+
+```ocaml
+(* in a.mli *)
+val f1 : foo:local_ int option -> unit
+val f2 : int -> unit
+
+(* in a.ml *)
+let f1 ~foo:_ = ()
+let f2 x = f1 ~foo:(Some x) (* [Some x] is stack allocated *)
+```
+
+<!-- See Note [Inference affects allocation in mli-less files] in [ocaml/testsuite/tests/typing-local/alloc_arg_with_mli.ml]
+     in the flambda-backend Git repo. The ensuing paragraph is related to that
+     note; we can remove this comment when the note is resolved.
+-->
+However, a missing mli *does* affect inference within the ml. As a conservative rule of thumb,
+function arguments in an mli-less file will be heap-allocated unless the function parameter or argument
+is annotated with `local_`. This is due to an implementation detail of the type-checker and
+is not fundamental, but for now, it's yet another reason to prefer writing mlis.
+
+```ocaml
+(* in a.ml; a.mli is missing *)
+let f1 ~foo:_ = ()
+let f2 x = f1 ~foo:(Some x) (* [Some x] is heap allocated *)
+```
+
 ## Regions
 
 Every local allocation takes places inside a _region_, which is a block of code

ocamltest/ocaml_actions.ml

@@ -291,6 +291,11 @@ let compile_module compiler module_ log env =
   let is_c = is_c_file module_with_filetype in
   let c_headers_flags =
     if is_c then Ocaml_flags.c_includes else "" in
+  let compile_only_flag_opt =
+    if Environments.lookup_as_bool Ocaml_variables.compile_only env = Some false
+    then ""
+    else " -c "
+  in
   let commandline =
   [
     compiler#name;
@@ -301,7 +306,8 @@ let compile_module compiler module_ log env =
     libraries compiler#target env;
     backend_default_flags env compiler#target;
     backend_flags env compiler#target;
-    "-c " ^ module_;
+    compile_only_flag_opt;
+    module_;
   ] in
   let exit_status =
     Actions_helpers.run_cmd

testsuite/tests/typing-local/alloc_arg_with_mli.ml

@@ -0,0 +1,138 @@
+(* TEST
+   * stack-allocation
+   ** native
+ *)
+
+(* Check whether functions that *could* take their argument
+   locally allow their callers to locally construct the argument.
+
+   Among other things, this checks how mode variables are defaulted in the
+   presence of an mli.
+
+   See the [..._no_mli.ml] version of this test for how mode variables
+   are defaulted in the absence of an mli.
+ *)
+
+let check_if_zero_alloc =
+  let measure_words f =
+    let before = Gc.allocated_bytes () in
+    ignore (f () : _);
+    let after  = Gc.allocated_bytes () in
+    int_of_float (after -. before) / (Sys.word_size / 8)
+  in
+  fun[@inline never] ~name ~f x ->
+    let words = measure_words (fun () -> f x) - measure_words ignore in
+    Printf.printf "%s: %d words%s\n" name words
+      (if words = 0 then "" else " (allocates!)")
+
+external escape : 'a -> 'a = "%identity"
+
+type t = { a : int; b : int }
+
+module _ : sig end = struct
+  let[@inline never] take_unrestricted { a; b } = a + b
+
+  let () =
+    check_if_zero_alloc ~name:"take unrestricted of global (not exposed)" 0 ~f:(fun x ->
+      take_unrestricted { a = x; b = x } [@nontail])
+end
+
+module _ : sig end = struct
+  let[@inline never] take_unrestricted { a; b } = a + b
+
+  let () =
+    check_if_zero_alloc ~name:"take unrestricted of local (not exposed)" 0 ~f:(fun x ->
+      take_unrestricted (local_ { a = x; b = x }) [@nontail])
+end
+
+module _ : sig end = struct
+  let[@inline never] take_local (local_ { a; b }) = a + b
+
+  let () =
+    check_if_zero_alloc ~name:"take local of local (not exposed)" 0 ~f:(fun x ->
+      take_local (local_ { a = x; b = x }) [@nontail])
+end
+
+module _ : sig end = struct
+  let[@inline never] take_local (local_ { a; b }) = a + b
+
+  let () =
+    check_if_zero_alloc ~name:"take local of global (not exposed)" 0 ~f:(fun x ->
+      take_local { a = x; b = x } [@nontail])
+end
+
+module _ : sig end = struct
+  let[@inline never] take_global ({ a; b } as t) = ignore (escape t); a + b
+
+  let () =
+    check_if_zero_alloc
+      ~name:"take global of global (not exposed; expected to allocate)"
+      0
+      ~f:(fun x -> take_global { a = x; b = x } [@nontail])
+end
+
+module M1 = struct
+  let[@inline never] take_unrestricted { a; b } = a + b
+
+  (* Note [Inference affects allocation in mli-less files]
+     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+     At the moment, this allocates in the without-mli case, but not for a
+     fundamental reason: when the type signature is inferred, the inferred
+     types and thus mode variables are reused. Instead, we could create new
+     copies of types with fresh mode variables related by submoding.
+   *)
+  let () =
+    check_if_zero_alloc ~name:"take unrestricted of global (exposed)" 0 ~f:(fun x ->
+      take_unrestricted { a = x; b = x } [@nontail])
+end
+
+module M2 = struct
+  let[@inline never] take_unrestricted { a; b } = a + b
+
+  let () =
+    check_if_zero_alloc ~name:"take unrestricted of local (exposed)" 0 ~f:(fun x ->
+      take_unrestricted (local_ { a = x; b = x }) [@nontail])
+end
+
+module M3 = struct
+  let[@inline never] take_local (local_ { a; b }) = a + b
+
+  let () =
+    check_if_zero_alloc ~name:"take local of local (exposed)" 0 ~f:(fun x ->
+      take_local (local_ { a = x; b = x }) [@nontail])
+end
+
+module M4 = struct
+  let[@inline never] take_local (local_ { a; b }) = a + b
+
+  let () =
+    check_if_zero_alloc ~name:"take local of global (exposed)" 0 ~f:(fun x ->
+      take_local { a = x; b = x } [@nontail])
+end
+
+module M5 = struct
+  let[@inline never] take_global ({ a; b } as t) = ignore (escape t); a + b
+
+  let () =
+    check_if_zero_alloc
+      ~name:"take global of global (exposed; expected to allocate)"
+      0
+      ~f:(fun x -> take_global { a = x; b = x } [@nontail])
+end
+
+module M6 = struct
+  let[@inline never] take_local__global_in_mli (local_ { a; b }) = a + b
+
+  let () =
+    check_if_zero_alloc ~name:"take local of local (exposed)" 0 ~f:(fun x ->
+      take_local__global_in_mli (local_ { a = x; b = x }) [@nontail])
+end
+
+module M7 = struct
+  let[@inline never] take_local__global_in_mli (local_ { a; b }) = a + b
+
+  let () =
+    check_if_zero_alloc ~name:"take local of global (exposed)" 0 ~f:(fun x ->
+      take_local__global_in_mli { a = x; b = x } [@nontail])
+end

testsuite/tests/typing-local/alloc_arg_with_mli.mli

@@ -0,0 +1,29 @@
+type t = { a : int; b : int }
+
+module M1 : sig
+  val take_unrestricted : t -> int
+end
+
+module M2 : sig
+  val take_unrestricted : t -> int
+end
+
+module M3 : sig
+  val take_local : local_ t -> int
+end
+
+module M4 : sig
+  val take_local : local_ t -> int
+end
+
+module M5 : sig
+  val take_global : t -> int
+end
+
+module M6 : sig
+  val take_local__global_in_mli : t -> int
+end
+
+module M7 : sig
+  val take_local__global_in_mli : t -> int
+end

testsuite/tests/typing-local/alloc_arg_with_mli.reference

@@ -0,0 +1,12 @@
+take unrestricted of global (not exposed): 0 words
+take unrestricted of local (not exposed): 0 words
+take local of local (not exposed): 0 words
+take local of global (not exposed): 0 words
+take global of global (not exposed; expected to allocate): 3 words (allocates!)
+take unrestricted of global (exposed): 0 words
+take unrestricted of local (exposed): 0 words
+take local of local (exposed): 0 words
+take local of global (exposed): 0 words
+take global of global (exposed; expected to allocate): 3 words (allocates!)
+take local of local (exposed): 0 words
+take local of global (exposed): 0 words

testsuite/tests/typing-local/alloc_arg_without_mli.ml

@@ -0,0 +1,20 @@
+(* TEST
+
+readonly_files = "alloc_arg_with_mli.ml"
+
+* stack-allocation
+** native
+compile_only = "false"
+flags = "-o ${test_build_directory}/alloc_arg_without_mli.opt"
+module = "alloc_arg_with_mli.ml"
+*)
+
+(* Check whether functions that *could* take their argument
+   locally allow their callers to locally construct the argument.
+
+   Among other things, this checks how mode variables are defaulted in the
+   absence of an mli.
+
+   See the [..._with_mli.ml] version of this test for how mode variables
+   are defaulted in the presence of an mli.
+ *)

testsuite/tests/typing-local/alloc_arg_without_mli.reference

@@ -0,0 +1,12 @@
+take unrestricted of global (not exposed): 0 words
+take unrestricted of local (not exposed): 0 words
+take local of local (not exposed): 0 words
+take local of global (not exposed): 0 words
+take global of global (not exposed; expected to allocate): 3 words (allocates!)
+take unrestricted of global (exposed): 3 words (allocates!)
+take unrestricted of local (exposed): 0 words
+take local of local (exposed): 0 words
+take local of global (exposed): 0 words
+take global of global (exposed; expected to allocate): 3 words (allocates!)
+take local of local (exposed): 0 words
+take local of global (exposed): 0 words