From ab58647b2e98e13136a1fe80308b6da1c933437a Mon Sep 17 00:00:00 2001
From: Ivan Lozano <ivanlozano@google.com>
Date: Wed, 15 May 2024 10:59:47 -0400
Subject: [PATCH] rust: Add an option to disable LTO for Rust

This adds an option to disable LTO when building a Rust module. This is
mostly intended to speedu p local prototyping, and LTO should not
normally be disabled for production builds.

Bug: 339628497
Test: m blueprint_tests && m rust
Change-Id: I21d5d4513a259a56f101ce8906e2bef7404e4efb
---
 rust/builder.go       | 18 ++++++++++++++----
 rust/compiler.go      | 15 +++++++++++++++
 rust/compiler_test.go | 29 +++++++++++++++++++++++++++++
 3 files changed, 58 insertions(+), 4 deletions(-)

diff --git a/rust/builder.go b/rust/builder.go
index abbf90d83..1ce92f4ec 100644
--- a/rust/builder.go
+++ b/rust/builder.go
@@ -158,7 +158,9 @@ func getTransformProperties(ctx ModuleContext, crateType string) transformProper
 
 func TransformSrcToBinary(ctx ModuleContext, mainSrc android.Path, deps PathDeps, flags Flags,
 	outputFile android.WritablePath) buildOutput {
-	flags.GlobalRustFlags = append(flags.GlobalRustFlags, "-C lto=thin")
+	if ctx.RustModule().compiler.Thinlto() {
+		flags.GlobalRustFlags = append(flags.GlobalRustFlags, "-C lto=thin")
+	}
 
 	return transformSrctoCrate(ctx, mainSrc, deps, flags, outputFile, getTransformProperties(ctx, "bin"))
 }
@@ -212,20 +214,28 @@ func TransformRlibstoStaticlib(ctx android.ModuleContext, mainSrc android.Path,
 
 func TransformSrctoDylib(ctx ModuleContext, mainSrc android.Path, deps PathDeps, flags Flags,
 	outputFile android.WritablePath) buildOutput {
-	flags.GlobalRustFlags = append(flags.GlobalRustFlags, "-C lto=thin")
+	if ctx.RustModule().compiler.Thinlto() {
+		flags.GlobalRustFlags = append(flags.GlobalRustFlags, "-C lto=thin")
+	}
 
 	return transformSrctoCrate(ctx, mainSrc, deps, flags, outputFile, getTransformProperties(ctx, "dylib"))
 }
 
 func TransformSrctoStatic(ctx ModuleContext, mainSrc android.Path, deps PathDeps, flags Flags,
 	outputFile android.WritablePath) buildOutput {
-	flags.GlobalRustFlags = append(flags.GlobalRustFlags, "-C lto=thin")
+	if ctx.RustModule().compiler.Thinlto() {
+		flags.GlobalRustFlags = append(flags.GlobalRustFlags, "-C lto=thin")
+	}
+
 	return transformSrctoCrate(ctx, mainSrc, deps, flags, outputFile, getTransformProperties(ctx, "staticlib"))
 }
 
 func TransformSrctoShared(ctx ModuleContext, mainSrc android.Path, deps PathDeps, flags Flags,
 	outputFile android.WritablePath) buildOutput {
-	flags.GlobalRustFlags = append(flags.GlobalRustFlags, "-C lto=thin")
+	if ctx.RustModule().compiler.Thinlto() {
+		flags.GlobalRustFlags = append(flags.GlobalRustFlags, "-C lto=thin")
+	}
+
 	return transformSrctoCrate(ctx, mainSrc, deps, flags, outputFile, getTransformProperties(ctx, "cdylib"))
 }
 
diff --git a/rust/compiler.go b/rust/compiler.go
index 5033fbac2..efc3deef3 100644
--- a/rust/compiler.go
+++ b/rust/compiler.go
@@ -47,6 +47,7 @@ type compiler interface {
 	edition() string
 	features() []string
 	rustdoc(ctx ModuleContext, flags Flags, deps PathDeps) android.OptionalPath
+	Thinlto() bool
 
 	// Output directory in which source-generated code from dependencies is
 	// copied. This is equivalent to Cargo's OUT_DIR variable.
@@ -231,6 +232,15 @@ type BaseCompilerProperties struct {
 
 	// If cargo_env_compat is true, sets the CARGO_PKG_VERSION env var to this value.
 	Cargo_pkg_version *string
+
+	// Control whether LTO is used for the final (Rust) linkage. This does not impact
+	// cross-language LTO.
+	Lto struct {
+		// Whether thin LTO should be enabled. By default this is true.
+		// LTO provides such a large code size benefit for Rust, this should always
+		// be enabled for production builds unless there's a clear need to disable it.
+		Thin *bool `android:"arch_variant"`
+	} `android:"arch_variant"`
 }
 
 type baseCompiler struct {
@@ -273,6 +283,11 @@ func (compiler *baseCompiler) Disabled() bool {
 	return false
 }
 
+// Thin LTO is enabled by default.
+func (compiler *baseCompiler) Thinlto() bool {
+	return BoolDefault(compiler.Properties.Lto.Thin, true)
+}
+
 func (compiler *baseCompiler) SetDisabled() {
 	panic("baseCompiler does not implement SetDisabled()")
 }
diff --git a/rust/compiler_test.go b/rust/compiler_test.go
index 89f4d1abf..4caa12b3e 100644
--- a/rust/compiler_test.go
+++ b/rust/compiler_test.go
@@ -63,6 +63,35 @@ func TestCfgsToFlags(t *testing.T) {
 	}
 }
 
+func TestLtoFlag(t *testing.T) {
+	ctx := testRust(t, `
+		rust_library_host {
+			name: "libfoo",
+			srcs: ["foo.rs"],
+			crate_name: "foo",
+			lto: {
+				thin: false,
+			}
+		}
+
+		rust_library_host {
+			name: "libfoo_lto",
+			srcs: ["foo.rs"],
+			crate_name: "foo",
+		}
+		`)
+
+	libfoo := ctx.ModuleForTests("libfoo", "linux_glibc_x86_64_dylib").Rule("rustc")
+	libfooLto := ctx.ModuleForTests("libfoo_lto", "linux_glibc_x86_64_dylib").Rule("rustc")
+
+	if strings.Contains(libfoo.Args["rustcFlags"], "-C lto=thin") {
+		t.Fatalf("libfoo expected to disable lto -- rustcFlags: %#v", libfoo.Args["rustcFlags"])
+	}
+	if !strings.Contains(libfooLto.Args["rustcFlags"], "-C lto=thin") {
+		t.Fatalf("libfoo expected to enable lto by default -- rustcFlags: %#v", libfooLto.Args["rustcFlags"])
+	}
+}
+
 // Test that we reject multiple source files.
 func TestEnforceSingleSourceFile(t *testing.T) {