feat: Add //rust/settings:lto (#3104)

Fixes #3045

This PR adds a new build setting `//rust/settings:lto=(off|thin|fat)`
which changes how we specify the following flags:

*
[`lto`](https://doc.rust-lang.org/rustc/codegen-options/index.html#lto)
*
[`embed-bitcode`](https://doc.rust-lang.org/rustc/codegen-options/index.html#embed-bitcode)
*
[`linker-plugin-lto`](https://doc.rust-lang.org/rustc/codegen-options/index.html#linker-plugin-lto)

The way we invoke the flags was based on how Cargo does it today
([code](https://github.com/rust-lang/cargo/blob/769f622e12db0001431d8ae36d1093fb8727c5d9/src/cargo/core/compiler/lto.rs#L4))
and based on suggestions from the [Rust
docs](https://doc.rust-lang.org/rustc/codegen-options/index.html#embed-bitcode).

When LTO is not enabled, we will specify `-Cembed-bitcode=no` which
tells `rustc` to skip embedding LLVM bitcode and should speed up builds.
Similarly when LTO is enabled we specify `-Clinker-plugin-lto` which
will cause `rustc` to skip generating objects files entirely, and
instead replace them with LLVM bitcode*.

*only when building an `rlib`, when building other crate types we
continue generating object files.

I added unit tests to make sure we pass the flags correctly, as well as
some docs describing the new setting. Please let me know if I should add
more!

rust/private/lto.bzl

@@ -0,0 +1,114 @@
+"""A module defining Rust link time optimization (lto) rules"""
+
+load("//rust/private:utils.bzl", "is_exec_configuration")
+
+_LTO_MODES = [
+    # Default. No mode has been explicitly set, rustc will do "thin local" LTO
+    # between the codegen units of a single crate.
+    "unspecified",
+    # LTO has been explicitly turned "off".
+    "off",
+    # Perform "thin" LTO. This is similar to "fat" but takes significantly less
+    # time to run, but provides similar performance improvements.
+    #
+    # See: <http://blog.llvm.org/2016/06/thinlto-scalable-and-incremental-lto.html>
+    "thin",
+    # Perform "fat"/full LTO.
+    "fat",
+]
+
+RustLtoInfo = provider(
+    doc = "A provider describing the link time optimization setting.",
+    fields = {"mode": "string: The LTO mode specified via a build setting."},
+)
+
+def _rust_lto_flag_impl(ctx):
+    value = ctx.build_setting_value
+
+    if value not in _LTO_MODES:
+        msg = "{NAME} build setting allowed to take values [{EXPECTED}], but was set to: {ACTUAL}".format(
+            NAME = ctx.label,
+            VALUES = ", ".join(["'{}'".format(m) for m in _LTO_MODES]),
+            ACTUAL = value,
+        )
+        fail(msg)
+
+    return RustLtoInfo(mode = value)
+
+rust_lto_flag = rule(
+    doc = "A build setting which specifies the link time optimization mode used when building Rust code. Allowed values are: ".format(_LTO_MODES),
+    implementation = _rust_lto_flag_impl,
+    build_setting = config.string(flag = True),
+)
+
+def _determine_lto_object_format(ctx, toolchain, crate_info):
+    """Determines if we should run LTO and what bitcode should get included in a built artifact.
+
+    Args:
+        ctx (ctx): The calling rule's context object.
+        toolchain (rust_toolchain): The current target's `rust_toolchain`.
+        crate_info (CrateInfo): The CrateInfo provider of the target crate.
+
+    Returns:
+        string: Returns one of only_object, only_bitcode, object_and_bitcode.
+    """
+
+    # Even if LTO is enabled don't use it for actions being built in the exec
+    # configuration, e.g. build scripts and proc-macros. This mimics Cargo.
+    if is_exec_configuration(ctx):
+        return "only_object"
+
+    mode = toolchain._lto.mode
+
+    if mode in ["off", "unspecified"]:
+        return "only_object"
+
+    perform_linking = crate_info.type in ["bin", "staticlib", "cdylib"]
+
+    # is_linkable = crate_info.type in ["lib", "rlib", "dylib", "proc-macro"]
+    is_dynamic = crate_info.type in ["dylib", "cdylib", "proc-macro"]
+    needs_object = perform_linking or is_dynamic
+
+    # At this point we know LTO is enabled, otherwise we would have returned above.
+
+    if not needs_object:
+        # If we're building an 'rlib' and LTO is enabled, then we can skip
+        # generating object files entirely.
+        return "only_bitcode"
+    elif crate_info.type == "dylib":
+        # If we're a dylib and we're running LTO, then only emit object code
+        # because 'rustc' doesn't currently support LTO with dylibs.
+        return "only_object"
+    else:
+        return "object_and_bitcode"
+
+def construct_lto_arguments(ctx, toolchain, crate_info):
+    """Returns a list of 'rustc' flags to configure link time optimization.
+
+    Args:
+        ctx (ctx): The calling rule's context object.
+        toolchain (rust_toolchain): The current target's `rust_toolchain`.
+        crate_info (CrateInfo): The CrateInfo provider of the target crate.
+
+    Returns:
+        list: A list of strings that are valid flags for 'rustc'.
+    """
+    mode = toolchain._lto.mode
+    format = _determine_lto_object_format(ctx, toolchain, crate_info)
+
+    args = []
+
+    if mode in ["thin", "fat", "off"] and not is_exec_configuration(ctx):
+        args.append("lto={}".format(mode))
+
+    if format in ["unspecified", "object_and_bitcode"]:
+        # Embedding LLVM bitcode in object files is `rustc's` default.
+        args.extend([])
+    elif format in ["off", "only_object"]:
+        args.extend(["embed-bitcode=no"])
+    elif format == "only_bitcode":
+        args.extend(["linker-plugin-lto"])
+    else:
+        fail("unrecognized LTO object format {}".format(format))
+
+    return ["-C{}".format(arg) for arg in args]

rust/private/rustc.bzl

@@ -24,6 +24,7 @@ load(
 )
 load("//rust/private:common.bzl", "rust_common")
 load("//rust/private:compat.bzl", "abs")
+load("//rust/private:lto.bzl", "construct_lto_arguments")
 load("//rust/private:providers.bzl", "RustcOutputDiagnosticsInfo", _BuildInfo = "BuildInfo")
 load("//rust/private:stamp.bzl", "is_stamping_enabled")
 load(
@@ -998,6 +999,7 @@ def construct_arguments(
     data_paths = depset(direct = getattr(attr, "data", []), transitive = [crate_info.compile_data_targets]).to_list()
 
     add_edition_flags(rustc_flags, crate_info)
+    _add_lto_flags(ctx, toolchain, rustc_flags, crate_info)
 
     # Link!
     if ("link" in emit and crate_info.type not in ["rlib", "lib"]) or add_flags_for_binary:
@@ -1583,6 +1585,18 @@ def _collect_nonstatic_linker_inputs(cc_info):
             ))
     return shared_linker_inputs
 
+def _add_lto_flags(ctx, toolchain, args, crate):
+    """Adds flags to an Args object to configure LTO for 'rustc'.
+
+    Args:
+        ctx (ctx): The calling rule's context object.
+        toolchain (rust_toolchain): The current target's `rust_toolchain`.
+        args (Args): A reference to an Args object
+        crate (CrateInfo): A CrateInfo provider
+    """
+    lto_args = construct_lto_arguments(ctx, toolchain, crate)
+    args.add_all(lto_args)
+
 def establish_cc_info(ctx, attr, crate_info, toolchain, cc_toolchain, feature_configuration, interface_library):
     """If the produced crate is suitable yield a CcInfo to allow for interop with cc rules
 

rust/settings/BUILD.bazel

@@ -14,6 +14,7 @@ load(
     "per_crate_rustc_flag",
     "rustc_output_diagnostics",
 )
+load("//rust/private:lto.bzl", "rust_lto_flag")
 load("//rust/private:unpretty.bzl", "rust_unpretty_flag")
 load(":incompatible.bzl", "incompatible_flag")
 
@@ -48,6 +49,12 @@ rust_unpretty_flag(
     visibility = ["//visibility:public"],
 )
 
+rust_lto_flag(
+    name = "lto",
+    build_setting_default = "unspecified",
+    visibility = ["//visibility:public"],
+)
+
 # A flag controlling whether to rename first-party crates such that their names
 # encode the Bazel package and target name, instead of just the target name.
 #

rust/toolchain.bzl

@@ -3,6 +3,7 @@
 load("@bazel_skylib//rules:common_settings.bzl", "BuildSettingInfo")
 load("//rust/platform:triple.bzl", "triple")
 load("//rust/private:common.bzl", "rust_common")
+load("//rust/private:lto.bzl", "RustLtoInfo")
 load("//rust/private:rust_analyzer.bzl", _rust_analyzer_toolchain = "rust_analyzer_toolchain")
 load(
     "//rust/private:rustfmt.bzl",
@@ -517,6 +518,7 @@ def _rust_toolchain_impl(ctx):
     third_party_dir = ctx.attr._third_party_dir[BuildSettingInfo].value
     pipelined_compilation = ctx.attr._pipelined_compilation[BuildSettingInfo].value
     no_std = ctx.attr._no_std[BuildSettingInfo].value
+    lto = ctx.attr._lto[RustLtoInfo]
 
     experimental_use_global_allocator = ctx.attr._experimental_use_global_allocator[BuildSettingInfo].value
     if _experimental_use_cc_common_link(ctx):
@@ -701,6 +703,7 @@ def _rust_toolchain_impl(ctx):
         _toolchain_generated_sysroot = ctx.attr._toolchain_generated_sysroot[BuildSettingInfo].value,
         _incompatible_do_not_include_data_in_compile_data = ctx.attr._incompatible_do_not_include_data_in_compile_data[IncompatibleFlagInfo].enabled,
         _no_std = no_std,
+        _lto = lto,
     )
     return [
         toolchain,
@@ -891,6 +894,10 @@ rust_toolchain = rule(
             default = Label("//rust/settings:incompatible_do_not_include_data_in_compile_data"),
             doc = "Label to a boolean build setting that controls whether to include data files in compile_data.",
         ),
+        "_lto": attr.label(
+            providers = [RustLtoInfo],
+            default = Label("//rust/settings:lto"),
+        ),
         "_no_std": attr.label(
             default = Label("//rust/settings:no_std"),
         ),

test/unit/lto/BUILD.bazel

@@ -0,0 +1,5 @@
+load(":lto_test_suite.bzl", "lto_test_suite")
+
+lto_test_suite(
+    name = "lto_test_suite",
+)

test/unit/lto/lto_test_suite.bzl

@@ -0,0 +1,123 @@
+"""Starlark tests for `//rust/settings/lto`"""
+
+load("@bazel_skylib//lib:unittest.bzl", "analysistest")
+load("@bazel_skylib//rules:write_file.bzl", "write_file")
+load("//rust:defs.bzl", "rust_library")
+load(
+    "//test/unit:common.bzl",
+    "assert_action_mnemonic",
+    "assert_argv_contains",
+    "assert_argv_contains_not",
+    "assert_argv_contains_prefix_not",
+)
+
+def _lto_test_impl(ctx, lto_setting, embed_bitcode, linker_plugin):
+    env = analysistest.begin(ctx)
+    target = analysistest.target_under_test(env)
+
+    action = target.actions[0]
+    assert_action_mnemonic(env, action, "Rustc")
+
+    # Check if LTO is enabled.
+    if lto_setting:
+        assert_argv_contains(env, action, "-Clto={}".format(lto_setting))
+    else:
+        assert_argv_contains_prefix_not(env, action, "-Clto")
+
+    # Check if we should embed bitcode.
+    if embed_bitcode:
+        assert_argv_contains(env, action, "-Cembed-bitcode={}".format(embed_bitcode))
+    else:
+        assert_argv_contains_prefix_not(env, action, "-Cembed-bitcode")
+
+    # Check if we should use linker plugin LTO.
+    if linker_plugin:
+        assert_argv_contains(env, action, "-Clinker-plugin-lto")
+    else:
+        assert_argv_contains_not(env, action, "-Clinker-plugin-lto")
+
+    return analysistest.end(env)
+
+def _lto_level_default(ctx):
+    return _lto_test_impl(ctx, None, "no", False)
+
+_lto_level_default_test = analysistest.make(
+    _lto_level_default,
+    config_settings = {},
+)
+
+def _lto_level_off(ctx):
+    return _lto_test_impl(ctx, "off", "no", False)
+
+_lto_level_off_test = analysistest.make(
+    _lto_level_off,
+    config_settings = {str(Label("//rust/settings:lto")): "off"},
+)
+
+def _lto_level_thin(ctx):
+    return _lto_test_impl(ctx, "thin", None, True)
+
+_lto_level_thin_test = analysistest.make(
+    _lto_level_thin,
+    config_settings = {str(Label("//rust/settings:lto")): "thin"},
+)
+
+def _lto_level_fat(ctx):
+    return _lto_test_impl(ctx, "fat", None, True)
+
+_lto_level_fat_test = analysistest.make(
+    _lto_level_fat,
+    config_settings = {str(Label("//rust/settings:lto")): "fat"},
+)
+
+def lto_test_suite(name):
+    """Entry-point macro called from the BUILD file.
+
+    Args:
+        name (str): The name of the test suite.
+    """
+    write_file(
+        name = "crate_lib",
+        out = "lib.rs",
+        content = [
+            "#[allow(dead_code)]",
+            "fn add() {}",
+            "",
+        ],
+    )
+
+    rust_library(
+        name = "lib",
+        srcs = [":lib.rs"],
+        edition = "2021",
+    )
+
+    _lto_level_default_test(
+        name = "lto_level_default_test",
+        target_under_test = ":lib",
+    )
+
+    _lto_level_off_test(
+        name = "lto_level_off_test",
+        target_under_test = ":lib",
+    )
+
+    _lto_level_thin_test(
+        name = "lto_level_thin_test",
+        target_under_test = ":lib",
+    )
+
+    _lto_level_fat_test(
+        name = "lto_level_fat_test",
+        target_under_test = ":lib",
+    )
+
+    native.test_suite(
+        name = name,
+        tests = [
+            ":lto_level_default_test",
+            ":lto_level_off_test",
+            ":lto_level_thin_test",
+            ":lto_level_fat_test",
+        ],
+    )