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

aconfig/codegen/java_aconfig_library.go

@@ -102,6 +102,13 @@ func (callbacks *JavaAconfigDeclarationsLibraryCallbacks) GenerateSourceJarBuild
 		},
 	})
 
+	// Mark our generated code as possibly needing jarjar repackaging
+	// TODO: Maybe control this with a property?
+	module.AddJarJarRenameRule(declarations.Package+".Flags", "")
+	module.AddJarJarRenameRule(declarations.Package+".FeatureFlags", "")
+	module.AddJarJarRenameRule(declarations.Package+".FeatureFlagsImpl", "")
+	module.AddJarJarRenameRule(declarations.Package+".FakeFeatureFlagsImpl", "")
+
 	return srcJarPath
 }
 

android/module.go

@@ -34,6 +34,7 @@ import (
 var (
 	DeviceSharedLibrary = "shared_library"
 	DeviceStaticLibrary = "static_library"
+	jarJarPrefixHandler func(ctx ModuleContext)
 )
 
 type Module interface {
@@ -1772,6 +1773,13 @@ func (m *ModuleBase) GenerateBuildActions(blueprintCtx blueprint.ModuleContext)
 			return
 		}
 
+		if jarJarPrefixHandler != nil {
+			jarJarPrefixHandler(ctx)
+			if ctx.Failed() {
+				return
+			}
+		}
+
 		m.module.GenerateAndroidBuildActions(ctx)
 		if ctx.Failed() {
 			return
@@ -1865,6 +1873,13 @@ func (m *ModuleBase) GenerateBuildActions(blueprintCtx blueprint.ModuleContext)
 	m.variables = ctx.variables
 }
 
+func SetJarJarPrefixHandler(handler func(ModuleContext)) {
+	if jarJarPrefixHandler != nil {
+		panic("jarJarPrefixHandler already set")
+	}
+	jarJarPrefixHandler = handler
+}
+
 func (m *ModuleBase) moduleInfoRegisterName(ctx ModuleContext, subName string) string {
 	name := m.BaseModuleName()
 

java/base.go

@@ -17,6 +17,8 @@ package java
 import (
 	"fmt"
 	"path/filepath"
+	"reflect"
+	"slices"
 	"strconv"
 	"strings"
 
@@ -89,6 +91,9 @@ type CommonProperties struct {
 	// 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 @@ type Module struct {
 	// 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 @@ type Module struct {
 	// 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 @@ type Module struct {
 
 	// 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) addGeneratedSrcJars(path android.Path) {
 }
 
 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 @@ func (j *Module) compile(ctx android.ModuleContext, extraSrcJars, extraClasspath
 			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 @@ func (j *Module) compile(ctx android.ModuleContext, extraSrcJars, extraClasspath
 			// 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 @@ func (j *Module) compile(ctx android.ModuleContext, extraSrcJars, extraClasspath
 		}
 	}
 
+	// 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 @@ func (j *Module) compile(ctx android.ModuleContext, extraSrcJars, extraClasspath
 
 	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 CheckKotlincFlags(ctx android.ModuleContext, flags []string) {
 
 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 @@ func (j *Module) compileJavaHeader(ctx android.ModuleContext, srcFiles, srcJars
 		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 @@ func (j *Module) compileJavaHeader(ctx android.ModuleContext, srcFiles, srcJars
 		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 @@ func (j *Module) collectDeps(ctx android.ModuleContext) deps {
 				}
 				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 @@ func (j *Module) collectDeps(ctx android.ModuleContext) deps {
 	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...)

java/dex.go

@@ -262,7 +262,7 @@ func (d *dexer) r8Flags(ctx android.ModuleContext, flags javaBuilderFlags) (r8Fl
 	var proguardRaiseDeps classpath
 	ctx.VisitDirectDepsWithTag(proguardRaiseTag, func(m android.Module) {
 		dep, _ := android.OtherModuleProvider(ctx, m, JavaInfoProvider)
-		proguardRaiseDeps = append(proguardRaiseDeps, dep.HeaderJars...)
+		proguardRaiseDeps = append(proguardRaiseDeps, dep.RepackagedHeaderJars...)
 	})
 
 	r8Flags = append(r8Flags, proguardRaiseDeps.FormJavaClassPath("-libraryjars"))

java/generated_java_library.go

@@ -107,3 +107,8 @@ func (module *GeneratedJavaLibraryModule) GenerateAndroidBuildActions(ctx androi
 	module.Library.properties.Generated_srcjars = append(module.Library.properties.Generated_srcjars, srcJarPath)
 	module.Library.GenerateAndroidBuildActions(ctx)
 }
+
+// Add a rule to the jarjar renaming rules.  See RepackageProviderData.
+func (module *GeneratedJavaLibraryModule) AddJarJarRenameRule(original string, renamed string) {
+	module.addJarJarRenameRule(original, renamed)
+}

java/java.go

@@ -248,6 +248,8 @@ type JavaInfo struct {
 	// against this module.  If empty, ImplementationJars should be used instead.
 	HeaderJars android.Paths
 
+	RepackagedHeaderJars android.Paths
+
 	// set of header jars for all transitive libs deps
 	TransitiveLibsHeaderJars *android.DepSet[android.Path]