Specify jnilib partition in Android-<target>.mk

Context
- Android-<target>.mk currently does not contain information about partition for its dependent unembedded jni libraries, but only lists the name of the unembedded jni libraries.
- If an android_app module depends on an unembedded jni library that is located in a different partition, make cannot find the library.

Implementation
- Create a string field partition in jniLib struct.
- Add variable "LOCAL_SOONG_JNI_LIBS_PARTITION_<target>", an array of mappings of the name of the jni library to its partition.

Bug: 154162945
Test: m
Change-Id: I6b8e1272ff59dc70e3dd6ce8c6c8e4686dad76df

android/paths.go

@@ -1638,6 +1638,10 @@ func (p InstallPath) PartitionDir() string {
 	}
 }
 
+func (p InstallPath) Partition() string {
+	return p.partition
+}
+
 // Join creates a new InstallPath with paths... joined with the current path. The
 // provided paths... may not use '..' to escape from the current path.
 func (p InstallPath) Join(ctx PathContext, paths ...string) InstallPath {

cc/cc.go

@@ -3835,6 +3835,15 @@ func (ks *kytheExtractAllSingleton) GenerateBuildActions(ctx android.SingletonCo
 	}
 }
 
+func (c *Module) Partition() string {
+	if p, ok := c.installer.(interface {
+		getPartition() string
+	}); ok {
+		return p.getPartition()
+	}
+	return ""
+}
+
 var Bool = proptools.Bool
 var BoolDefault = proptools.BoolDefault
 var BoolPtr = proptools.BoolPtr

cc/library.go

@@ -2222,6 +2222,10 @@ func (library *libraryDecorator) makeUninstallable(mod *Module) {
 	mod.ModuleBase.MakeUninstallable()
 }
 
+func (library *libraryDecorator) getPartition() string {
+	return library.path.Partition()
+}
+
 func (library *libraryDecorator) getAPIListCoverageXMLPath() android.ModuleOutPath {
 	return library.apiListCoverageXmlPath
 }

cc/linkable.go

