Add a trace_references_from dex optimization property

Introduce a new property `trace_references_from` in the optimize block.
This field allows specifying a set of downstream targets from which to
trace references when optimizing the current target.

In practice, this requires explicit breaking of any cycles that might
occur. A follow-up effort will explore doing this dynamically and
implicitly for `libs` references onto an optimized target that
enables this feature.

Bug: 212737576
Test: m nothing
Test: go test ./java
Flag: EXEMPT bugfix
Change-Id: I088bd8bf89403661a084a43b2822a6efbbc0294d
diff --git a/java/base.go b/java/base.go
index 0833831..d49a9bc 100644
--- a/java/base.go
+++ b/java/base.go
@@ -888,6 +888,10 @@
 	// Add dependency on libraries that provide additional hidden api annotations.
 	ctx.AddVariationDependencies(nil, hiddenApiAnnotationsTag, j.properties.Hiddenapi_additional_annotations...)
 
+	// Add dependency on (soft) downstream libs from which to trace references during optimization.
+	traceRefs := j.dexProperties.Optimize.Trace_references_from.GetOrDefault(ctx, []string{})
+	ctx.AddVariationDependencies(nil, traceReferencesTag, traceRefs...)
+
 	// For library dependencies that are component libraries (like stubs), add the implementation
 	// as a dependency (dexpreopt needs to be against the implementation library, not stubs).
 	for _, dep := range libDeps {
diff --git a/java/dex.go b/java/dex.go
index b32d5ae..ed9c82b 100644
--- a/java/dex.go
+++ b/java/dex.go
@@ -116,6 +116,21 @@
 		//
 		// By default all classes are compiled using R8 when Optimize.Enabled is set.
 		Exclude *string `android:"path"`
+
+		// Optional list of downstream (Java) libraries from which to trace and preserve references
+		// when optimizing. Note that this requires that the source reference does *not* have
+		// a strict lib dependency on this target; dependencies should be on intermediate targets
+		// statically linked into this target, e.g., if A references B, and we want to trace and
+		// keep references from A when optimizing B, you would create an intermediate B.impl (
+		// containing all static code), have A depend on `B.impl` via libs, and set
+		// `trace_references_from: ["A"]` on B.
+		//
+		// Also note that these are *not* inherited across targets, they must be specified at the
+		// top-level target that is optimized.
+		//
+		// TODO(b/212737576): Handle this implicitly using bottom-up deps mutation and implicit
+		// creation of a proxy `.impl` library.
+		Trace_references_from proptools.Configurable[[]string] `android:"arch_variant"`
 	}
 
 	// Keep the data uncompressed. We always need uncompressed dex for execution,
@@ -458,6 +473,20 @@
 
 	flagFiles = append(flagFiles, android.PathsForModuleSrc(ctx, opt.Proguard_flags_files)...)
 
+	traceReferencesSources := android.Paths{}
+	ctx.VisitDirectDepsProxyWithTag(traceReferencesTag, func(m android.ModuleProxy) {
+		if dep, ok := android.OtherModuleProvider(ctx, m, JavaInfoProvider); ok {
+			traceReferencesSources = append(traceReferencesSources, dep.ImplementationJars...)
+		}
+	})
+	if len(traceReferencesSources) > 0 {
+		traceTarget := dexParams.classesJar
+		traceLibs := android.FirstUniquePaths(append(flags.bootClasspath.Paths(), flags.dexClasspath.Paths()...))
+		traceReferencesFlags := android.PathForModuleOut(ctx, "proguard", "trace_references.flags")
+		TraceReferences(ctx, traceReferencesSources, traceTarget, traceLibs, traceReferencesFlags)
+		flagFiles = append(flagFiles, traceReferencesFlags)
+	}
+
 	flagFiles = android.FirstUniquePaths(flagFiles)
 
 	r8Flags = append(r8Flags, android.JoinWithPrefix(flagFiles.Strings(), "-include "))
diff --git a/java/dex_test.go b/java/dex_test.go
index 66d801d..e94864b 100644
--- a/java/dex_test.go
+++ b/java/dex_test.go
@@ -866,3 +866,46 @@
 		})
 	}
 }
+
+func TestTraceReferences(t *testing.T) {
+	t.Parallel()
+	bp := `
+		android_app {
+			name: "app",
+			libs: ["lib.impl"],
+			srcs: ["foo.java"],
+			platform_apis: true,
+		}
+
+		java_library {
+			name: "lib",
+			optimize: {
+				enabled: true,
+				trace_references_from: ["app"],
+			},
+			srcs: ["bar.java"],
+			static_libs: ["lib.impl"],
+			installable: true,
+		}
+
+		java_library {
+			name: "lib.impl",
+			srcs: ["baz.java"],
+		}
+	`
+	result := android.GroupFixturePreparers(
+		PrepareForTestWithJavaDefaultModules,
+	).RunTestWithBp(t, bp)
+
+	appJar := result.ModuleForTests(t, "app", "android_common").Output("combined/app.jar").Output
+	libJar := result.ModuleForTests(t, "lib", "android_common").Output("combined/lib.jar").Output
+	libTraceRefs := result.ModuleForTests(t, "lib", "android_common").Rule("traceReferences")
+	libR8 := result.ModuleForTests(t, "lib", "android_common").Rule("r8")
+
+	android.AssertStringDoesContain(t, "expected trace reference source from app jar",
+		libTraceRefs.Args["sources"], "--source "+appJar.String())
+	android.AssertStringEquals(t, "expected trace reference target into lib jar",
+		libJar.String(), libTraceRefs.Input.String())
+	android.AssertStringDoesContain(t, "expected trace reference proguard flags in lib r8 flags",
+		libR8.Args["r8Flags"], "trace_references.flags")
+}
diff --git a/java/dexpreopt.go b/java/dexpreopt.go
index b21cfc9..e8e1cd4 100644
--- a/java/dexpreopt.go
+++ b/java/dexpreopt.go
@@ -650,6 +650,9 @@
 		// for now just exclude any known irrelevant dependencies that would lead to incorrect errors.
 		if _, ok := tag.(bootclasspathDependencyTag); ok {
 			return false
+		} else if tag == traceReferencesTag {
+			// Allow ordering inversion if the dependency is purely for tracing references.
+			return false
 		}
 		depIndex := jars.IndexOfJar(dep.Name())
 		if jarIndex < depIndex && !config.BrokenSuboptimalOrderOfSystemServerJars {
diff --git a/java/java.go b/java/java.go
index 215fbbd..372e043 100644
--- a/java/java.go
+++ b/java/java.go
@@ -578,6 +578,7 @@
 	extraLintCheckTag       = dependencyTag{name: "extra-lint-check", toolchain: true}
 	jniLibTag               = dependencyTag{name: "jnilib", runtimeLinked: true}
 	r8LibraryJarTag         = dependencyTag{name: "r8-libraryjar", runtimeLinked: true}
+	traceReferencesTag      = dependencyTag{name: "trace-references"}
 	syspropPublicStubDepTag = dependencyTag{name: "sysprop public stub"}
 	javaApiContributionTag  = dependencyTag{name: "java-api-contribution"}
 	aconfigDeclarationTag   = dependencyTag{name: "aconfig-declaration"}
@@ -608,6 +609,7 @@
 		kotlinPluginTag,
 		syspropPublicStubDepTag,
 		instrumentationForTag,
+		traceReferencesTag,
 	}
 )