Refactor rustc_compile_action to return dict of provider types to providers

Author: brianduff

extensions/prost/private/prost.bzl

@@ -125,25 +125,22 @@ def _compile_proto(
     return lib_rs, package_info_file
 
 def _get_crate_info(providers):
-    """Finds the CrateInfo provider in the list of providers."""
-    for provider in providers:
-        if hasattr(provider, "name"):
-            return provider
-    fail("Couldn't find a CrateInfo in the list of providers")
+    """Finds the CrateInfo provider in the dict of providers."""
+    if rust_common.crate_info in providers:
+        return providers[rust_common.crate_info]
+    fail("Couldn't find a CrateInfo in the providers")
 
 def _get_dep_info(providers):
-    """Finds the DepInfo provider in the list of providers."""
-    for provider in providers:
-        if hasattr(provider, "direct_crates"):
-            return provider
-    fail("Couldn't find a DepInfo in the list of providers")
+    """Finds the DepInfo provider in the dict of providers."""
+    if rust_common.dep_info in providers:
+        return providers[rust_common.dep_info]
+    fail("Couldn't find a DepInfo in the providers")
 
 def _get_cc_info(providers):
-    """Finds the CcInfo provider in the list of providers."""
-    for provider in providers:
-        if hasattr(provider, "linking_context"):
-            return provider
-    fail("Couldn't find a CcInfo in the list of providers")
+    """Finds the CcInfo provider in the dict of providers."""
+    if CcInfo in providers:
+        return providers[CcInfo]
+    fail("Couldn't find a CcInfo in the providers")
 
 def _compile_rust(
         *,

extensions/protobuf/proto.bzl

@@ -246,8 +246,10 @@ def _rust_proto_compile(protos, descriptor_sets, imports, crate_name, ctx, is_gr
         ),
         output_hash = output_hash,
     )
-    providers.append(OutputGroupInfo(rust_generated_srcs = srcs))
-    return providers
+    # Convert dict to list and add OutputGroupInfo
+    result = list(providers.values())
+    result.append(OutputGroupInfo(rust_generated_srcs = srcs))
+    return result
 
 def _rust_protogrpc_library_impl(ctx, is_grpc):
     """Implementation of the rust_(proto|grpc)_library.

extensions/wasm_bindgen/private/wasm_bindgen_test.bzl

@@ -173,18 +173,21 @@ def _rust_wasm_bindgen_test_impl(ctx):
     # Force the use of a browser for now as there is no node integration.
     env["WASM_BINDGEN_USE_BROWSER"] = "1"
 
+    # Extract DefaultInfo and create a modified version with the wrapper as executable
+    default_info = crate_providers[DefaultInfo]
+    files = default_info.files.to_list()
+    if len(files) != 1:
+        fail("Unexpected number of output files for `{}`: {}".format(ctx.label, files))
+    wasm_file = files[0]
+    env["TEST_WASM_BINARY"] = _rlocationpath(files[0], ctx.workspace_name)
+
+    # Build providers list from the dict, replacing DefaultInfo with our modified version
     providers = []
-
-    for prov in crate_providers:
-        if type(prov) == "DefaultInfo":
-            files = prov.files.to_list()
-            if len(files) != 1:
-                fail("Unexpected number of output files for `{}`: {}".format(ctx.label, files))
-            wasm_file = files[0]
-            env["TEST_WASM_BINARY"] = _rlocationpath(files[0], ctx.workspace_name)
+    for key, prov in crate_providers.items():
+        if key == DefaultInfo:
             providers.append(DefaultInfo(
-                files = prov.files,
-                runfiles = prov.default_runfiles.merge(ctx.runfiles(files = [wasm_file], transitive_files = wb_toolchain.all_test_files)),
+                files = default_info.files,
+                runfiles = default_info.default_runfiles.merge(ctx.runfiles(files = [wasm_file], transitive_files = wb_toolchain.all_test_files)),
                 executable = wrapper,
             ))
         else:

rust/private/rust.bzl

@@ -218,7 +218,7 @@ def _rust_library_common(ctx, crate_type):
     deps = transform_deps(ctx.attr.deps)
     proc_macro_deps = transform_deps(ctx.attr.proc_macro_deps + get_import_macro_deps(ctx))
 
-    return rustc_compile_action(
+    return list(rustc_compile_action(
         ctx = ctx,
         attr = ctx.attr,
         toolchain = toolchain,
@@ -246,7 +246,7 @@ def _rust_library_common(ctx, crate_type):
             owner = ctx.label,
             cfgs = _collect_cfgs(ctx, toolchain, crate_root, crate_type, crate_is_test = False),
         ),
-    )
+    ).values())
 
 def _rust_binary_impl(ctx):
     """The implementation of the `rust_binary` rule
@@ -312,16 +312,16 @@ def _rust_binary_impl(ctx):
         ),
     )
 
