Introduce afdo

A new configuration rule for sampling PGO. This differs from the
original pgo.go rule in the following ways:
 * Automatic propagation to static dependencies
 * Simpler configuration (just put `afdo: true` to optimsation targets)

http://go/android-afdo

Test: build
Bug: 79161490
Change-Id: Ie194824cd523bca19e10ced41d2078fc598f13b3
diff --git a/android/config.go b/android/config.go
index 5c0e5ae..0187a8a 100644
--- a/android/config.go
+++ b/android/config.go
@@ -1245,6 +1245,10 @@
 	return coverage
 }
 
+func (c *deviceConfig) AfdoAdditionalProfileDirs() []string {
+	return c.config.productVariables.AfdoAdditionalProfileDirs
+}
+
 func (c *deviceConfig) PgoAdditionalProfileDirs() []string {
 	return c.config.productVariables.PgoAdditionalProfileDirs
 }
diff --git a/android/variable.go b/android/variable.go
index bc93835..7bccc91 100644
--- a/android/variable.go
+++ b/android/variable.go
@@ -327,7 +327,8 @@
 
 	NamespacesToExport []string `json:",omitempty"`
 
-	PgoAdditionalProfileDirs []string `json:",omitempty"`
+	AfdoAdditionalProfileDirs []string `json:",omitempty"`
+	PgoAdditionalProfileDirs  []string `json:",omitempty"`
 
 	VndkUseCoreVariant         *bool `json:",omitempty"`
 	VndkSnapshotBuildArtifacts *bool `json:",omitempty"`
diff --git a/cc/Android.bp b/cc/Android.bp
index 07aa7cb..0bf0045 100644
--- a/cc/Android.bp
+++ b/cc/Android.bp
@@ -19,10 +19,11 @@
         "soong-tradefed",
     ],
     srcs: [
+        "afdo.go",
         "androidmk.go",
         "api_level.go",
-        "builder.go",
         "bp2build.go",
+        "builder.go",
         "cc.go",
         "ccdeps.go",
         "check.go",
diff --git a/cc/afdo.go b/cc/afdo.go
new file mode 100644
index 0000000..022f283
--- /dev/null
+++ b/cc/afdo.go
@@ -0,0 +1,194 @@
+// Copyright 2021 Google Inc. All rights reserved.
+//
+// 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 cc
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/google/blueprint/proptools"
+
+	"android/soong/android"
+)
+
+var (
+	globalAfdoProfileProjects = []string{
+		"vendor/google_data/pgo_profile/sampling/",
+		"toolchain/pgo-profiles/sampling/",
+	}
+)
+
+var afdoProfileProjectsConfigKey = android.NewOnceKey("AfdoProfileProjects")
+
+const afdoCFlagsFormat = "-fprofile-sample-accurate -fprofile-sample-use=%s"
+
+func getAfdoProfileProjects(config android.DeviceConfig) []string {
+	return config.OnceStringSlice(afdoProfileProjectsConfigKey, func() []string {
+		return append(globalAfdoProfileProjects, config.AfdoAdditionalProfileDirs()...)
+	})
+}
+
+func recordMissingAfdoProfileFile(ctx BaseModuleContext, missing string) {
+	getNamedMapForConfig(ctx.Config(), modulesMissingProfileFileKey).Store(missing, true)
+}
+
+type AfdoProperties struct {
+	Afdo bool
+
+	AfdoTarget *string  `blueprint:"mutated"`
+	AfdoDeps   []string `blueprint:"mutated"`
+}
+
+type afdo struct {
+	Properties AfdoProperties
+}
+
+func (afdo *afdo) props() []interface{} {
+	return []interface{}{&afdo.Properties}
+}
+
+func (afdo *afdo) AfdoEnabled() bool {
+	return afdo != nil && afdo.Properties.Afdo && afdo.Properties.AfdoTarget != nil
+}
+
+// Get list of profile file names, ordered by level of specialisation. For example:
+//   1. libfoo_arm64.afdo
+//   2. libfoo.afdo
+// Add more specialisation as needed.
+func getProfileFiles(ctx 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 {
+	// 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()) {
+			path := android.ExistentPathForSource(ctx, profileProject, profileFile)
+			if path.Valid() {
+				return path
+			}
+		}
+	}
+
+	// Record that this module's profile file is absent
+	missing := ctx.ModuleDir() + ":" + module
+	recordMissingAfdoProfileFile(ctx, missing)
+
+	return android.OptionalPathForPath(nil)
+}
+
+func (afdo *afdo) begin(ctx BaseModuleContext) {
+	if afdo.Properties.Afdo && !ctx.static() && !ctx.Host() {
+		module := ctx.ModuleName()
+		if afdo.Properties.getAfdoProfileFile(ctx, module).Valid() {
+			afdo.Properties.AfdoTarget = proptools.StringPtr(module)
+		}
+	}
+}
+
+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() {
+			profileFilePath := profileFile.Path()
+
+			profileUseFlag := fmt.Sprintf(afdoCFlagsFormat, profileFile)
+			flags.Local.CFlags = append(flags.Local.CFlags, profileUseFlag)
+			flags.Local.LdFlags = append(flags.Local.LdFlags, profileUseFlag)
+			flags.Local.LdFlags = append(flags.Local.LdFlags, "-Wl,-mllvm,-no-warn-sample-unused=true")
+
+			// Update CFlagsDeps and LdFlagsDeps so the module is rebuilt
+			// if profileFile gets updated
+			flags.CFlagsDeps = append(flags.CFlagsDeps, profileFilePath)
+			flags.LdFlagsDeps = append(flags.LdFlagsDeps, profileFilePath)
+		}
+	}
+
+	return flags
+}
+
+// Propagate afdo requirements down from binaries
+func afdoDepsMutator(mctx android.TopDownMutatorContext) {
+	if m, ok := mctx.Module().(*Module); ok && m.afdo.AfdoEnabled() {
+		afdoTarget := *m.afdo.Properties.AfdoTarget
+		mctx.WalkDeps(func(dep android.Module, parent android.Module) bool {
+			tag := mctx.OtherModuleDependencyTag(dep)
+			libTag, isLibTag := tag.(libraryDependencyTag)
+
+			// Do not recurse down non-static dependencies
+			if isLibTag {
+				if !libTag.static() {
+					return false
+				}
+			} else {
+				if tag != objDepTag && tag != reuseObjTag {
+					return false
+				}
+			}
+
+			if dep, ok := dep.(*Module); ok {
+				dep.afdo.Properties.AfdoDeps = append(dep.afdo.Properties.AfdoDeps, afdoTarget)
+			}
+
+			return true
+		})
+	}
+}
+
+// Create afdo variants for modules that need them
+func afdoMutator(mctx android.BottomUpMutatorContext) {
+	if m, ok := mctx.Module().(*Module); ok && m.afdo != nil {
+		if m.afdo.AfdoEnabled() && !m.static() {
+			afdoTarget := *m.afdo.Properties.AfdoTarget
+			mctx.SetDependencyVariation(encodeTarget(afdoTarget))
+		}
+
+		variationNames := []string{""}
+		afdoDeps := android.FirstUniqueStrings(m.afdo.Properties.AfdoDeps)
+		for _, dep := range afdoDeps {
+			variationNames = append(variationNames, encodeTarget(dep))
+		}
+		if len(variationNames) > 1 {
+			modules := mctx.CreateVariations(variationNames...)
+			for i, name := range variationNames {
+				if name == "" {
+					continue
+				}
+				variation := modules[i].(*Module)
+				variation.Properties.PreventInstall = true
+				variation.Properties.HideFromMake = true
+				variation.afdo.Properties.AfdoTarget = proptools.StringPtr(decodeTarget(name))
+			}
+		}
+	}
+}
+
+// Encode target name to variation name.
+func encodeTarget(target string) string {
+	if target == "" {
+		return ""
+	}
+	return "afdo-" + target
+}
+
+// Decode target name from variation name.
+func decodeTarget(variation string) string {
+	if variation == "" {
+		return ""
+	}
+	return strings.TrimPrefix(variation, "afdo-")
+}
diff --git a/cc/cc.go b/cc/cc.go
index 22baf30..5f63615 100644
--- a/cc/cc.go
+++ b/cc/cc.go
@@ -64,6 +64,9 @@
 
 		ctx.BottomUp("coverage", coverageMutator).Parallel()
 
