Enforce partition property on apex system server jars

This is a followup AI from https://r.android.com/3288083, which installs
dexpreopt files of apex system server jars in the same partition as the
top-level apex. This CL enforces that the partition properties of the
apex and the java library match.

This has been implemented in a bottomup manner.
systemserverclasspath_fragment module sets
`LibraryNameToPartitionInfoProvider`. The top-level apex uses this info
and raies an exception if the partition properties do not match.

This enforcement is done only for source apexes for now. It is not
needed for prebuilts since
- The dexpreopt rules of jars in prebuilt apexes are generated in the
  top-level `prebuiltApex`/ `apexSet`
- We do not have prebuilts of system_ext apexes today
  (com.android.compos is not part of mainline)

Test: m nothing --no-skip-soong-tests

Change-Id: Ib4df634457f315d5421681bdb0afebc2b1bf92d9
diff --git a/apex/apex.go b/apex/apex.go
index d7dc6d7..d3e7eee 100644
--- a/apex/apex.go
+++ b/apex/apex.go
@@ -29,6 +29,7 @@
 	"android/soong/android"
 	"android/soong/bpf"
 	"android/soong/cc"
+	"android/soong/dexpreopt"
 	prebuilt_etc "android/soong/etc"
 	"android/soong/filesystem"
 	"android/soong/java"
@@ -1919,6 +1920,32 @@
 	})
 }
 
+// enforcePartitionTagOnApexSystemServerJar checks that the partition tags of an apex system server jar  matches
+// the partition tags of the top-level apex.
+// e.g. if the top-level apex sets system_ext_specific to true, the javalib must set this property to true as well.
+// This check ensures that the dexpreopt artifacts of the apex system server jar is installed in the same partition
+// as the apex.
+func (a *apexBundle) enforcePartitionTagOnApexSystemServerJar(ctx android.ModuleContext) {
+	global := dexpreopt.GetGlobalConfig(ctx)
+	ctx.VisitDirectDepsWithTag(sscpfTag, func(child android.Module) {
+		info, ok := android.OtherModuleProvider(ctx, child, java.LibraryNameToPartitionInfoProvider)
+		if !ok {
+			ctx.ModuleErrorf("Could not find partition info of apex system server jars.")
+		}
+		apexPartition := ctx.Module().PartitionTag(ctx.DeviceConfig())
+		for javalib, javalibPartition := range info.LibraryNameToPartition {
+			if !global.AllApexSystemServerJars(ctx).ContainsJar(javalib) {
+				continue // not an apex system server jar
+			}
+			if apexPartition != javalibPartition {
+				ctx.ModuleErrorf(`
+%s is an apex systemserver jar, but its partition does not match the partition of its containing apex. Expected %s, Got %s`,
+					javalib, apexPartition, javalibPartition)
+			}
+		}
+	})
+}
+
 func (a *apexBundle) depVisitor(vctx *visitorContext, ctx android.ModuleContext, child, parent android.Module) bool {
 	depTag := ctx.OtherModuleDependencyTag(child)
 	if _, ok := depTag.(android.ExcludeFromApexContentsTag); ok {
@@ -2341,6 +2368,7 @@
 	a.required = append(a.required, a.VintfFragmentModuleNames(ctx)...)
 
 	a.setOutputFiles(ctx)
+	a.enforcePartitionTagOnApexSystemServerJar(ctx)
 }
 
 // Set prebuiltInfoProvider. This will be used by `apex_prebuiltinfo_singleton` to print out a metadata file