Retrieve dex implementation jars from apex for java_sdk_library_import

Bug: 181267622
Test: m nothing
Change-Id: Idd6af2482f48bd3a05db88b8a06dbbbdee01ef78
diff --git a/apex/apex_test.go b/apex/apex_test.go
index 3f56047..750e68b 100644
--- a/apex/apex_test.go
+++ b/apex/apex_test.go
@@ -191,6 +191,7 @@
 		"AppSet.apks":                                nil,
 		"foo.rs":                                     nil,
 		"libfoo.jar":                                 nil,
+		"libbar.jar":                                 nil,
 	}
 
 	cc.GatherRequiredFilesForTest(fs)
@@ -4366,14 +4367,15 @@
 		// Make sure the import has been given the correct path to the dex jar.
 		p := ctx.ModuleForTests(name, "android_common_myapex").Module().(java.UsesLibraryDependency)
 		dexJarBuildPath := p.DexJarBuildPath()
-		if expected, actual := ".intermediates/myapex.deapexer/android_common/deapexer/javalib/libfoo.jar", android.NormalizePathForTesting(dexJarBuildPath); actual != expected {
+		stem := android.RemoveOptionalPrebuiltPrefix(name)
+		if expected, actual := ".intermediates/myapex.deapexer/android_common/deapexer/javalib/"+stem+".jar", android.NormalizePathForTesting(dexJarBuildPath); actual != expected {
 			t.Errorf("Incorrect DexJarBuildPath value '%s', expected '%s'", actual, expected)
 		}
 	}
 
-	ensureNoSourceVariant := func(t *testing.T, ctx *android.TestContext) {
+	ensureNoSourceVariant := func(t *testing.T, ctx *android.TestContext, name string) {
 		// Make sure that an apex variant is not created for the source module.
-		if expected, actual := []string{"android_common"}, ctx.ModuleVariantsForTests("libfoo"); !reflect.DeepEqual(expected, actual) {
+		if expected, actual := []string{"android_common"}, ctx.ModuleVariantsForTests(name); !reflect.DeepEqual(expected, actual) {
 			t.Errorf("invalid set of variants for %q: expected %q, found %q", "libfoo", expected, actual)
 		}
 	}
@@ -4390,19 +4392,28 @@
 					src: "myapex-arm.apex",
 				},
 			},
-			exported_java_libs: ["libfoo"],
+			exported_java_libs: ["libfoo", "libbar"],
 		}
 
 		java_import {
 			name: "libfoo",
 			jars: ["libfoo.jar"],
 		}
+
+		java_sdk_library_import {
+			name: "libbar",
+			public: {
+				jars: ["libbar.jar"],
+			},
+		}
 	`
 
 		// Make sure that dexpreopt can access dex implementation files from the prebuilt.
 		ctx := testDexpreoptWithApexes(t, bp, "", transform)
 
 		checkDexJarBuildPath(t, ctx, "libfoo")
+
+		checkDexJarBuildPath(t, ctx, "libbar")
 	})
 
 	t.Run("prebuilt with source preferred", func(t *testing.T) {
@@ -4418,7 +4429,7 @@
 					src: "myapex-arm.apex",
 				},
 			},
-			exported_java_libs: ["libfoo"],
+			exported_java_libs: ["libfoo", "libbar"],
 		}
 
 		java_import {
@@ -4429,13 +4440,29 @@
 		java_library {
 			name: "libfoo",
 		}
+
+		java_sdk_library_import {
+			name: "libbar",
+			public: {
+				jars: ["libbar.jar"],
+			},
+		}
+
+		java_sdk_library {
+			name: "libbar",
+			srcs: ["foo/bar/MyClass.java"],
+			unsafe_ignore_missing_latest_api: true,
+		}
 	`
 
 		// Make sure that dexpreopt can access dex implementation files from the prebuilt.
 		ctx := testDexpreoptWithApexes(t, bp, "", transform)
 
 		checkDexJarBuildPath(t, ctx, "prebuilt_libfoo")
-		ensureNoSourceVariant(t, ctx)
+		ensureNoSourceVariant(t, ctx, "libfoo")
+
+		checkDexJarBuildPath(t, ctx, "prebuilt_libbar")
+		ensureNoSourceVariant(t, ctx, "libbar")
 	})
 
 	t.Run("prebuilt preferred with source", func(t *testing.T) {
@@ -4450,7 +4477,7 @@
 					src: "myapex-arm.apex",
 				},
 			},
