Added support for rustc self-profiling (bazelbuild#4073)

Introduced a new build setting,
`--//rust/settings:zself_profile_events`, to enable rustc's
`-Zself-profile` and `-Zself-profile-events` flag. Rust `nightly`
toolchain is required for using this feature.
The setting accepts patterns to match against crate labels or execution
paths, allowing fine-grained control over which crates are profiled.
Profiling data is output to a `<crate_name>_self-profile` directory.

Co-authored-by: Vipul Gupta <ivip@google.com>

Author: 1vipulgupta

rust/private/rust.bzl

@@ -32,6 +32,7 @@ load(
 )
 load(
     ":rustc.bzl",
+    "UnstableSelfProfileInfo",
     "collect_extra_rustc_flags",
     "is_no_std",
     "rustc_compile_action",
@@ -853,6 +854,11 @@ _COMMON_ATTRS = {
         doc = "A version to inject in the cargo environment variable.",
         default = "0.0.0",
     ),
+    "zself_profile_events": attr.label(
+        doc = "Passes -Zself-profile and -Zself-profile-events flag to rustc, requires a nightly toolchain.",
+        providers = [UnstableSelfProfileInfo],
+        default = None,
+    ),
     "_collect_cfgs": attr.label(
         doc = "Enable collection of cfg flags with results stored in CrateInfo.cfgs.",
         default = Label("//rust/settings:collect_cfgs"),

rust/private/rustc.bzl

@@ -53,6 +53,7 @@ load(
     "is_exec_configuration",
     "is_std_dylib",
     "make_static_lib_symlink",
+    "matches_prefix_filter",
     "parse_env_strings",
     "relativize",
 )
@@ -105,6 +106,19 @@ PerCrateRustcFlagsInfo = provider(
     fields = {"per_crate_rustc_flags": "List[string] Extra flags to pass to rustc in non-exec configuration"},
 )
 
+UnstableSelfProfileInfo = provider(
+    doc = "Passes -Zself-profile and -Zself-profile-events flags to matching rust crates.",
+    fields = {
+        "events": (
+            "List[tuple[str, str]]: A list of `(pattern, event_types)` pairs. The `pattern` " +
+            "matches against a target's label (with leading `@//` stripped) or " +
+            "its execution path (an empty `pattern` matches all targets). The `event_types` " +
+            "specifies comma-separated categories of self-profile events to pass to " +
+            "`-Zself-profile-events` (e.g., `all`)."
+        ),
+    },
+)
+
 def _get_rustc_env(attr, toolchain, crate_name):
     """Gathers rustc environment variables
 
@@ -1455,6 +1469,45 @@ def collect_extra_rustc_flags(ctx, toolchain, crate_root, crate_type):
 
     return flags
 
+def setup_zself_profile(ctx, crate_info):
+    """Sets up rustc self-profiling if enabled by zself_profile_events.
+
+    Args:
+        ctx (ctx): The current rule's context object.
+        crate_info (CrateInfo): The CrateInfo provider of the target crate.
+
+    Returns:
+        tuple: A tuple containing:
+            - File: The declared self-profile directory, or None if disabled.
+            - list[str]: The self-profile flags to pass to rustc.
+    """
+    if not getattr(ctx.attr, "zself_profile_events", None) or UnstableSelfProfileInfo not in ctx.attr.zself_profile_events:
+        return None, []
+
+    events_info = ctx.attr.zself_profile_events[UnstableSelfProfileInfo].events
+
+    is_self_profile_enabled = False
+    event_types_to_use = None
+
+    # Check if the current crate matches any of the specified prefix filters.
+    # Matching works by comparing against the target's label or its execution path.
+    for pattern, event_types in events_info:
+        if matches_prefix_filter(ctx.label, crate_info.root.path, pattern):
+            is_self_profile_enabled = True
+            event_types_to_use = event_types
+            break
+
+    if not is_self_profile_enabled:
+        return None, []
+
+    profiling_dir = ctx.actions.declare_directory(crate_info.output.basename + "_self-profile", sibling = crate_info.output)
+
+    profiling_flags = ["-Zself-profile=%s" % profiling_dir.path]
+    if event_types_to_use:
+        profiling_flags.append("-Zself-profile-events=%s" % event_types_to_use)
+
+    return profiling_dir, profiling_flags
+
 def rustc_compile_action(
         *,
         ctx,
@@ -1542,6 +1595,9 @@ def rustc_compile_action(
         rust_flags = rust_flags + ctx.attr.lint_config[LintsInfo].rustc_lint_flags
         lint_files = lint_files + ctx.attr.lint_config[LintsInfo].rustc_lint_files
 
+    profiling_dir, profiling_flags = setup_zself_profile(ctx, crate_info)
+    rust_flags = rust_flags + profiling_flags
+
     compile_inputs, out_dir, build_env_files, build_flags_files, linkstamp_outs, ambiguous_libs = collect_inputs(
         ctx = ctx,
         file = ctx.file,
@@ -1676,6 +1732,8 @@ def rustc_compile_action(
     action_outputs = list(outputs)
     if rustc_output:
         action_outputs.append(rustc_output)
+    if profiling_dir:
+        action_outputs.append(profiling_dir)
 
     # Get the compilation mode for the current target.
     compilation_mode = get_compilation_mode_opts(ctx, toolchain)
@@ -1950,7 +2008,8 @@ def rustc_compile_action(
             output_group_info["rustc_rmeta_output"] = depset([rustc_rmeta_output])
     if rustc_output:
         output_group_info["rustc_output"] = depset([rustc_output])
-
+    if profiling_dir:
+        output_group_info["self_profile"] = depset([profiling_dir])
     if output_group_info:
         providers.append(OutputGroupInfo(**output_group_info))
 
@@ -2737,19 +2796,7 @@ def _collect_per_crate_rustc_flags(ctx, crate_root, per_crate_rustc_flags):
         if not flag:
             fail("per_crate_rustc_flag '{}' does not follow the expected format: prefix_filter@flag".format(per_crate_rustc_flag))
 
-        label_string = str(ctx.label)
-        if label_string.startswith("@//"):
-            label = label_string[1:]
-        elif label_string.startswith(
-            # buildifier: disable=canonical-repository
-            "@@//",
-        ):
-            label = label_string[2:]
-        else:
-            label = label_string
-        execution_path = crate_root.path
-
-        if label.startswith(prefix_filter) or execution_path.startswith(prefix_filter):
+        if matches_prefix_filter(ctx.label, crate_root.path, prefix_filter):
             flags.append(flag)
 
     return flags
@@ -2946,3 +2993,38 @@ no_std = rule(
     },
     implementation = _no_std_impl,
 )
+
+def _zself_profile_events_impl(ctx):
+    events = []
+    for val in ctx.build_setting_value:
+        if not val:
+            continue
+        if "@" in val:
+            pattern, event_types = val.split("@", 1)
+            events.append((pattern, event_types))
+        else:
+            fail("zself_profile_events '{}' does not follow the expected format: prefix_filter@comma_separated_flag".format(val))
+    return [UnstableSelfProfileInfo(events = events)]
+
+zself_profile_events = rule(
+    doc = (
+        "Passes -Zself-profile and -Zself-profile-events flags to matching Rust crates." +
+        "This feature allows end-users to profile rustc compiler performance on specific crates " +
+        "using rustc's self-profiler. Because these flags are unstable, using them requires a " +
+        "nightly compiler toolchain. The setting is configured from the command line via " +
+        "`--@rules_rust//rust/settings:zself_profile_events`." +
+        "The expected value format is `<prefix_filter>@<events_specification>`. Multiple uses of " +
+        "this flag are accumulated, however only first <events_specification> will be applied for same" +
+        "<prefix_filter>." +
+        "If the target prefix matches with <prefix_filter>, `-Zself-profile` and `-Zself-profile-events` " +
+        "with values as `<crate_name>.self-profile/` and <events_specification> respectively " +
+        "is passed to rustc compiler. The generated profile files (e.g., `.mm_profdata`) are placed" +
+        "under `bazel-out/bin/path/to/package/crate_name_self-profile/` which can be seen by passing" +
+        " `--output_groups=self_profile` flag." +
+        "blaze build //my/project:my_lib \\" +
+        "--@rules_rust//rust/settings:zself_profile_events=//my/project@all \\" +
+        "--output_groups=self_profile"
+    ),
+    implementation = _zself_profile_events_impl,
+    build_setting = config.string_list(flag = True, repeatable = True),
+)

rust/private/utils.bzl

@@ -978,3 +978,32 @@ def is_std_dylib(file):
         # for windows
         basename.startswith("std-") and basename.endswith(".dll")
     )
+
+def matches_prefix_filter(label, crate_root_path, pattern):
+    """Determines if a target matches a prefix filter pattern.
+
+    Matching works by comparing against the target's label or its execution path.
+
+    Args:
+        label (Label): The target's label.
+        crate_root_path (str): The target's execution path (e.g. crate root path).
+        pattern (str): The prefix pattern to match against.
+
+    Returns:
+        bool: True if the target matches the pattern, False otherwise.
+    """
+    if not pattern:
+        return True
+
+    label_string = str(label)
+    if label_string.startswith("@//"):
+        target_label = label_string[1:]
+    elif label_string.startswith(
+        # buildifier: disable=canonical-repository
+        "@@//",
+    ):
+        target_label = label_string[2:]
+    else:
+        target_label = label_string
+
+    return target_label.startswith(pattern) or crate_root_path.startswith(pattern)

rust/rust_common.bzl

@@ -30,6 +30,7 @@ load(
     _TestCrateInfo = "TestCrateInfo",
     _UnstableRustFeaturesInfo = "UnstableRustFeaturesInfo",
 )
+load("//rust/private:rustc.bzl", _UnstableSelfProfileInfo = "UnstableSelfProfileInfo")
 
 BuildInfo = _BuildInfo
 ClippyInfo = _ClippyInfo
@@ -39,7 +40,7 @@ DepInfo = _DepInfo
 DepVariantInfo = _DepVariantInfo
 TestCrateInfo = _TestCrateInfo
 UnstableRustFeaturesInfo = _UnstableRustFeaturesInfo
-
+UnstableSelfProfileInfo = _UnstableSelfProfileInfo
 COMMON_PROVIDERS = _COMMON_PROVIDERS
 
 rust_common = _rust_common

rust/settings/BUILD.bazel

@@ -40,6 +40,7 @@ load(
     "toolchain_generated_sysroot",
     "toolchain_linker_preference",
     "unpretty",
+    "zself_profile_events",
 )
 
 package(default_visibility = ["//visibility:public"])
@@ -137,3 +138,5 @@ toolchain_generated_sysroot()
 toolchain_linker_preference()
 
 unpretty()
+
+zself_profile_events()

rust/settings/settings.bzl

@@ -30,6 +30,7 @@ load(
     _no_std = "no_std",
     _per_crate_rustc_flag = "per_crate_rustc_flag",
     _rustc_output_diagnostics = "rustc_output_diagnostics",
+    _zself_profile_events = "zself_profile_events",
 )
 load("//rust/private:unpretty.bzl", "UNPRETTY_MODES", "rust_unpretty_flag")
 load(":incompatible.bzl", "incompatible_flag")
@@ -575,3 +576,11 @@ def collect_cfgs():
         scope = "universal",
         build_setting_default = False,
     )
+
+def zself_profile_events(name = "zself_profile_events"):
+    """Passes -Zself-profile and -Zself-profile-events flags to rustc, requires a nightly toolchain.
+    """
+    _zself_profile_events(
+        name = name,
+        build_setting_default = [],
+    )

test/integration/unstable_rust_features/BUILD.bazel

@@ -1,4 +1,5 @@
-load("@rules_rust//rust:defs.bzl", "rust_binary")
+load("@rules_rust//rust:defs.bzl", "rust_binary", "rust_test")
+load(":self_profile_provider.bzl", "unstable_self_profiling_flags")
 load(":unstable_features_for_test.bzl", "unstable_rust_features_for_test_rule")
 
 unstable_rust_features_for_test_rule(name = "unstable_rust_features_for_test")
@@ -8,3 +9,26 @@ rust_binary(
     srcs = ["main.rs"],
     unstable_rust_features_config = ":unstable_rust_features_for_test",
 )
+
+unstable_self_profiling_flags(name = "unstable_self_profiling_flags")
+
+rust_binary(
+    name = "sample_binary",
+    srcs = ["hello.rs"],
+    zself_profile_events = ":unstable_self_profiling_flags",
+)
+
+filegroup(
+    name = "self_profile_data",
+    srcs = [":sample_binary"],
+    output_group = "self_profile",
+)
+
+rust_test(
+    name = "test",
+    srcs = ["validate.rs"],
+    data = [":self_profile_data"],
+    edition = "2021",
+    env = {"PROFILE_PATH": "$(rootpath :sample_binary)"},
+    deps = [":sample_binary"],
+)

test/integration/unstable_rust_features/hello.rs

@@ -0,0 +1,7 @@
+fn say_hello(){
+    println!("Hello, World!");
+}
+
+pub fn main(){
+    say_hello()
+}

test/integration/unstable_rust_features/self_profile_provider.bzl

@@ -0,0 +1,14 @@
+"""Defines a test rule providing UnstableSelfProfileInfo"""
+
+load("@rules_rust//rust:rust_common.bzl", "UnstableSelfProfileInfo")
+
+def _get_self_profiling_flag_impl(_ctx):
+    return UnstableSelfProfileInfo(
+        events = [("//:sample_binary", "all")],
+    )
+
+unstable_self_profiling_flags = rule(
+    attrs = {},
+    provides = [UnstableSelfProfileInfo],
+    implementation = _get_self_profiling_flag_impl,
+)

test/integration/unstable_rust_features/validate.rs

@@ -0,0 +1,15 @@
+#[test]
+fn test_genquery_file_empty() {
+    let profile_base_path = std::path::PathBuf::from(std::env::var_os("PROFILE_PATH").unwrap());
+    let profile_path = profile_base_path.parent().unwrap().join("sample_binary_self-profile");
+    let profile_path = std::path::Path::new(&profile_path);
+    assert!(profile_path.is_dir());
+
+    let files_count = match std::fs::read_dir(profile_path) {
+        Ok(entries) => {
+            entries.filter_map(Result::ok).count()
+        }
+        Err(_) => 0,
+    };
+    assert!(files_count > 0);
+}