@@ -253,6 +253,9 @@ type LinkableInterface interface {
 
 	// VndkVersion returns the VNDK version string for this module.
 	VndkVersion() string
+
+	// Partition returns the partition string for this module.
+	Partition() string
 }
 
 var (

java/androidmk.go

@@ -17,6 +17,7 @@ package java
 import (
 	"fmt"
 	"io"
+	"strings"
 
 	"android/soong/android"
 )
@@ -398,6 +399,19 @@ func (app *AndroidApp) AndroidMkEntries() []android.AndroidMkEntries {
 				} else {
 					for _, jniLib := range app.jniLibs {
 						entries.AddStrings("LOCAL_SOONG_JNI_LIBS_"+jniLib.target.Arch.ArchType.String(), jniLib.name)
+						var partitionTag string
+
+						// Mimic the creation of partition_tag in build/make,
+						// which defaults to an empty string when the partition is system.
+						// Otherwise, capitalize with a leading _
+						if jniLib.partition == "system" {
+							partitionTag = ""
+						} else {
+							split := strings.Split(jniLib.partition, "/")
+							partitionTag = "_" + strings.ToUpper(split[len(split)-1])
+						}
+						entries.AddStrings("LOCAL_SOONG_JNI_LIBS_PARTITION_"+jniLib.target.Arch.ArchType.String(),
+							jniLib.name+":"+partitionTag)
 					}
 				}
 

java/androidmk_test.go

@@ -19,6 +19,9 @@ import (
 	"testing"
 
 	"android/soong/android"
+	"android/soong/cc"
+
+	"github.com/google/blueprint/proptools"
 )
 
 func TestRequired(t *testing.T) {
@@ -252,3 +255,149 @@ func TestGetOverriddenPackages(t *testing.T) {
 		android.AssertDeepEquals(t, "overrides property", expected.overrides, actual)
 	}
 }
+
+func TestJniPartition(t *testing.T) {
+	bp := `
+		cc_library {
+			name: "libjni_system",
+			system_shared_libs: [],
+			sdk_version: "current",
+			stl: "none",
+		}
+
+		cc_library {
+			name: "libjni_system_ext",
+			system_shared_libs: [],
+			sdk_version: "current",
+			stl: "none",
+			system_ext_specific: true,
+		}
+
+		cc_library {
+			name: "libjni_odm",
+			system_shared_libs: [],
+			sdk_version: "current",
+			stl: "none",
+			device_specific: true,
+		}
+
+		cc_library {
+			name: "libjni_product",
+			system_shared_libs: [],
+			sdk_version: "current",
+			stl: "none",
+			product_specific: true,
+		}
+
+		cc_library {
+			name: "libjni_vendor",
+			system_shared_libs: [],
+			sdk_version: "current",
+			stl: "none",
+			soc_specific: true,
+		}
+
+		android_app {
+			name: "test_app_system_jni_system",
+			privileged: true,
+			platform_apis: true,
+			certificate: "platform",
+			jni_libs: ["libjni_system"],
+		}
+
+		android_app {
+			name: "test_app_system_jni_system_ext",
+			privileged: true,
+			platform_apis: true,
+			certificate: "platform",
+			jni_libs: ["libjni_system_ext"],
+		}
+
+		android_app {
+			name: "test_app_system_ext_jni_system",
+			privileged: true,
+			platform_apis: true,
+			certificate: "platform",
+			jni_libs: ["libjni_system"],
+			system_ext_specific: true
+		}
+
+		android_app {
+			name: "test_app_system_ext_jni_system_ext",
+			sdk_version: "core_platform",
+			jni_libs: ["libjni_system_ext"],
+			system_ext_specific: true
+		}
+
+		android_app {
+			name: "test_app_product_jni_product",
+			sdk_version: "core_platform",
+			jni_libs: ["libjni_product"],
+			product_specific: true
+		}
+
+		android_app {
+			name: "test_app_vendor_jni_odm",
+			sdk_version: "core_platform",
+			jni_libs: ["libjni_odm"],
+			soc_specific: true
+		}
+
+		android_app {
+			name: "test_app_odm_jni_vendor",
+			sdk_version: "core_platform",
+			jni_libs: ["libjni_vendor"],
+			device_specific: true
+		}
+		android_app {
+			name: "test_app_system_jni_multiple",
+			privileged: true,
+			platform_apis: true,
+			certificate: "platform",
+			jni_libs: ["libjni_system", "libjni_system_ext"],
+		}
+		android_app {
+			name: "test_app_vendor_jni_multiple",
+			sdk_version: "core_platform",
+			jni_libs: ["libjni_odm", "libjni_vendor"],
+			soc_specific: true
+		}
+		`
+	arch := "arm64"
+	ctx := android.GroupFixturePreparers(
+		PrepareForTestWithJavaDefaultModules,
+		cc.PrepareForTestWithCcDefaultModules,
+		android.PrepareForTestWithAndroidMk,
+		android.FixtureModifyConfig(func(config android.Config) {
+			config.TestProductVariables.DeviceArch = proptools.StringPtr(arch)
+		}),
+	).
+		RunTestWithBp(t, bp)
+	testCases := []struct {
+		name           string
+		partitionNames []string
+		partitionTags  []string
+	}{
+		{"test_app_system_jni_system", []string{"libjni_system"}, []string{""}},
+		{"test_app_system_jni_system_ext", []string{"libjni_system_ext"}, []string{"_SYSTEM_EXT"}},
+		{"test_app_system_ext_jni_system", []string{"libjni_system"}, []string{""}},
+		{"test_app_system_ext_jni_system_ext", []string{"libjni_system_ext"}, []string{"_SYSTEM_EXT"}},
+		{"test_app_product_jni_product", []string{"libjni_product"}, []string{"_PRODUCT"}},
+		{"test_app_vendor_jni_odm", []string{"libjni_odm"}, []string{"_ODM"}},
+		{"test_app_odm_jni_vendor", []string{"libjni_vendor"}, []string{"_VENDOR"}},
+		{"test_app_system_jni_multiple", []string{"libjni_system", "libjni_system_ext"}, []string{"", "_SYSTEM_EXT"}},
+		{"test_app_vendor_jni_multiple", []string{"libjni_odm", "libjni_vendor"}, []string{"_ODM", "_VENDOR"}},
+	}
+
+	for _, test := range testCases {
+		t.Run(test.name, func(t *testing.T) {
+			mod := ctx.ModuleForTests(test.name, "android_common").Module()
+			entry := android.AndroidMkEntriesForTest(t, ctx.TestContext, mod)[0]
+			for i := range test.partitionNames {
+				actual := entry.EntryMap["LOCAL_SOONG_JNI_LIBS_PARTITION_"+arch][i]
+				expected := test.partitionNames[i] + ":" + test.partitionTags[i]
+				android.AssertStringEquals(t, "Expected and actual differ", expected, actual)
+			}
+		})
+	}
+}

java/app.go

@@ -777,6 +777,7 @@ func collectAppDeps(ctx android.ModuleContext, app appDepsInterface,
 						target:         module.Target(),
 						coverageFile:   dep.CoverageOutputFile(),
 						unstrippedFile: dep.UnstrippedOutputFile(),
+						partition:      dep.Partition(),
 					})
 				} else {
 					ctx.ModuleErrorf("dependency %q missing output file", otherName)

java/java.go

@@ -418,6 +418,7 @@ type jniLib struct {
 	target         android.Target
 	coverageFile   android.OptionalPath
 	unstrippedFile android.Path
+	partition      string
 }
 
 func sdkDeps(ctx android.BottomUpMutatorContext, sdkContext android.SdkContext, d dexer) {

rust/rust.go

@@ -1664,6 +1664,10 @@ func (k kytheExtractRustSingleton) GenerateBuildActions(ctx android.SingletonCon
 	}
 }
 
+func (c *Module) Partition() string {
+	return ""
+}
+
 var Bool = proptools.Bool
 var BoolDefault = proptools.BoolDefault
 var String = proptools.String