Allow passing default_roots to rust_test to disambiguate

Summary:
In the general case, it's ambiguous whether a usage of `rust_test()` is to declare a test target for a *library* or for a *binary*. This makes it tricky to determine the crate root as we must consider both `lib.rs` and `main.rs`. However, for our use-case, we tend to add unit test targets as part of a `rust_library` or `rust_binary` macro, where we have the information to disambiguate.

This change proposes exposing `default_roots` on `rust_test` so that we can pass the correct root candidate for our automatic test targets.

See also https://fb.workplace.com/groups/rustonmobile/posts/3858526464393667

Reviewed By: dtolnay

Differential Revision: D62585839

fbshipit-source-id: d00b890d7b051e92ce4d16e388fe520ef05a8248

prelude/decls/rust_common.bzl

@@ -95,6 +95,15 @@ def _crate_root():
 """),
     }
 
+def _default_roots_arg():
+    return {
+        "default_roots": attrs.option(attrs.list(attrs.string()), default = None, doc = """
+    Set the candidate source names to consider for crate root. Typically used to disambiguate between
+     lib.rs or main.rs for rust_test, which may be declare a test suite for either library or binary
+     rules. Has no effect if an explicit `crate_root` is provided.
+"""),
+    }
+
 def _env_arg():
     return {
         "env": attrs.dict(key = attrs.string(), value = attrs.arg(), sorted = False, default = {}, doc = """
@@ -172,6 +181,7 @@ rust_common = struct(
     exported_linker_flags_arg = _exported_linker_flags_arg,
     crate = _crate,
     crate_root = _crate_root,
+    default_roots_arg = _default_roots_arg,
     env_arg = _env_arg,
     run_env_arg = _run_env_arg,
     build_and_run_env_arg = _build_and_run_env_arg,

prelude/decls/rust_rules.bzl

@@ -260,6 +260,7 @@ rust_test = prelude_rule(
         rust_common.rustc_flags_arg() |
         rust_common.crate(crate_type = attrs.option(attrs.string(), default = None)) |
         rust_common.crate_root() |
+        rust_common.default_roots_arg() |
         rust_common.run_env_arg() |
         rust_common.build_and_run_env_arg() |
         _rust_binary_attrs_group(prefix = "") |

prelude/rust/rust_binary.bzl

@@ -423,7 +423,10 @@ def rust_test_impl(ctx: AnalysisContext) -> list[Provider]:
     providers, args = _rust_binary_common(
         ctx = ctx,
         compile_ctx = compile_ctx,
-        default_roots = ["main.rs", "lib.rs"],
+        # Unless default_roots are provided, it is ambiguous whether this test rule is invoked
+        # to test a binary, or to test a library. As such, we must consider both main.rs and
+        # lib.rs as potential candidates.
+        default_roots = ctx.attrs.default_roots or ["main.rs", "lib.rs"],
         extra_flags = extra_flags,
         allow_cache_upload = False,
     )