AFDO for Rust

Bug: 195134194
Bug: 165018530
Test: build
Change-Id: I30932a22dc0b22716cdc925a3fcc5f9a169fcec4
diff --git a/cc/afdo.go b/cc/afdo.go
index 022f283..f7639fa 100644
--- a/cc/afdo.go
+++ b/cc/afdo.go
@@ -40,7 +40,7 @@
 	})
 }
 
-func recordMissingAfdoProfileFile(ctx BaseModuleContext, missing string) {
+func recordMissingAfdoProfileFile(ctx android.BaseModuleContext, missing string) {
 	getNamedMapForConfig(ctx.Config(), modulesMissingProfileFileKey).Store(missing, true)
 }
 
@@ -67,14 +67,14 @@
 //   1. libfoo_arm64.afdo
 //   2. libfoo.afdo
 // Add more specialisation as needed.
-func getProfileFiles(ctx BaseModuleContext, moduleName string) []string {
+func getProfileFiles(ctx android.BaseModuleContext, moduleName string) []string {
 	var files []string
 	files = append(files, moduleName+"_"+ctx.Arch().ArchType.String()+".afdo")
 	files = append(files, moduleName+".afdo")
 	return files
 }
 
-func (props *AfdoProperties) getAfdoProfileFile(ctx BaseModuleContext, module string) android.OptionalPath {
+func (props *AfdoProperties) GetAfdoProfileFile(ctx android.BaseModuleContext, module string) android.OptionalPath {
 	// Test if the profile_file is present in any of the Afdo profile projects
 	for _, profileFile := range getProfileFiles(ctx, module) {
 		for _, profileProject := range getAfdoProfileProjects(ctx.DeviceConfig()) {
@@ -95,7 +95,7 @@
 func (afdo *afdo) begin(ctx BaseModuleContext) {
 	if afdo.Properties.Afdo && !ctx.static() && !ctx.Host() {
 		module := ctx.ModuleName()
-		if afdo.Properties.getAfdoProfileFile(ctx, module).Valid() {
+		if afdo.Properties.GetAfdoProfileFile(ctx, module).Valid() {
 			afdo.Properties.AfdoTarget = proptools.StringPtr(module)
 		}
 	}
@@ -103,7 +103,7 @@
 
 func (afdo *afdo) flags(ctx ModuleContext, flags Flags) Flags {
 	if profile := afdo.Properties.AfdoTarget; profile != nil {
-		if profileFile := afdo.Properties.getAfdoProfileFile(ctx, *profile); profileFile.Valid() {
+		if profileFile := afdo.Properties.GetAfdoProfileFile(ctx, *profile); profileFile.Valid() {
 			profileFilePath := profileFile.Path()
 
 			profileUseFlag := fmt.Sprintf(afdoCFlagsFormat, profileFile)
diff --git a/rust/Android.bp b/rust/Android.bp
index 5e14da8..3fd68e5 100644
--- a/rust/Android.bp
+++ b/rust/Android.bp
@@ -14,6 +14,7 @@
         "soong-snapshot",
     ],
     srcs: [
+        "afdo.go",
         "androidmk.go",
         "benchmark.go",
         "binary.go",
diff --git a/rust/afdo.go b/rust/afdo.go
new file mode 100644
index 0000000..996fd7e
--- /dev/null
+++ b/rust/afdo.go
@@ -0,0 +1,48 @@
+// Copyright 2022 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package rust
+
+import (
+	"fmt"
+
+	"android/soong/cc"
+)
+
+const afdoFlagFormat = "-Zprofile-sample-use=%s"
+
+type afdo struct {
+	Properties cc.AfdoProperties
+}
+
+func (afdo *afdo) props() []interface{} {
+	return []interface{}{&afdo.Properties}
+}
+
+func (afdo *afdo) flags(ctx ModuleContext, flags Flags, deps PathDeps) (Flags, PathDeps) {
+	if ctx.Host() {
+		return flags, deps
+	}
+
+	if afdo != nil && afdo.Properties.Afdo {
+		if profileFile := afdo.Properties.GetAfdoProfileFile(ctx, ctx.ModuleName()); profileFile.Valid() {
+			profileUseFlag := fmt.Sprintf(afdoFlagFormat, profileFile)
+			flags.RustFlags = append(flags.RustFlags, profileUseFlag)
+
+			profileFilePath := profileFile.Path()
+			deps.AfdoProfiles = append(deps.AfdoProfiles, profileFilePath)
+		}
+	}
+	return flags, deps
+}
diff --git a/rust/builder.go b/rust/builder.go
index a7efc28..e66a6f0 100644
--- a/rust/builder.go
+++ b/rust/builder.go
@@ -245,6 +245,7 @@
 	implicits = append(implicits, deps.StaticLibs...)
 	implicits = append(implicits, deps.SharedLibDeps...)
 	implicits = append(implicits, deps.srcProviderFiles...)
+	implicits = append(implicits, deps.AfdoProfiles...)
 
 	if deps.CrtBegin.Valid() {
 		implicits = append(implicits, deps.CrtBegin.Path(), deps.CrtEnd.Path())
diff --git a/rust/rust.go b/rust/rust.go
index cba92c3..0f7b768 100644
--- a/rust/rust.go
+++ b/rust/rust.go
@@ -148,6 +148,7 @@
 
 	makeLinkType string
 
+	afdo             *afdo
 	compiler         compiler
 	coverage         *coverage
 	clippy           *clippy
@@ -403,6 +404,7 @@
 	SharedLibDeps android.Paths
 	StaticLibs    android.Paths
 	ProcMacros    RustLibraries
+	AfdoProfiles  android.Paths
 
 	// depFlags and depLinkFlags are rustc and linker (clang) flags.
 	depFlags     []string
@@ -551,6 +553,7 @@
 	module.AddProperties(props...)
 	module.AddProperties(
 		&BaseProperties{},
+		&cc.AfdoProperties{},
 		&cc.VendorProperties{},
 		&BenchmarkProperties{},
 		&BindgenProperties{},
@@ -688,6 +691,9 @@
 	mod.AddProperties(&mod.Properties)
 	mod.AddProperties(&mod.VendorProperties)
 
+	if mod.afdo != nil {
+		mod.AddProperties(mod.afdo.props()...)
+	}
 	if mod.compiler != nil {
 		mod.AddProperties(mod.compiler.compilerProps()...)
 	}
@@ -719,6 +725,7 @@
 }
 func newModule(hod android.HostOrDeviceSupported, multilib android.Multilib) *Module {
 	module := newBaseModule(hod, multilib)
+	module.afdo = &afdo{}
 	module.coverage = &coverage{}
 	module.clippy = &clippy{}
 	module.sanitize = &sanitize{}
@@ -856,6 +863,9 @@
 	}
 
 	// Calculate rustc flags
+	if mod.afdo != nil {
+		flags, deps = mod.afdo.flags(ctx, flags, deps)
+	}
 	if mod.compiler != nil {
 		flags = mod.compiler.compilerFlags(ctx, flags)
 		flags = mod.compiler.cfgFlags(ctx, flags)