Arch-specific source for prebuilt_apex

Arch-specific source can be specified for prebuilt_apex as follows.

arch: {
    arm64: {
        src: "myapex-arm64.apex",
    },
},

A note on the implementation. The Src property was not tagged as
`android:"arch_variant"` as usual. Instead, multiple Src properties are
explicitly declared like

struct Arch {
    struct Arm {
        struct Src *string
    }
    ...
}

Corresponding Src property is manually selected according to the
MultiTargets()[0].

This is because prebuilt_apex is mutated only for android_common, in
order to have the same arch variant with the apex module type.
Therefore, we can't rely on the arch_variant tag.

Bug: 127789981
Test: m (apex_test amended)

Change-Id: I77dbe626171d8975f549bdb4af3c487232cf05f7

2 files changed, 61 insertions, 6 deletions

diff --git a/apex/apex.go b/apex/apex.go
index 685a7746..e07fae06 100644
--- a/apex/apex.go
+++ b/apex/apex.go
@@ -1300,7 +1300,23 @@ type Prebuilt struct {
 
 type PrebuiltProperties struct {
 	// the path to the prebuilt .apex file to import.
-	Src string `android:"arch_variant"`
+	Source string `blueprint:"mutated"`
+
+	Src  *string
+	Arch struct {
+		Arm struct {
+			Src *string
+		}
+		Arm64 struct {
+			Src *string
+		}
+		X86 struct {
+			Src *string
+		}
+		X86_64 struct {
+			Src *string
+		}
+	}
 
 	// the name of the apex_key module that contains the matching public key to be installed.
 	Key *string
@@ -1312,11 +1328,37 @@ func (p *Prebuilt) DepsMutator(ctx android.BottomUpMutatorContext) {
 		return
 	}
 	ctx.AddDependency(ctx.Module(), keyTag, *p.properties.Key)
+
+	// This is called before prebuilt_select and prebuilt_postdeps mutators
+	// The mutators requires that src to be set correctly for each arch so that
+	// arch variants are disabled when src is not provided for the arch.
+	if len(ctx.MultiTargets()) != 1 {
+		ctx.ModuleErrorf("compile_multilib shouldn't be \"both\" for prebuilt_apex")
+		return
+	}
+	var src string
+	switch ctx.MultiTargets()[0].Arch.ArchType {
+	case android.Arm:
+		src = String(p.properties.Arch.Arm.Src)
+	case android.Arm64:
+		src = String(p.properties.Arch.Arm64.Src)
+	case android.X86:
+		src = String(p.properties.Arch.X86.Src)
+	case android.X86_64:
+		src = String(p.properties.Arch.X86_64.Src)
+	default:
+		ctx.ModuleErrorf("prebuilt_apex does not support %q", ctx.MultiTargets()[0].Arch.String())
+		return
+	}
+	if src == "" {
+		src = String(p.properties.Src)
+	}
+	p.properties.Source = src
 }
 
 func (p *Prebuilt) GenerateAndroidBuildActions(ctx android.ModuleContext) {
 	// TODO(jungjw): Check the key validity.
-	p.inputApex = p.prebuilt.SingleSourcePath(ctx)
+	p.inputApex = p.Prebuilt().SingleSourcePath(ctx)
 	p.installDir = android.PathForModuleInstall(ctx, "apex")
 	ctx.InstallFile(p.installDir, ctx.ModuleName()+imageApexSuffix, p.inputApex)
 }
@@ -1348,7 +1390,7 @@ func (p *Prebuilt) AndroidMk() android.AndroidMkData {
 func PrebuiltFactory() android.Module {
 	module := &Prebuilt{}
 	module.AddProperties(&module.properties)
-	android.InitSingleSourcePrebuiltModule(module, &module.properties.Src)
-	android.InitAndroidArchModule(module, android.DeviceSupported, android.MultilibCommon)
+	android.InitSingleSourcePrebuiltModule(module, &module.properties.Source)
+	android.InitAndroidMultiTargetsArchModule(module, android.DeviceSupported, android.MultilibCommon)
 	return module
 }
diff --git a/apex/apex_test.go b/apex/apex_test.go
index 2d9cca6f..1e8d5b44 100644
--- a/apex/apex_test.go
+++ b/apex/apex_test.go
@@ -171,7 +171,8 @@ func testApex(t *testing.T, bp string) *android.TestContext {
 		"custom_notice":                        nil,
 		"testkey2.avbpubkey":                   nil,
 		"testkey2.pem":                         nil,
-		"myapex.apex":                          nil,
+		"myapex-arm64.apex":                    nil,
+		"myapex-arm.apex":                      nil,
 	})
 	_, errs := ctx.ParseFileList(".", []string{"Android.bp"})
 	android.FailIfErrored(t, errs)
@@ -1243,7 +1244,14 @@ func TestPrebuilt(t *testing.T) {
 	ctx := testApex(t, `
 		prebuilt_apex {
 			name: "myapex",
-			src: "myapex.apex",
+			arch: {
+				arm64: {
+					src: "myapex-arm64.apex",
+				},
+				arm: {
+					src: "myapex-arm.apex",
+				},
+			},
 			key: "myapex.key"
 		}
 
@@ -1257,6 +1265,11 @@ func TestPrebuilt(t *testing.T) {
 
 	prebuilt := ctx.ModuleForTests("myapex", "android_common").Module().(*Prebuilt)
 
+	expectedInput := "myapex-arm64.apex"
+	if prebuilt.inputApex.String() != expectedInput {
+		t.Errorf("inputApex invalid. expected: %q, actual: %q", expectedInput, prebuilt.inputApex.String())
+	}
+
 	// Check if the key module is added as a required module.
 	buf := &bytes.Buffer{}
 	prebuilt.AndroidMk().Extra[0](buf, nil)