Install dexpreopt artifacts of apex system server jars in same partition

e.g. If a system server jar providing apex is installed in /system_ext,
the dexpreopt files of that jar will also be installed in /system_ext.
Currently, all of these artifacts are installed in /system

This behavior will be flag guarded by
RELEASE_INSTALL_APEX_SYSTEMSERVER_DEXPREOPT_SAME_PARTITION. This is
necessary because the ART runtime needs to be updated to look in the new
/system_ext location. Since some release configs build with ART
prebuilts, the runtime in those prebuilt apexes will not have the
additional search path introduced in https://r.android.com/3287191

Test: Verified that "Could not check odex file" does not appear for
service-compos (a system_ext apex systemserver jar)

Bug: 369678122
Change-Id: I752bdc7f5f69226b503800ce25726a211302cb07
diff --git a/dexpreopt/config.go b/dexpreopt/config.go
index fe6317c..84d4f10 100644
--- a/dexpreopt/config.go
+++ b/dexpreopt/config.go
@@ -191,6 +191,10 @@
 	ForceCreateAppImage bool
 
 	PresignedPrebuilt bool
+
+	// ApexPartition is the partition in which the dexpreopt files of apex system server jars (if any) are installed.
+	// This is a noop unless the module is apex system server jar.
+	ApexPartition string
 }
 
 type globalSoongConfigSingleton struct{}
diff --git a/dexpreopt/dexpreopt.go b/dexpreopt/dexpreopt.go
index 5616483..7a39fa1 100644
--- a/dexpreopt/dexpreopt.go
+++ b/dexpreopt/dexpreopt.go
@@ -219,9 +219,9 @@
 }
 
 // Returns the location to the odex file for the dex file at `path`.
-func ToOdexPath(path string, arch android.ArchType) string {
+func ToOdexPath(path string, arch android.ArchType, partition string) string {
 	if strings.HasPrefix(path, "/apex/") {
-		return filepath.Join("/system/framework/oat", arch.String(),
+		return filepath.Join(partition, "framework/oat", arch.String(),
 			strings.ReplaceAll(path[1:], "/", "@")+"@classes.odex")
 	}
 
@@ -245,7 +245,7 @@
 
 	odexPath := module.BuildPath.InSameDir(ctx, "oat", arch.String(), pathtools.ReplaceExtension(base, "odex"))
 	odexSymbolsPath := odexPath.ReplaceExtension(ctx, "symbols.odex")
-	odexInstallPath := ToOdexPath(module.DexLocation, arch)
+	odexInstallPath := ToOdexPath(module.DexLocation, arch, module.ApexPartition)
 	if odexOnSystemOther(module, global) {
 		odexInstallPath = filepath.Join(SystemOtherPartition, odexInstallPath)
 	}
diff --git a/dexpreopt/dexpreopt_test.go b/dexpreopt/dexpreopt_test.go
index 6f7d3bb..7b0f51f 100644
--- a/dexpreopt/dexpreopt_test.go
+++ b/dexpreopt/dexpreopt_test.go
@@ -42,12 +42,14 @@
 }
 
 func testApexModuleConfig(ctx android.PathContext, name, apexName string) *ModuleConfig {
-	return createTestModuleConfig(
+	ret := createTestModuleConfig(
 		name,
 		fmt.Sprintf("/apex/%s/javalib/%s.jar", apexName, name),
 		android.PathForOutput(ctx, fmt.Sprintf("%s/dexpreopt/%s.jar", name, name)),
 		android.PathForOutput(ctx, fmt.Sprintf("%s/aligned/%s.jar", name, name)),
 		android.PathForOutput(ctx, fmt.Sprintf("%s/enforce_uses_libraries.status", name)))
+	ret.ApexPartition = "/system"
+	return ret
 }
 
 func testPlatformSystemServerModuleConfig(ctx android.PathContext, name string) *ModuleConfig {
@@ -221,6 +223,49 @@
 	DexpreoptRunningInSoong = oldDexpreoptRunningInSoong
 }
 
+// Same as `TestDexPreoptApexSystemServerJars`, but the apex jar is in /system_ext
+func TestDexPreoptApexSystemServerJarsSystemExt(t *testing.T) {
+	// modify the global variable for test
+	var oldDexpreoptRunningInSoong = DexpreoptRunningInSoong
+	DexpreoptRunningInSoong = true
+
+	// test begin
+	config := android.TestConfig("out", nil, "", nil)
+	ctx := android.BuilderContextForTesting(config)
+	globalSoong := globalSoongConfigForTests(ctx)
+	global := GlobalConfigForTests(ctx)
+	module := testApexModuleConfig(ctx, "service-A", "com.android.apex1")
+	module.ApexPartition = "/system_ext"
+	productPackages := android.PathForTesting("product_packages.txt")
+
+	global.ApexSystemServerJars = android.CreateTestConfiguredJarList(
+		[]string{"com.android.apex1:service-A"})
+
+	rule, err := GenerateDexpreoptRule(ctx, globalSoong, global, module, productPackages, true)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	wantInstalls := android.RuleBuilderInstalls{
+		{android.PathForOutput(ctx, "service-A/dexpreopt/oat/arm/javalib.odex"), "/system_ext/framework/oat/arm/apex@com.android.apex1@javalib@service-A.jar@classes.odex"},
+		{android.PathForOutput(ctx, "service-A/dexpreopt/oat/arm/javalib.vdex"), "/system_ext/framework/oat/arm/apex@com.android.apex1@javalib@service-A.jar@classes.vdex"},
+	}
+
+	android.AssertStringEquals(t, "installs", wantInstalls.String(), rule.Installs().String())
+
+	android.AssertStringListContains(t, "apex sscp jar copy", rule.Outputs().Strings(), "out/soong/system_server_dexjars/service-A.jar")
+
+	// rule with apex sscp cp as false
+	rule, err = GenerateDexpreoptRule(ctx, globalSoong, global, module, productPackages, false)
+	if err != nil {
+		t.Fatal(err)
+	}
+	android.AssertStringListDoesNotContain(t, "apex sscp jar copy", rule.Outputs().Strings(), "out/soong/system_server_dexjars/service-A.jar")
+
+	// cleanup the global variable for test
+	DexpreoptRunningInSoong = oldDexpreoptRunningInSoong
+}
+
 func TestDexPreoptStandaloneSystemServerJars(t *testing.T) {
 	config := android.TestConfig("out", nil, "", nil)
 	ctx := android.BuilderContextForTesting(config)