-			exported_java_libs: ["libfoo"],
+			exported_java_libs: ["libfoo", "libbar"],
 		}
 
 		java_import {
@@ -4462,13 +4489,30 @@
 		java_library {
 			name: "libfoo",
 		}
+
+		java_sdk_library_import {
+			name: "libbar",
+			prefer: true,
+			public: {
+				jars: ["libbar.jar"],
+			},
+		}
+
+		java_sdk_library {
+			name: "libbar",
+			srcs: ["foo/bar/MyClass.java"],
+			unsafe_ignore_missing_latest_api: true,
+		}
 	`
 
 		// Make sure that dexpreopt can access dex implementation files from the prebuilt.
 		ctx := testDexpreoptWithApexes(t, bp, "", transform)
 
 		checkDexJarBuildPath(t, ctx, "prebuilt_libfoo")
-		ensureNoSourceVariant(t, ctx)
+		ensureNoSourceVariant(t, ctx, "libfoo")
+
+		checkDexJarBuildPath(t, ctx, "prebuilt_libbar")
+		ensureNoSourceVariant(t, ctx, "libbar")
 	})
 }
 
@@ -6319,6 +6363,9 @@
 	}
 	cc.GatherRequiredFilesForTest(fs)
 
+	for k, v := range filesForSdkLibrary {
+		fs[k] = v
+	}
 	config := android.TestArchConfig(buildDir, nil, bp, fs)
 
 	ctx := android.NewTestArchContext(config)
diff --git a/java/sdk_library.go b/java/sdk_library.go
index 90823a0..e0100d4 100644
--- a/java/sdk_library.go
+++ b/java/sdk_library.go
@@ -1786,6 +1786,9 @@
 	// The reference to the xml permissions module created by the source module.
 	// Is nil if the source module does not exist.
 	xmlPermissionsFileModule *sdkLibraryXml
+
+	// Path to the dex implementation jar obtained from the prebuilt_apex, if any.
+	dexJarFile android.Path
 }
 
 var _ SdkLibraryDependency = (*SdkLibraryImport)(nil)
@@ -1982,6 +1985,8 @@
 func (module *SdkLibraryImport) GenerateAndroidBuildActions(ctx android.ModuleContext) {
 	module.generateCommonBuildActions(ctx)
 
+	var deapexerModule android.Module
+
 	// Record the paths to the prebuilt stubs library and stubs source.
 	ctx.VisitDirectDeps(func(to android.Module) {
 		tag := ctx.OtherModuleDependencyTag(to)
@@ -2007,6 +2012,11 @@
 				ctx.ModuleErrorf("xml permissions file module must be of type *sdkLibraryXml but was %T", to)
 			}
 		}
+
+		// Save away the `deapexer` module on which this depends, if any.
+		if tag == android.DeapexerTag {
+			deapexerModule = to
+		}
 	})
 
 	// Populate the scope paths with information from the properties.
@@ -2019,6 +2029,30 @@
 		paths.currentApiFilePath = android.OptionalPathForModuleSrc(ctx, scopeProperties.Current_api)
 		paths.removedApiFilePath = android.OptionalPathForModuleSrc(ctx, scopeProperties.Removed_api)
 	}
+
+	if ctx.Device() {
+		// If this is a variant created for a prebuilt_apex then use the dex implementation jar
+		// obtained from the associated deapexer module.
+		ai := ctx.Provider(android.ApexInfoProvider).(android.ApexInfo)
+		if ai.ForPrebuiltApex {
+			if deapexerModule == nil {
+				// This should never happen as a variant for a prebuilt_apex is only created if the
+				// deapxer module has been configured to export the dex implementation jar for this module.
+				ctx.ModuleErrorf("internal error: module %q does not depend on a `deapexer` module for prebuilt_apex %q",
+					module.Name(), ai.ApexVariationName)
+			}
+
+			// Get the path of the dex implementation jar from the `deapexer` module.
+			di := ctx.OtherModuleProvider(deapexerModule, android.DeapexerProvider).(android.DeapexerInfo)
+			if dexOutputPath := di.PrebuiltExportPath(module.BaseModuleName(), ".dexjar"); dexOutputPath != nil {
+				module.dexJarFile = dexOutputPath
+			} else {
+				// This should never happen as a variant for a prebuilt_apex is only created if the
+				// prebuilt_apex has been configured to export the java library dex file.
+				ctx.ModuleErrorf("internal error: no dex implementation jar available from prebuilt_apex %q", deapexerModule.Name())
+			}
+		}
+	}
 }
 
 func (module *SdkLibraryImport) sdkJars(ctx android.BaseModuleContext, sdkVersion sdkSpec, headerJars bool) android.Paths {
@@ -2051,6 +2085,11 @@
 
 // to satisfy UsesLibraryDependency interface
 func (module *SdkLibraryImport) DexJarBuildPath() android.Path {
+	// The dex implementation jar extracted from the .apex file should be used in preference to the
+	// source.
+	if module.dexJarFile != nil {
+		return module.dexJarFile
+	}
 	if module.implLibraryModule == nil {
 		return nil
 	} else {