Add arch variant support to android_app_import.

Bug: 128610294
Fixes: 138792623
Test: app_test.go
Change-Id: I47c80ec283ce58a0ce9b7d0af40844bd73e9d3f1

4 files changed, 120 insertions, 30 deletions

diff --git a/android/arch.go b/android/arch.go
index 957a659c..6e8c4bdb 100644
--- a/android/arch.go
+++ b/android/arch.go
@@ -557,6 +557,10 @@ func newArch(name, multilib string) ArchType {
 	return archType
 }
 
+func ArchTypeList() []ArchType {
+	return append([]ArchType(nil), archTypeList...)
+}
+
 func (a ArchType) String() string {
 	return a.Name
 }
diff --git a/java/app.go b/java/app.go
index bcbf6f31..cea45572 100644
--- a/java/app.go
+++ b/java/app.go
@@ -38,6 +38,8 @@ func init() {
 	android.RegisterModuleType("android_app_certificate", AndroidAppCertificateFactory)
 	android.RegisterModuleType("override_android_app", OverrideAndroidAppModuleFactory)
 	android.RegisterModuleType("android_app_import", AndroidAppImportFactory)
+
+	initAndroidAppImportVariantGroupTypes()
 }
 
 // AndroidManifest.xml merging
@@ -685,8 +687,9 @@ type AndroidAppImport struct {
 	android.DefaultableModuleBase
 	prebuilt android.Prebuilt
 
-	properties  AndroidAppImportProperties
-	dpiVariants interface{}
+	properties   AndroidAppImportProperties
+	dpiVariants  interface{}
+	archVariants interface{}
 
 	outputFile  android.Path
 	certificate *Certificate
@@ -724,8 +727,8 @@ type AndroidAppImportProperties struct {
 	Filename *string
 }
 
-// Chooses a source APK path to use based on the module and product specs.
-func (a *AndroidAppImport) updateSrcApkPath(ctx android.LoadHookContext) {
+// Updates properties with variant-specific values.
+func (a *AndroidAppImport) processVariants(ctx android.LoadHookContext) {
 	config := ctx.Config()
 
 	dpiProps := reflect.ValueOf(a.dpiVariants).Elem().FieldByName("Dpi_variants")
@@ -733,24 +736,22 @@ func (a *AndroidAppImport) updateSrcApkPath(ctx android.LoadHookContext) {
 	// overwrites everything else.
 	// TODO(jungjw): Can we optimize this by making it priority order?
 	for i := len(config.ProductAAPTPrebuiltDPI()) - 1; i >= 0; i-- {
-		dpi := config.ProductAAPTPrebuiltDPI()[i]
-		if inList(dpi, supportedDpis) {
-			MergePropertiesFromVariant(ctx, &a.properties, dpiProps, dpi, "dpi_variants")
-		}
+		MergePropertiesFromVariant(ctx, &a.properties, dpiProps, config.ProductAAPTPrebuiltDPI()[i])
 	}
 	if config.ProductAAPTPreferredConfig() != "" {
-		dpi := config.ProductAAPTPreferredConfig()
-		if inList(dpi, supportedDpis) {
-			MergePropertiesFromVariant(ctx, &a.properties, dpiProps, dpi, "dpi_variants")
-		}
+		MergePropertiesFromVariant(ctx, &a.properties, dpiProps, config.ProductAAPTPreferredConfig())
 	}
+
+	archProps := reflect.ValueOf(a.archVariants).Elem().FieldByName("Arch")
+	archType := ctx.Config().Targets[android.Android][0].Arch.ArchType
+	MergePropertiesFromVariant(ctx, &a.properties, archProps, archType.Name)
 }
 
 func MergePropertiesFromVariant(ctx android.BaseModuleContext,
-	dst interface{}, variantGroup reflect.Value, variant, variantGroupPath string) {
+	dst interface{}, variantGroup reflect.Value, variant string) {
 	src := variantGroup.FieldByName(proptools.FieldNameForProperty(variant))
 	if !src.IsValid() {
-		ctx.ModuleErrorf("field %q does not exist", variantGroupPath+"."+variant)
+		return
 	}
 
 	err := proptools.ExtendMatchingProperties([]interface{}{dst}, src.Interface(), nil, proptools.OrderAppend)
@@ -876,26 +877,46 @@ func (a *AndroidAppImport) Name() string {
 	return a.prebuilt.Name(a.ModuleBase.Name())
 }
 
-// Populates dpi_variants property and its fields at creation time.
-func (a *AndroidAppImport) addDpiVariants() {
-	// TODO(jungjw): Do we want to do some filtering here?
-	props := reflect.ValueOf(&a.properties).Type()
+var dpiVariantGroupType reflect.Type
+var archVariantGroupType reflect.Type
+
+func initAndroidAppImportVariantGroupTypes() {
+	dpiVariantGroupType = createVariantGroupType(supportedDpis, "Dpi_variants")
+
+	archNames := make([]string, len(android.ArchTypeList()))
+	for i, archType := range android.ArchTypeList() {
+		archNames[i] = archType.Name
+	}
+	archVariantGroupType = createVariantGroupType(archNames, "Arch")
+}
 
-	dpiFields := make([]reflect.StructField, len(supportedDpis))
-	for i, dpi := range supportedDpis {
-		dpiFields[i] = reflect.StructField{
-			Name: proptools.FieldNameForProperty(dpi),
+// Populates all variant struct properties at creation time.
+func (a *AndroidAppImport) populateAllVariantStructs() {
+	a.dpiVariants = reflect.New(dpiVariantGroupType).Interface()
+	a.AddProperties(a.dpiVariants)
+
+	a.archVariants = reflect.New(archVariantGroupType).Interface()
+	a.AddProperties(a.archVariants)
+}
+
+func createVariantGroupType(variants []string, variantGroupName string) reflect.Type {
+	props := reflect.TypeOf((*AndroidAppImportProperties)(nil))
+
+	variantFields := make([]reflect.StructField, len(variants))
+	for i, variant := range variants {
+		variantFields[i] = reflect.StructField{
+			Name: proptools.FieldNameForProperty(variant),
 			Type: props,
 		}
 	}
-	dpiStruct := reflect.StructOf(dpiFields)
-	a.dpiVariants = reflect.New(reflect.StructOf([]reflect.StructField{
+
+	variantGroupStruct := reflect.StructOf(variantFields)
+	return reflect.StructOf([]reflect.StructField{
 		{
-			Name: "Dpi_variants",
-			Type: dpiStruct,
+			Name: variantGroupName,
+			Type: variantGroupStruct,
 		},
-	})).Interface()
-	a.AddProperties(a.dpiVariants)
+	})
 }
 
 // android_app_import imports a prebuilt apk with additional processing specified in the module.
@@ -918,9 +939,9 @@ func AndroidAppImportFactory() android.Module {
 	module := &AndroidAppImport{}
 	module.AddProperties(&module.properties)
 	module.AddProperties(&module.dexpreoptProperties)
-	module.addDpiVariants()
+	module.populateAllVariantStructs()
 	android.AddLoadHook(module, func(ctx android.LoadHookContext) {
-		module.updateSrcApkPath(ctx)
+		module.processVariants(ctx)
 	})
 
 	InitJavaModule(module, android.DeviceSupported)
diff --git a/java/app_test.go b/java/app_test.go
index 8b254ad0..a5c4fbc7 100644
--- a/java/app_test.go
+++ b/java/app_test.go
@@ -1362,3 +1362,66 @@ func TestAndroidAppImport_Filename(t *testing.T) {
 		}
 	}
 }
+
+func TestAndroidAppImport_ArchVariants(t *testing.T) {
+	// The test config's target arch is ARM64.
+	testCases := []struct {
+		name     string
+		bp       string
+		expected string
+	}{
+		{
+			name: "matching arch",
+			bp: `
+				android_app_import {
+					name: "foo",
+					apk: "prebuilts/apk/app.apk",
+					arch: {
+						arm64: {
+							apk: "prebuilts/apk/app_arm64.apk",
+						},
+					},
+					certificate: "PRESIGNED",
+					dex_preopt: {
+						enabled: true,
+					},
+				}
+			`,
+			expected: "prebuilts/apk/app_arm64.apk",
+		},
+		{
+			name: "no matching arch",
+			bp: `
+				android_app_import {
+					name: "foo",
+					apk: "prebuilts/apk/app.apk",
+					arch: {
+						arm: {
+							apk: "prebuilts/apk/app_arm.apk",
+						},
+					},
+					certificate: "PRESIGNED",
+					dex_preopt: {
+						enabled: true,
+					},
+				}
+			`,
+			expected: "prebuilts/apk/app.apk",
+		},
+	}
+
+	jniRuleRe := regexp.MustCompile("^if \\(zipinfo (\\S+)")
+	for _, test := range testCases {
+		ctx := testJava(t, test.bp)
+
+		variant := ctx.ModuleForTests("foo", "android_common")
+		jniRuleCommand := variant.Output("jnis-uncompressed/foo.apk").RuleParams.Command
+		matches := jniRuleRe.FindStringSubmatch(jniRuleCommand)
+		if len(matches) != 2 {
+			t.Errorf("failed to extract the src apk path from %q", jniRuleCommand)
+		}
+		if test.expected != matches[1] {
+			t.Errorf("wrong src apk, expected: %q got: %q", test.expected, matches[1])
+		}
+	}
+}
diff --git a/java/java_test.go b/java/java_test.go
index 96ef7e00..bafb466c 100644
--- a/java/java_test.go
+++ b/java/java_test.go
@@ -169,6 +169,8 @@ func testContext(config android.Config, bp string,
 		"prebuilts/sdk/Android.bp":                    []byte(`prebuilt_apis { name: "sdk", api_dirs: ["14", "28", "current"],}`),
 
 		"prebuilts/apk/app.apk":        nil,
+		"prebuilts/apk/app_arm.apk":    nil,
+		"prebuilts/apk/app_arm64.apk":  nil,
 		"prebuilts/apk/app_xhdpi.apk":  nil,
 		"prebuilts/apk/app_xxhdpi.apk": nil,