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
diff --git a/apex/apex.go b/apex/apex.go
index 685a774..e07fae0 100644
--- a/apex/apex.go
+++ b/apex/apex.go
@@ -1300,7 +1300,23 @@
 
 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 @@
 		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 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
 }