+		ctx.TopDown("afdo_deps", afdoDepsMutator)
+		ctx.BottomUp("afdo", afdoMutator).Parallel()
+
 		ctx.TopDown("lto_deps", ltoDepsMutator)
 		ctx.BottomUp("lto", ltoMutator).Parallel()
 
@@ -810,6 +813,7 @@
 	sabi     *sabi
 	vndkdep  *vndkdep
 	lto      *lto
+	afdo     *afdo
 	pgo      *pgo
 
 	library libraryInterface
@@ -1143,6 +1147,9 @@
 	if c.lto != nil {
 		c.AddProperties(c.lto.props()...)
 	}
+	if c.afdo != nil {
+		c.AddProperties(c.afdo.props()...)
+	}
 	if c.pgo != nil {
 		c.AddProperties(c.pgo.props()...)
 	}
@@ -1620,6 +1627,7 @@
 	module.sabi = &sabi{}
 	module.vndkdep = &vndkdep{}
 	module.lto = &lto{}
+	module.afdo = &afdo{}
 	module.pgo = &pgo{}
 	return module
 }
@@ -1815,6 +1823,9 @@
 	if c.lto != nil {
 		flags = c.lto.flags(ctx, flags)
 	}
+	if c.afdo != nil {
+		flags = c.afdo.flags(ctx, flags)
+	}
 	if c.pgo != nil {
 		flags = c.pgo.flags(ctx, flags)
 	}
@@ -1942,6 +1953,9 @@
 	if c.lto != nil {
 		c.lto.begin(ctx)
 	}
+	if c.afdo != nil {
+		c.afdo.begin(ctx)
+	}
 	if c.pgo != nil {
 		c.pgo.begin(ctx)
 	}
@@ -3534,6 +3548,7 @@
 		&SAbiProperties{},
 		&VndkProperties{},
 		&LTOProperties{},
+		&AfdoProperties{},
 		&PgoProperties{},
 		&android.ProtoProperties{},
 		// RustBindgenProperties is included here so that cc_defaults can be used for rust_bindgen modules.