From 46d66de1c1a52ae3842bb2453de026e64a2d54da Mon Sep 17 00:00:00 2001
From: Jihoon Kang <jihoonkang@google.com>
Date: Wed, 22 May 2024 22:42:39 +0000
Subject: [PATCH] Propagate DirectlyInAnyApex to transitive dependencies

`UpdateDirectlyInAnyApex` is used to propagate the ApexProperties
`DirectlyInAnyApex` and `InAnyApex` to the direct dependencies of the
direct dependencies of an apex bundle. In other words, this will be
propagated only to two-levels max dependency from the apex bundle.

However, the implementation library of the sdk library can have longer
dependency chain from the apex bundle than two levels:
e.g. apex -> bcpf -> sdk_lib -> sdk_lib impl lib

Therefore, even if the implementation library of the sdk library is
registered as a dependency using the "CopyDirectlyInAnyApexTag"
dependency tag, the ApexProperties would not be propagated to the
implementation library. In order to resolve this issue and recognize
the implementation library to be directly in apex and allow
instrumentation for the implementation library, this change proposes
propagating `DirectlyInAnyApex` and `InAnyApex` to transitive
dependencies on top of direct dependencies.

Test: DIST_DIR=dist_dir TARGET_BUILD_VARIANT=userdebug PRODUCT=mainline_modules_x86_64 COVERAGE_MODULES="uwb" ./vendor/google/build/build_unbundled_coverage_mainline_module.sh && \
unzip -l out/target/product/module_x86_64/jacoco-report-classes-all.jar and ensure that framework-uwb is included
Bug: 341170242

Change-Id: I27d7a74f6e5bc3e0a044d13c619f4897b6b2eb57
---
 android/apex.go     | 19 ++++++++++++-------
 java/base.go        |  6 +-----
 java/sdk_library.go |  6 ++++++
 3 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/android/apex.go b/android/apex.go
index fcbd13ea2..2ac6ed08d 100644
--- a/android/apex.go
+++ b/android/apex.go
@@ -350,7 +350,8 @@ type SkipApexAllowedDependenciesCheck interface {
 // ApexModuleBase provides the default implementation for the ApexModule interface. APEX-aware
 // modules are expected to include this struct and call InitApexModule().
 type ApexModuleBase struct {
-	ApexProperties ApexProperties
+	ApexProperties     ApexProperties
+	apexPropertiesLock sync.Mutex // protects ApexProperties during parallel apexDirectlyInAnyMutator
 
 	canHaveApexVariants bool
 
@@ -761,18 +762,22 @@ func UpdateUniqueApexVariationsForDeps(mctx BottomUpMutatorContext, am ApexModul
 
 // UpdateDirectlyInAnyApex uses the final module to store if any variant of this module is directly
 // in any APEX, and then copies the final value to all the modules. It also copies the
-// DirectlyInAnyApex value to any direct dependencies with a CopyDirectlyInAnyApexTag dependency
-// tag.
+// DirectlyInAnyApex value to any transitive dependencies with a CopyDirectlyInAnyApexTag
+// dependency tag.
 func UpdateDirectlyInAnyApex(mctx BottomUpMutatorContext, am ApexModule) {
 	base := am.apexModuleBase()
-	// Copy DirectlyInAnyApex and InAnyApex from any direct dependencies with a
+	// Copy DirectlyInAnyApex and InAnyApex from any transitive dependencies with a
 	// CopyDirectlyInAnyApexTag dependency tag.
-	mctx.VisitDirectDeps(func(dep Module) {
-		if _, ok := mctx.OtherModuleDependencyTag(dep).(CopyDirectlyInAnyApexTag); ok {
-			depBase := dep.(ApexModule).apexModuleBase()
+	mctx.WalkDeps(func(child, parent Module) bool {
+		if _, ok := mctx.OtherModuleDependencyTag(child).(CopyDirectlyInAnyApexTag); ok {
+			depBase := child.(ApexModule).apexModuleBase()
+			depBase.apexPropertiesLock.Lock()
+			defer depBase.apexPropertiesLock.Unlock()
 			depBase.ApexProperties.DirectlyInAnyApex = base.ApexProperties.DirectlyInAnyApex
 			depBase.ApexProperties.InAnyApex = base.ApexProperties.InAnyApex
+			return true
 		}
+		return false
 	})
 
 	if base.ApexProperties.DirectlyInAnyApex {
diff --git a/java/base.go b/java/base.go
index d04e97cb1..b9e236cf0 100644
--- a/java/base.go
+++ b/java/base.go
@@ -716,11 +716,7 @@ func (j *Module) shouldInstrumentInApex(ctx android.BaseModuleContext) bool {
 	apexInfo, _ := android.ModuleProvider(ctx, android.ApexInfoProvider)
 	isJacocoAgent := ctx.ModuleName() == "jacocoagent"
 
-	isApexVariantSdkLibImplLib := j.SdkLibraryName() != nil &&
-		strings.HasSuffix(j.Name(), ".impl") &&
-		len(apexInfo.InApexVariants) > 0
-
-	if (j.DirectlyInAnyApex() || isApexVariantSdkLibImplLib) && !isJacocoAgent && !apexInfo.IsForPlatform() {
+	if j.DirectlyInAnyApex() && !isJacocoAgent && !apexInfo.IsForPlatform() {
 		if !inList(ctx.ModuleName(), config.InstrumentFrameworkModules) {
 			return true
 		} else if ctx.Config().IsEnvTrue("EMMA_INSTRUMENT_FRAMEWORK") {
diff --git a/java/sdk_library.go b/java/sdk_library.go
index 8c91288c3..72eb6e346 100644
--- a/java/sdk_library.go
+++ b/java/sdk_library.go
@@ -1506,6 +1506,12 @@ var implLibraryTag = sdkLibraryComponentTag{name: "impl-library"}
 
 var _ android.InstallNeededDependencyTag = sdkLibraryComponentTag{}
 
+// To satisfy the CopyDirectlyInAnyApexTag interface. Implementation library of the sdk library
+// in an apex is considered to be directly in the apex, as if it was listed in java_libs.
+func (t sdkLibraryComponentTag) CopyDirectlyInAnyApex() {}
+
+var _ android.CopyDirectlyInAnyApexTag = implLibraryTag
+
 func (t sdkLibraryComponentTag) InstallDepNeeded() bool {
 	return t.name == "xml-permissions-file" || t.name == "impl-library"
 }