-    providers.append(RunEnvironmentInfo(
+    providers[RunEnvironmentInfo] = RunEnvironmentInfo(
         environment = expand_dict_value_locations(
             ctx,
             ctx.attr.env,
             ctx.attr.data,
             {},
         ),
-    ))
+    )
 
-    return providers
+    return list(providers.values())
 
 def get_rust_test_flags(attr):
     """Determine the desired rustc flags for test targets.
@@ -520,19 +520,14 @@ def _rust_test_impl(ctx):
     # If sharding is enabled and we're using libtest harness, wrap the test binary
     # with a script that handles test enumeration and shard partitioning
     if ctx.attr.experimental_enable_sharding and ctx.attr.use_libtest_harness:
-        # Find DefaultInfo and CrateInfo in the providers list
-        # DefaultInfo is first, CrateInfo follows (or TestCrateInfo for staticlib/cdylib)
-        default_info = providers[0]
-        default_info_index = 0
-
-        # Get the test binary from CrateInfo - it's the output of the compiled test
-        crate_info_provider = None
-        for p in providers:
-            if hasattr(p, "output") and hasattr(p, "is_test"):
-                crate_info_provider = p
-                break
+        default_info = providers[DefaultInfo]
 
+        # Get the test binary from CrateInfo (or TestCrateInfo for staticlib/cdylib)
+        crate_info_provider = providers.get(rust_common.crate_info) or providers.get(rust_common.test_crate_info)
         if crate_info_provider:
+            # TestCrateInfo wraps the actual CrateInfo
+            if hasattr(crate_info_provider, "crate"):
+                crate_info_provider = crate_info_provider.crate
             test_binary = crate_info_provider.output
 
             # Select the appropriate wrapper template based on target OS
@@ -559,7 +554,7 @@ def _rust_test_impl(ctx):
             )
 
             # Replace DefaultInfo with wrapper as executable
-            providers[default_info_index] = DefaultInfo(
+            providers[DefaultInfo] = DefaultInfo(
                 files = default_info.files,
                 runfiles = new_runfiles,
                 executable = wrapper,
@@ -593,12 +588,12 @@ def _rust_test_impl(ctx):
         env["RUST_LLVM_PROFDATA"] = llvm_profdata_path
     components = "{}/{}".format(ctx.label.workspace_root, ctx.label.package).split("/")
     env["CARGO_MANIFEST_DIR"] = "/".join([c for c in components if c])
-    providers.append(RunEnvironmentInfo(
+    providers[RunEnvironmentInfo] = RunEnvironmentInfo(
         environment = env,
         inherited_environment = ctx.attr.env_inherit,
-    ))
+    )
 
-    return providers
+    return list(providers.values())
 
 def _rust_library_group_impl(ctx):
     dep_variant_infos = []

rust/private/rustc.bzl

@@ -1301,10 +1301,14 @@ def rustc_compile_action(
         include_coverage (bool, optional): Whether to generate coverage information or not.
 
     Returns:
-        list: A list of the following providers:
-            - (CrateInfo): info for the crate we just built; same as `crate_info` parameter.
-            - (DepInfo): The transitive dependencies of this crate.
-            - (DefaultInfo): The output file for this crate, and its runfiles.
+        dict: A dict mapping provider types to provider instances. Keys include:
+            - DefaultInfo: The output file for this crate, and its runfiles.
+            - CrateInfo: info for the crate we just built (or TestCrateInfo for staticlib/cdylib).
+            - DepInfo: The transitive dependencies of this crate.
+            - InstrumentedFilesInfo: Coverage information (if include_coverage is True).
+            - CcInfo: C/C++ interop info (if applicable).
+            - OutputGroupInfo: Additional output groups (if any).
+        Callers should convert to a list via `list(providers.values())` when returning from a rule.
     """
     deps = crate_info_dict.pop("deps")
     proc_macro_deps = crate_info_dict.pop("proc_macro_deps")
@@ -1695,25 +1699,23 @@ def rustc_compile_action(
             "metadata_files": coverage_runfiles + [executable] if executable else [],
         })
 
