Automatically propagate jarjar rules for aconfig libraries

Test: treehugger
Bug: 310504781
(cherry picked from https://googleplex-android-review.googlesource.com/q/commit:97c03a6dc659102ff40793759fb3f0f18164a85b)
Merged-In: I639d12ff33175b7bed7e7d0595a40dd9b0d99367
Change-Id: I639d12ff33175b7bed7e7d0595a40dd9b0d99367
diff --git a/java/base.go b/java/base.go
index e52cedd..284ec99 100644
--- a/java/base.go
+++ b/java/base.go
@@ -17,6 +17,8 @@
 import (
 	"fmt"
 	"path/filepath"
+	"reflect"
+	"slices"
 	"strconv"
 	"strings"
 
@@ -89,6 +91,9 @@
 	// if not blank, run jarjar using the specified rules file
 	Jarjar_rules *string `android:"path,arch_variant"`
 
+	// if not blank, used as prefix to generate repackage rule
+	Jarjar_prefix *string
+
 	// If not blank, set the java version passed to javac as -source and -target
 	Java_version *string
 
@@ -425,6 +430,8 @@
 	// inserting into the bootclasspath/classpath of another compile
 	headerJarFile android.Path
 
+	repackagedHeaderJarFile android.Path
+
 	// jar file containing implementation classes including static library dependencies but no
 	// resources
 	implementationJarFile android.Path
@@ -489,6 +496,9 @@
 	// expanded Jarjar_rules
 	expandJarjarRules android.Path
 
+	// jarjar rule for inherited jarjar rules
+	repackageJarjarRules android.Path
+
 	// Extra files generated by the module type to be added as java resources.
 	extraResources android.Paths
 
@@ -518,6 +528,10 @@
 
 	// Single aconfig "cache file" merged from this module and all dependencies.
 	mergedAconfigFiles map[string]android.Paths
+
+	// Values that will be set in the JarJarProvider data for jarjar repackaging,
+	// and merged with our dependencies' rules.
+	jarjarRenameRules map[string]string
 }
 
 func (j *Module) CheckStableSdkVersion(ctx android.BaseModuleContext) error {
@@ -1072,6 +1086,19 @@
 }
 
 func (j *Module) compile(ctx android.ModuleContext, extraSrcJars, extraClasspathJars, extraCombinedJars android.Paths) {
+
+	// Auto-propagating jarjar rules
+	jarjarProviderData := j.collectJarJarRules(ctx)
+	if jarjarProviderData != nil {
+		android.SetProvider(ctx, JarJarProvider, *jarjarProviderData)
+		text := getJarJarRuleText(jarjarProviderData)
+		if text != "" {
+			ruleTextFile := android.PathForModuleOut(ctx, "repackaged-jarjar", "repackaging.txt")
+			android.WriteFileRule(ctx, ruleTextFile, text)
+			j.repackageJarjarRules = ruleTextFile
+		}
+	}
+
 	j.exportAidlIncludeDirs = android.PathsForModuleSrc(ctx, j.deviceProperties.Aidl.Export_include_dirs)
 
 	deps := j.collectDeps(ctx)
@@ -1170,7 +1197,7 @@
 			ctx.ModuleErrorf("headers_only is enabled but Turbine is disabled.")
 		}
 
-		_, j.headerJarFile =
+		_, j.headerJarFile, _ =
 			j.compileJavaHeader(ctx, uniqueJavaFiles, srcJars, deps, flags, jarName,
 				extraCombinedJars)
 		if ctx.Failed() {
@@ -1285,7 +1312,7 @@
 			// with sharding enabled. See: b/77284273.
 		}
 		extraJars := append(android.CopyOf(extraCombinedJars), kotlinHeaderJars...)
-		headerJarFileWithoutDepsOrJarjar, j.headerJarFile =
+		headerJarFileWithoutDepsOrJarjar, j.headerJarFile, j.repackagedHeaderJarFile =
 			j.compileJavaHeader(ctx, uniqueJavaFiles, srcJars, deps, flags, jarName, extraJars)
 		if ctx.Failed() {
 			return
@@ -1509,6 +1536,16 @@
 		}
 	}
 
