arch specific dependencies are supported for apex

The apex module type now supports arch specific native dependencies.

apex {
    name: "myapex",
    arch: {
        arm64: {
	    native_shared_libraries: ["arm64_only_lib"],
	},
    },
}

Bug: 174639241
Test: m nothing
Change-Id: I3f00aa87c480d4127b27d33e9620c4336824e937
diff --git a/apex/apex.go b/apex/apex.go
index 2a81b9b..7475742 100644
--- a/apex/apex.go
+++ b/apex/apex.go
@@ -237,6 +237,23 @@
 	}
 }
 
+type apexArchBundleProperties struct {
+	Arch struct {
+		Arm struct {
+			ApexNativeDependencies
+		}
+		Arm64 struct {
+			ApexNativeDependencies
+		}
+		X86 struct {
+			ApexNativeDependencies
+		}
+		X86_64 struct {
+			ApexNativeDependencies
+		}
+	}
+}
+
 // These properties can be used in override_apex to override the corresponding properties in the
 // base apex.
 type overridableProperties struct {
@@ -273,6 +290,7 @@
 	// Properties
 	properties            apexBundleProperties
 	targetProperties      apexTargetBundleProperties
+	archProperties        apexArchBundleProperties
 	overridableProperties overridableProperties
 	vndkProperties        apexVndkProperties // only for apex_vndk modules
 
@@ -653,6 +671,20 @@
 			}
 		}
 
+		// Add native modules targeting a specific arch variant
+		switch target.Arch.ArchType {
+		case android.Arm:
+			depsList = append(depsList, a.archProperties.Arch.Arm.ApexNativeDependencies)
+		case android.Arm64:
+			depsList = append(depsList, a.archProperties.Arch.Arm64.ApexNativeDependencies)
+		case android.X86:
+			depsList = append(depsList, a.archProperties.Arch.X86.ApexNativeDependencies)
+		case android.X86_64:
+			depsList = append(depsList, a.archProperties.Arch.X86_64.ApexNativeDependencies)
+		default:
+			panic(fmt.Errorf("unsupported arch %v\n", ctx.Arch().ArchType))
+		}
+
 		for _, d := range depsList {
 			addDependenciesForNativeModules(ctx, d, target, imageVariation)
 		}
@@ -1910,6 +1942,7 @@
 
 	module.AddProperties(&module.properties)
 	module.AddProperties(&module.targetProperties)
+	module.AddProperties(&module.archProperties)
 	module.AddProperties(&module.overridableProperties)
 
 	android.InitAndroidMultiTargetsArchModule(module, android.HostAndDeviceSupported, android.MultilibCommon)
diff --git a/apex/apex_test.go b/apex/apex_test.go
index ab8ae16..ea2e251 100644
--- a/apex/apex_test.go
+++ b/apex/apex_test.go
@@ -3890,6 +3890,64 @@
 	ensureListContains(t, ctx.ModuleVariantsForTests("mylib2"), "android_arm64_armv8-a_shared")
 }
 
+func TestApexWithArch(t *testing.T) {
+	ctx, _ := testApex(t, `
+		apex {
+			name: "myapex",
+			key: "myapex.key",
+			arch: {
+				arm64: {
+					native_shared_libs: ["mylib.arm64"],
+				},
+				x86_64: {
+					native_shared_libs: ["mylib.x64"],
+				},
+			}
+		}
+
+		apex_key {
+			name: "myapex.key",
+			public_key: "testkey.avbpubkey",
+			private_key: "testkey.pem",
+		}
+
+		cc_library {
+			name: "mylib.arm64",
+			srcs: ["mylib.cpp"],
+			system_shared_libs: [],
+			stl: "none",
+			// TODO: remove //apex_available:platform
+			apex_available: [
+				"//apex_available:platform",
+				"myapex",
+			],
+		}
+
+		cc_library {
+			name: "mylib.x64",
+			srcs: ["mylib.cpp"],
+			system_shared_libs: [],
+			stl: "none",
+			// TODO: remove //apex_available:platform
+			apex_available: [
+				"//apex_available:platform",
+				"myapex",
+			],
+		}
+	`)
+
+	apexRule := ctx.ModuleForTests("myapex", "android_common_myapex_image").Rule("apexRule")
+	copyCmds := apexRule.Args["copy_commands"]
+
+	// Ensure that apex variant is created for the direct dep
+	ensureListContains(t, ctx.ModuleVariantsForTests("mylib.arm64"), "android_arm64_armv8-a_shared_apex10000")
+	ensureListNotContains(t, ctx.ModuleVariantsForTests("mylib.x64"), "android_arm64_armv8-a_shared_apex10000")
+
+	// Ensure that both direct and indirect deps are copied into apex
+	ensureContains(t, copyCmds, "image.apex/lib64/mylib.arm64.so")
+	ensureNotContains(t, copyCmds, "image.apex/lib64/mylib.x64.so")
+}
+
 func TestApexWithShBinary(t *testing.T) {
 	ctx, _ := testApex(t, `
 		apex {