-    providers = [
-        DefaultInfo(
+    providers = {
+        DefaultInfo: DefaultInfo(
             # nb. This field is required for cc_library to depend on our output.
             files = depset(outputs),
             runfiles = runfiles,
             executable = executable,
         ),
-    ]
+    }
 
     # When invoked by aspects (and when running `bazel coverage`), the
     # baseline_coverage.dat created here will conflict with the baseline_coverage.dat of the
     # underlying target, which is a build failure. So we add an option to disable it so that this
     # function can be invoked from aspects for rules that have its own InstrumentedFilesInfo.
     if include_coverage:
-        providers.append(
-            coverage_common.instrumented_files_info(
-                ctx,
-                **instrumented_files_kwargs
-            ),
+        providers["InstrumentedFilesInfo"] = coverage_common.instrumented_files_info(
+            ctx,
+            **instrumented_files_kwargs
         )
 
     if crate_info_dict != None:
@@ -1732,11 +1734,20 @@ def rustc_compile_action(
         # as such they shouldn't provide a CrateInfo. However, one may still want to
         # write a rust_test for them, so we provide the CrateInfo wrapped in a provider
         # that rust_test understands.
-        providers.extend([rust_common.test_crate_info(crate = crate_info), dep_info])
+        providers[rust_common.test_crate_info] = rust_common.test_crate_info(crate = crate_info)
     else:
-        providers.extend([crate_info, dep_info])
+        providers[rust_common.crate_info] = crate_info
+
+    providers[rust_common.dep_info] = dep_info
 
-    providers += establish_cc_info(ctx, attr, crate_info, toolchain, cc_toolchain, feature_configuration, interface_library)
+    cc_info_providers = establish_cc_info(ctx, attr, crate_info, toolchain, cc_toolchain, feature_configuration, interface_library)
+    for cc_provider in cc_info_providers:
+        # establish_cc_info returns CcInfo and optionally AllocatorLibrariesImplInfo
+        if type(cc_provider) == "CcInfo":
+            providers[CcInfo] = cc_provider
+        else:
+            # AllocatorLibrariesImplInfo
+            providers[AllocatorLibrariesImplInfo] = cc_provider
 
     output_group_info = {}
 
@@ -1752,12 +1763,12 @@ def rustc_compile_action(
         output_group_info["rustc_output"] = depset([rustc_output])
 
     if output_group_info:
-        providers.append(OutputGroupInfo(**output_group_info))
+        providers[OutputGroupInfo] = OutputGroupInfo(**output_group_info)
 
     # A bit unfortunate, but sidecar the lints info so rustdoc can access the
     # set of lints from the target it is documenting.
     if hasattr(ctx.attr, "lint_config") and ctx.attr.lint_config:
-        providers.append(ctx.attr.lint_config[LintsInfo])
+        providers[LintsInfo] = ctx.attr.lint_config[LintsInfo]
 
     return providers
 

test/unit/consistent_crate_name/with_modified_crate_name.bzl

@@ -31,7 +31,7 @@ def _with_modified_crate_name_impl(ctx):
     ) for dep in ctx.attr.deps]
 
     rust_lib = ctx.actions.declare_file(rust_lib_name)
-    return rustc_compile_action(
+    return list(rustc_compile_action(
         ctx = ctx,
         attr = ctx.attr,
         toolchain = toolchain,
@@ -52,7 +52,7 @@ def _with_modified_crate_name_impl(ctx):
             is_test = False,
         ),
         output_hash = output_hash,
-    )
+    ).values())
 
 with_modified_crate_name = rule(
     implementation = _with_modified_crate_name_impl,

test/unit/force_all_deps_direct/generator.bzl

@@ -47,7 +47,7 @@ def _generator_impl(ctx):
     ) for dep in ctx.attr.deps]
 
     rust_lib = ctx.actions.declare_file(rust_lib_name)
-    return rustc_compile_action(
+    return list(rustc_compile_action(
         ctx = ctx,
         attr = ctx.attr,
         toolchain = toolchain,
@@ -69,7 +69,7 @@ def _generator_impl(ctx):
         ),
         output_hash = output_hash,
         force_all_deps_direct = True,
-    )
+    ).values())
 
 generator = rule(
     implementation = _generator_impl,

test/unit/pipelined_compilation/wrap.bzl

@@ -59,7 +59,7 @@ def _wrap_impl(ctx):
     rust_metadata = None
     if ctx.attr.generate_metadata:
         rust_metadata = ctx.actions.declare_file(rust_metadata_name)
-    return rustc_compile_action(
+    return list(rustc_compile_action(
         ctx = ctx,
         attr = ctx.attr,
         toolchain = toolchain,
@@ -83,7 +83,7 @@ def _wrap_impl(ctx):
             is_test = False,
         ),
         output_hash = output_hash,
-    )
+    ).values())
 
 wrap = rule(
     implementation = _wrap_impl,