+	// Automatic jarjar rules propagation
+	if j.repackageJarjarRules != nil {
+		repackagedJarjarFile := android.PathForModuleOut(ctx, "repackaged-jarjar", jarName).OutputPath
+		TransformJarJar(ctx, repackagedJarjarFile, outputFile, j.repackageJarjarRules)
+		outputFile = repackagedJarjarFile
+		if ctx.Failed() {
+			return
+		}
+	}
+
 	// Check package restrictions if necessary.
 	if len(j.properties.Permitted_packages) > 0 {
 		// Time stamp file created by the package check rule.
@@ -1678,6 +1715,7 @@
 
 	android.SetProvider(ctx, JavaInfoProvider, JavaInfo{
 		HeaderJars:                     android.PathsIfNonNil(j.headerJarFile),
+		RepackagedHeaderJars:           android.PathsIfNonNil(j.repackagedHeaderJarFile),
 		TransitiveLibsHeaderJars:       j.transitiveLibsHeaderJars,
 		TransitiveStaticLibsHeaderJars: j.transitiveStaticLibsHeaderJars,
 		ImplementationAndResourcesJars: android.PathsIfNonNil(j.implementationAndResourcesJar),
@@ -1813,7 +1851,7 @@
 
 func (j *Module) compileJavaHeader(ctx android.ModuleContext, srcFiles, srcJars android.Paths,
 	deps deps, flags javaBuilderFlags, jarName string,
-	extraJars android.Paths) (headerJar, jarjarAndDepsHeaderJar android.Path) {
+	extraJars android.Paths) (headerJar, jarjarAndDepsHeaderJar, jarjarAndDepsRepackagedHeaderJar android.Path) {
 
 	var jars android.Paths
 	if len(srcFiles) > 0 || len(srcJars) > 0 {
@@ -1821,7 +1859,7 @@
 		turbineJar := android.PathForModuleOut(ctx, "turbine", jarName)
 		TransformJavaToHeaderClasses(ctx, turbineJar, srcFiles, srcJars, flags)
 		if ctx.Failed() {
-			return nil, nil
+			return nil, nil, nil
 		}
 		jars = append(jars, turbineJar)
 		headerJar = turbineJar
@@ -1846,11 +1884,22 @@
 		TransformJarJar(ctx, jarjarFile, jarjarAndDepsHeaderJar, j.expandJarjarRules)
 		jarjarAndDepsHeaderJar = jarjarFile
 		if ctx.Failed() {
-			return nil, nil
+			return nil, nil, nil
 		}
 	}
 
-	return headerJar, jarjarAndDepsHeaderJar
+	if j.repackageJarjarRules != nil {
+		repackagedJarjarFile := android.PathForModuleOut(ctx, "repackaged-turbine-jarjar", jarName)
+		TransformJarJar(ctx, repackagedJarjarFile, jarjarAndDepsHeaderJar, j.repackageJarjarRules)
+		jarjarAndDepsRepackagedHeaderJar = repackagedJarjarFile
+		if ctx.Failed() {
+			return nil, nil, nil
+		}
+	} else {
+		jarjarAndDepsRepackagedHeaderJar = jarjarAndDepsHeaderJar
+	}
+
+	return headerJar, jarjarAndDepsHeaderJar, jarjarAndDepsRepackagedHeaderJar
 }
 
 func (j *Module) instrument(ctx android.ModuleContext, flags javaBuilderFlags,
@@ -2207,6 +2256,10 @@
 				}
 				deps.classpath = append(deps.classpath, dep.HeaderJars...)
 				deps.dexClasspath = append(deps.dexClasspath, dep.HeaderJars...)
+				if len(dep.RepackagedHeaderJars) == 1 && !slices.Contains(dep.HeaderJars, dep.RepackagedHeaderJars[0]) {
+					deps.classpath = append(deps.classpath, dep.RepackagedHeaderJars...)
+					deps.dexClasspath = append(deps.dexClasspath, dep.RepackagedHeaderJars...)
+				}
 				deps.aidlIncludeDirs = append(deps.aidlIncludeDirs, dep.AidlIncludeDirs...)
 				addPlugins(&deps, dep.ExportedPlugins, dep.ExportedPluginClasses...)
 				deps.disableTurbine = deps.disableTurbine || dep.ExportedPluginDisableTurbine
@@ -2311,6 +2364,187 @@
 	return deps
 }
 
+// Provider for jarjar renaming rules.
+//
+// Modules can set their jarjar renaming rules with addJarJarRenameRule, and those renamings will be
+// passed to all rdeps.  The typical way that these renamings will NOT be inherited is when a module
+// links against stubs -- these are not passed through stubs. The classes will remain unrenamed on
+// classes until a module with jarjar_prefix is reached, and all as yet unrenamed classes will then
+// be renamed from that module.
+// TODO: Add another property to suppress the forwarding of
+type JarJarProviderData struct {
+	// Mapping of class names: original --> renamed.  If the value is "", the class will be
+	// renamed by the next rdep that has the jarjar_prefix attribute (or this module if it has
+	// attribute). Rdeps of that module will inherit the renaming.
+	Rename map[string]string
+}
+
+func (this JarJarProviderData) GetDebugString() string {
+	result := ""
+	for k, v := range this.Rename {
+		if strings.Contains(k, "android.companion.virtual.flags.FakeFeatureFlagsImpl") {
+			result += k + "-->" + v + ";"
+		}
+	}
+	return result
+}
+
+var JarJarProvider = blueprint.NewProvider[JarJarProviderData]()
+
+var overridableJarJarPrefix = "com.android.internal.hidden_from_bootclasspath"
+
+func init() {
+	android.SetJarJarPrefixHandler(mergeJarJarPrefixes)
+}
+
+// BaseJarJarProviderData contains information that will propagate across dependencies regardless of
+// whether they are java modules or not.
+type BaseJarJarProviderData struct {
+	JarJarProviderData JarJarProviderData
+}
+
+func (this BaseJarJarProviderData) GetDebugString() string {
+	return this.JarJarProviderData.GetDebugString()
+}
+
+var BaseJarJarProvider = blueprint.NewProvider[BaseJarJarProviderData]()
+
+// mergeJarJarPrefixes is called immediately before module.GenerateAndroidBuildActions is called.
+// Since there won't be a JarJarProvider, we create the BaseJarJarProvider if any of our deps have
+// either JarJarProvider or BaseJarJarProvider.
+func mergeJarJarPrefixes(ctx android.ModuleContext) {
+	mod := ctx.Module()
+	// Explicitly avoid propagating into some module types.
+	switch reflect.TypeOf(mod).String() {
+	case "*java.Droidstubs":
+		return
+	}
+	jarJarData := collectDirectDepsProviders(ctx)
+	if jarJarData != nil {
+		providerData := BaseJarJarProviderData{
+			JarJarProviderData: *jarJarData,
+		}
+		android.SetProvider(ctx, BaseJarJarProvider, providerData)
+	}
+
+}
+
+// Add a jarjar renaming rule to this module, to be inherited to all dependent modules.
+func (module *Module) addJarJarRenameRule(original string, renamed string) {
+	if module.jarjarRenameRules == nil {
+		module.jarjarRenameRules = make(map[string]string)
+	}
+	module.jarjarRenameRules[original] = renamed
+}
+
+func collectDirectDepsProviders(ctx android.ModuleContext) (result *JarJarProviderData) {
+	// Gather repackage information from deps
+	// If the dep jas a JarJarProvider, it is used.  Otherwise, any BaseJarJarProvider is used.
+	ctx.VisitDirectDepsIgnoreBlueprint(func(m android.Module) {
+		merge := func(theirs *JarJarProviderData) {
+			for orig, renamed := range theirs.Rename {
+				if result == nil {
+					result = &JarJarProviderData{
+						Rename: make(map[string]string),
+					}
+				}
+				if preexisting, exists := (*result).Rename[orig]; !exists || preexisting == "" {
+					result.Rename[orig] = renamed
+				} else if preexisting != "" && renamed != "" && preexisting != renamed {
+					if strings.HasPrefix(preexisting, overridableJarJarPrefix) {
+						result.Rename[orig] = renamed
+					} else if !strings.HasPrefix(renamed, overridableJarJarPrefix) {
+						ctx.ModuleErrorf("1. Conflicting jarjar rules inherited for class: %s (%s and %s)", orig, renamed, preexisting, ctx.ModuleName(), m.Name())
+						continue
+					}
+				}
+			}
+		}
+		if theirs, ok := android.OtherModuleProvider(ctx, m, JarJarProvider); ok {
+			merge(&theirs)
+		} else if theirs, ok := android.OtherModuleProvider(ctx, m, BaseJarJarProvider); ok {
+			// TODO: if every java.Module should have a JarJarProvider, and we find only the
+			// BaseJarJarProvider, then there is a bug.  Consider seeing if m can be cast
+			// to java.Module.
+			merge(&theirs.JarJarProviderData)
+		}
+	})
+	return
+}
+
+func (this Module) GetDebugString() string {
+	return "sdk_version=" + proptools.String(this.deviceProperties.Sdk_version)
+}
+
+// Merge the jarjar rules we inherit from our dependencies, any that have been added directly to
+// us, and if it's been set, apply the jarjar_prefix property to rename them.
+func (module *Module) collectJarJarRules(ctx android.ModuleContext) *JarJarProviderData {
+	// Gather repackage information from deps
+	result := collectDirectDepsProviders(ctx)
+
+	// Update that with entries we've stored for ourself
+	for orig, renamed := range module.jarjarRenameRules {
+		if result == nil {
+			result = &JarJarProviderData{
+				Rename: make(map[string]string),
+			}
+		}
+		if renamed != "" {
+			if preexisting, exists := (*result).Rename[orig]; exists && preexisting != renamed {
+				ctx.ModuleErrorf("Conflicting jarjar rules inherited for class: %s (%s and %s)", orig, renamed, preexisting)
+				continue
+			}
+		}
+		(*result).Rename[orig] = renamed
+	}
+
+	// If there are no renamings, then jarjar_prefix does nothing, so skip the extra work.
+	if result == nil {
+		return nil
+	}
+
+	// If they've given us a jarjar_prefix property, then we will use that to rename any classes
+	// that have not yet been renamed.
+	prefix := proptools.String(module.properties.Jarjar_prefix)
+	if prefix != "" {
+		if prefix[0] == '.' {
+			ctx.PropertyErrorf("jarjar_prefix", "jarjar_prefix can not start with '.'")
+			return nil
+		}
+		if prefix[len(prefix)-1] == '.' {
+			ctx.PropertyErrorf("jarjar_prefix", "jarjar_prefix can not end with '.'")
+			return nil
+		}
+
+		var updated map[string]string
+		for orig, renamed := range (*result).Rename {
+			if renamed == "" {
+				if updated == nil {
+					updated = make(map[string]string)
+				}
+				updated[orig] = prefix + "." + orig
+			}
+		}
+		for orig, renamed := range updated {
+			(*result).Rename[orig] = renamed
+		}
+	}
+
+	return result
+}
+
+// Get the jarjar rule text for a given provider for the fully resolved rules. Classes that map
+// to "" won't be in this list because they shouldn't be renamed yet.
+func getJarJarRuleText(provider *JarJarProviderData) string {
+	result := ""
+	for orig, renamed := range provider.Rename {
+		if renamed != "" {
+			result += "rule " + orig + " " + renamed + "\n"
+		}
+	}
+	return result
+}
+
 func addPlugins(deps *deps, pluginJars android.Paths, pluginClasses ...string) {
 	deps.processorPath = append(deps.processorPath, pluginJars...)
 	deps.processorClasses = append(deps.processorClasses, pluginClasses...)