Create variants for each image type

Create variant of image, zip or flattened according to
TARGET_FLATTEN_APEX and payload type.
If payload type is zip, only zip variant is created because flattened
apex is not supported. And if payload type is image, image and flattened
variants are created.

Bug: 139053989
Test: m -j
Change-Id: Ibde18490d23ec602c4cca97cf97db90a562e014e
diff --git a/apex/apex.go b/apex/apex.go
index 4ad2680..4cbd762 100644
--- a/apex/apex.go
+++ b/apex/apex.go
@@ -129,9 +129,11 @@
 const (
 	imageApexSuffix = ".apex"
 	zipApexSuffix   = ".zipapex"
+	flattenedSuffix = ".flattened"
 
-	imageApexType = "image"
-	zipApexType   = "zip"
+	imageApexType     = "image"
+	zipApexType       = "zip"
+	flattenedApexType = "flattened"
 
 	vndkApexNamePrefix = "com.android.vndk.v"
 )
@@ -316,13 +318,30 @@
 
 func apexFlattenedMutator(mctx android.BottomUpMutatorContext) {
 	if ab, ok := mctx.Module().(*apexBundle); ok {
-		if !mctx.Config().FlattenApex() || mctx.Config().UnbundledBuild() {
-			modules := mctx.CreateLocalVariations("", "flattened")
-			modules[0].(*apexBundle).SetFlattened(false)
-			modules[1].(*apexBundle).SetFlattened(true)
-		} else {
-			ab.SetFlattened(true)
-			ab.SetFlattenedConfigValue()
+		var variants []string
+		switch proptools.StringDefault(ab.properties.Payload_type, "image") {
+		case "image":
+			variants = append(variants, imageApexType, flattenedApexType)
+		case "zip":
+			variants = append(variants, zipApexType)
+		case "both":
+			variants = append(variants, imageApexType, zipApexType, flattenedApexType)
+		default:
+			mctx.PropertyErrorf("type", "%q is not one of \"image\" or \"zip\".", *ab.properties.Payload_type)
+			return
+		}
+
+		modules := mctx.CreateLocalVariations(variants...)
+
+		for i, v := range variants {
+			switch v {
+			case imageApexType:
+				modules[i].(*apexBundle).properties.ApexType = imageApex
+			case zipApexType:
+				modules[i].(*apexBundle).properties.ApexType = zipApex
+			case flattenedApexType:
+				modules[i].(*apexBundle).properties.ApexType = flattenedApex
+			}
 		}
 	}
 }
@@ -438,13 +457,9 @@
 	// List of APKs to package inside APEX
 	Apps []string
 
-	// To distinguish between flattened and non-flattened apex.
-	// if set true, then output files are flattened.
-	Flattened bool `blueprint:"mutated"`
-
-	// if true, it means that TARGET_FLATTEN_APEX is true and
-	// TARGET_BUILD_APPS is false
-	FlattenedConfigValue bool `blueprint:"mutated"`
+	// package format of this apex variant; could be non-flattened, flattened, or zip.
+	// imageApex, zipApex or flattened
+	ApexType apexPackaging `blueprint:"mutated"`
 
 	// List of SDKs that are used to build this APEX. A reference to an SDK should be either
 	// `name#version` or `name` which is an alias for `name#current`. If left empty, `platform#current`
@@ -501,33 +516,16 @@
 const (
 	imageApex apexPackaging = iota
 	zipApex
-	both
+	flattenedApex
 )
 
-func (a apexPackaging) image() bool {
-	switch a {
-	case imageApex, both:
-		return true
-	}
-	return false
-}
-
-func (a apexPackaging) zip() bool {
-	switch a {
-	case zipApex, both:
-		return true
-	}
-	return false
-}
-
+// The suffix for the output "file", not the module
 func (a apexPackaging) suffix() string {
 	switch a {
 	case imageApex:
 		return imageApexSuffix
 	case zipApex:
 		return zipApexSuffix
-	case both:
-		panic(fmt.Errorf("must be either zip or image"))
 	default:
 		panic(fmt.Errorf("unknown APEX type %d", a))
 	}
@@ -539,8 +537,6 @@
 		return imageApexType
 	case zipApex:
 		return zipApexType
-	case both:
-		panic(fmt.Errorf("must be either zip or image"))
 	default:
 		panic(fmt.Errorf("unknown APEX type %d", a))
 	}
@@ -590,11 +586,8 @@
 	targetProperties apexTargetBundleProperties
 	vndkProperties   apexVndkProperties
 
-	apexTypes apexPackaging
-
 	bundleModuleFile android.WritablePath
-	outputFiles      map[apexPackaging]android.WritablePath
-	flattenedOutput  android.InstallPath
+	outputFile       android.WritablePath
 	installDir       android.InstallPath
 
 	prebuiltFileToDelete string
@@ -611,8 +604,9 @@
 	// list of module names that this APEX is depending on
 	externalDeps []string
 
-	testApex bool
-	vndkApex bool
+	testApex        bool
+	vndkApex        bool
+	primaryApexType bool
 
 	// intermediate path for apex_manifest.json
 	manifestOut android.WritablePath
@@ -622,6 +616,10 @@
 	// apex package itself(for unflattened build) or apex_manifest.json(for flattened build)
 	// so that compat symlinks are always installed regardless of TARGET_FLATTEN_APEX setting.
 	compatSymlinks []string
+
+	// Suffix of module name in Android.mk
+	// ".flattened", ".apex", ".zipapex", or ""
+	suffix string
 }
 
 func addDependenciesForNativeModules(ctx android.BottomUpMutatorContext,
@@ -808,18 +806,7 @@
 func (a *apexBundle) OutputFiles(tag string) (android.Paths, error) {
 	switch tag {
 	case "":
-		if file, ok := a.outputFiles[imageApex]; ok {
-			return android.Paths{file}, nil
-		} else {
-			return nil, nil
-		}
-	case ".flattened":
-		if a.properties.Flattened {
-			flattenedApexPath := a.flattenedOutput
-			return android.Paths{flattenedApexPath}, nil
-		} else {
-			return nil, nil
-		}
+		return android.Paths{a.outputFile}, nil
 	default:
 		return nil, fmt.Errorf("unsupported module reference tag %q", tag)
 	}
@@ -876,24 +863,6 @@
 	a.properties.HideFromMake = true
 }
 
-func (a *apexBundle) SetFlattened(flattened bool) {
-	a.properties.Flattened = flattened
-}
-
-func (a *apexBundle) SetFlattenedConfigValue() {
-	a.properties.FlattenedConfigValue = true
-}
-
-// isFlattenedVariant returns true when the current module is the flattened
-// variant of an apex that has both a flattened and an unflattened variant.
-// It returns false when the current module is flattened but there is no
-// unflattened variant, which occurs when ctx.Config().FlattenedApex() returns
-// true. It can be used to avoid collisions between the install paths of the
-// flattened and unflattened variants.
-func (a *apexBundle) isFlattenedVariant() bool {
-	return a.properties.Flattened && !a.properties.FlattenedConfigValue
-}
-
 func getCopyManifestForNativeLibrary(ccMod *cc.Module, config android.Config, handleSpecialLibs bool) (fileToCopy android.Path, dirInApex string) {
 	// Decide the APEX-local directory by the multilib of the library
 	// In the future, we may query this to the module.
@@ -1012,15 +981,29 @@
 func (a *apexBundle) GenerateAndroidBuildActions(ctx android.ModuleContext) {
 	filesInfo := []apexFile{}
 
-	if a.properties.Payload_type == nil || *a.properties.Payload_type == "image" {
-		a.apexTypes = imageApex
-	} else if *a.properties.Payload_type == "zip" {
-		a.apexTypes = zipApex
-	} else if *a.properties.Payload_type == "both" {
-		a.apexTypes = both
-	} else {
-		ctx.PropertyErrorf("type", "%q is not one of \"image\", \"zip\", or \"both\".", *a.properties.Payload_type)
-		return
+	buildFlattenedAsDefault := ctx.Config().FlattenApex() && !ctx.Config().UnbundledBuild()
+	switch a.properties.ApexType {
+	case imageApex:
+		if buildFlattenedAsDefault {
+			a.suffix = imageApexSuffix
+		} else {
+			a.suffix = ""
+			a.primaryApexType = true
+		}
+	case zipApex:
+		if proptools.String(a.properties.Payload_type) == "zip" {
+			a.suffix = ""
+			a.primaryApexType = true
+		} else {
+			a.suffix = zipApexSuffix
+		}
+	case flattenedApex:
+		if buildFlattenedAsDefault {
+			a.suffix = ""
+			a.primaryApexType = true
+		} else {
+			a.suffix = flattenedSuffix
+		}
 	}
 
 	if len(a.properties.Tests) > 0 && !a.testApex {
@@ -1266,7 +1249,7 @@
 	// prepend the name of this APEX to the module names. These names will be the names of
 	// modules that will be defined if the APEX is flattened.
 	for i := range filesInfo {
-		filesInfo[i].moduleName = filesInfo[i].moduleName + "." + ctx.ModuleName()
+		filesInfo[i].moduleName = filesInfo[i].moduleName + "." + ctx.ModuleName() + a.suffix
 	}
 
 	a.installDir = android.PathForModuleInstall(ctx, "apex")
@@ -1297,27 +1280,14 @@
 		},
 	})
 
-	// Temporarily wrap the original `ctx` into a `flattenedApexContext` to have it
-	// reply true to `InstallBypassMake()` (thus making the call
-	// `android.PathForModuleInstall` below use `android.pathForInstallInMakeDir`
-	// instead of `android.PathForOutput`) to return the correct path to the flattened
-	// APEX (as its contents is installed by Make, not Soong).
-	factx := flattenedApexContext{ctx}
-	apexName := proptools.StringDefault(a.properties.Apex_name, ctx.ModuleName())
-	a.flattenedOutput = android.PathForModuleInstall(&factx, "apex", apexName)
-
-	if a.apexTypes.zip() {
-		a.buildUnflattenedApex(ctx, zipApex)
-	}
-	if a.apexTypes.image() {
-		// Build rule for unflattened APEX is created even when ctx.Config().FlattenApex()
-		// is true. This is to support referencing APEX via ":<module_name>" syntax
-		// in other modules. It is in AndroidMk where the selection of flattened
-		// or unflattened APEX is made.
-		a.buildUnflattenedApex(ctx, imageApex)
+	a.setCertificateAndPrivateKey(ctx)
+	if a.properties.ApexType == flattenedApex {
 		a.buildFlattenedApex(ctx)
+	} else {
+		a.buildUnflattenedApex(ctx)
 	}
 
+	apexName := proptools.StringDefault(a.properties.Apex_name, ctx.ModuleName())
 	a.compatSymlinks = makeCompatSymlinks(apexName, ctx)
 }
 
@@ -1343,18 +1313,7 @@
 	return android.BuildNoticeOutput(ctx, a.installDir, apexFileName, android.FirstUniquePaths(noticeFiles)).HtmlGzOutput
 }
 
-func (a *apexBundle) buildUnflattenedApex(ctx android.ModuleContext, apexType apexPackaging) {
-	cert := String(a.properties.Certificate)
-	if cert != "" && android.SrcIsModule(cert) == "" {
-		defaultDir := ctx.Config().DefaultAppCertificateDir(ctx)
-		a.container_certificate_file = defaultDir.Join(ctx, cert+".x509.pem")
-		a.container_private_key_file = defaultDir.Join(ctx, cert+".pk8")
-	} else if cert == "" {
-		pem, key := ctx.Config().DefaultAppCertificate(ctx)
-		a.container_certificate_file = pem
-		a.container_private_key_file = key
-	}
-
+func (a *apexBundle) buildUnflattenedApex(ctx android.ModuleContext) {
 	var abis []string
 	for _, target := range ctx.MultiTargets() {
 		if len(target.Arch.Abi) > 0 {
@@ -1364,6 +1323,7 @@
 
 	abis = android.FirstUniqueStrings(abis)
 
+	apexType := a.properties.ApexType
 	suffix := apexType.suffix()
 	unsignedOutputFile := android.PathForModuleOut(ctx, ctx.ModuleName()+suffix+".unsigned")
 
@@ -1424,7 +1384,7 @@
 	outHostBinDir := android.PathForOutput(ctx, "host", ctx.Config().PrebuiltOS(), "bin").String()
 	prebuiltSdkToolsBinDir := filepath.Join("prebuilts", "sdk", "tools", runtime.GOOS, "bin")
 
-	if apexType.image() {
+	if apexType == imageApex {
 		// files and dirs that will be created in APEX
 		var readOnlyPaths []string
 		var executablePaths []string // this also includes dirs
@@ -1570,11 +1530,11 @@
 		})
 	}
 
-	a.outputFiles[apexType] = android.PathForModuleOut(ctx, ctx.ModuleName()+suffix)
+	a.outputFile = android.PathForModuleOut(ctx, ctx.ModuleName()+suffix)
 	ctx.Build(pctx, android.BuildParams{
 		Rule:        java.Signapk,
 		Description: "signapk",
-		Output:      a.outputFiles[apexType],
+		Output:      a.outputFile,
 		Input:       unsignedOutputFile,
 		Implicits: []android.Path{
 			a.container_certificate_file,
@@ -1587,16 +1547,43 @@
 	})
 
 	// Install to $OUT/soong/{target,host}/.../apex
-	if a.installable() && (!ctx.Config().FlattenApex() || apexType.zip()) && !a.isFlattenedVariant() {
-		ctx.InstallFile(a.installDir, ctx.ModuleName()+suffix, a.outputFiles[apexType])
+	if a.installable() {
+		ctx.InstallFile(a.installDir, ctx.ModuleName()+suffix, a.outputFile)
 	}
+	a.buildFilesInfo(ctx)
 }
 
 func (a *apexBundle) buildFlattenedApex(ctx android.ModuleContext) {
+	// Temporarily wrap the original `ctx` into a `flattenedApexContext` to have it
+	// reply true to `InstallBypassMake()` (thus making the call
+	// `android.PathForModuleInstall` below use `android.pathForInstallInMakeDir`
+	// instead of `android.PathForOutput`) to return the correct path to the flattened
+	// APEX (as its contents is installed by Make, not Soong).
+	factx := flattenedApexContext{ctx}
+	apexName := proptools.StringDefault(a.properties.Apex_name, ctx.ModuleName())
+	a.outputFile = android.PathForModuleInstall(&factx, "apex", apexName)
+
+	a.buildFilesInfo(ctx)
+}
+
+func (a *apexBundle) setCertificateAndPrivateKey(ctx android.ModuleContext) {
+	cert := String(a.properties.Certificate)
+	if cert != "" && android.SrcIsModule(cert) == "" {
+		defaultDir := ctx.Config().DefaultAppCertificateDir(ctx)
+		a.container_certificate_file = defaultDir.Join(ctx, cert+".x509.pem")
+		a.container_private_key_file = defaultDir.Join(ctx, cert+".pk8")
+	} else if cert == "" {
+		pem, key := ctx.Config().DefaultAppCertificate(ctx)
+		a.container_certificate_file = pem
+		a.container_private_key_file = key
+	}
+}
+
+func (a *apexBundle) buildFilesInfo(ctx android.ModuleContext) {
 	if a.installable() {
 		// For flattened APEX, do nothing but make sure that apex_manifest.json and apex_pubkey are also copied along
 		// with other ordinary files.
-		a.filesInfo = append(a.filesInfo, apexFile{a.manifestOut, "apex_manifest.json." + ctx.ModuleName(), ".", etc, nil, nil})
+		a.filesInfo = append(a.filesInfo, apexFile{a.manifestOut, "apex_manifest.json." + ctx.ModuleName() + a.suffix, ".", etc, nil, nil})
 
 		// rename to apex_pubkey
 		copiedPubkey := android.PathForModuleOut(ctx, "apex_pubkey")
@@ -1605,9 +1592,9 @@
 			Input:  a.public_key_file,
 			Output: copiedPubkey,
 		})
-		a.filesInfo = append(a.filesInfo, apexFile{copiedPubkey, "apex_pubkey." + ctx.ModuleName(), ".", etc, nil, nil})
+		a.filesInfo = append(a.filesInfo, apexFile{copiedPubkey, "apex_pubkey." + ctx.ModuleName() + a.suffix, ".", etc, nil, nil})
 
-		if ctx.Config().FlattenApex() {
+		if a.properties.ApexType == flattenedApex {
 			apexName := proptools.StringDefault(a.properties.Apex_name, ctx.ModuleName())
 			for _, fi := range a.filesInfo {
 				dir := filepath.Join("apex", apexName, fi.installDir)
@@ -1627,12 +1614,7 @@
 		}
 	}
 	writers := []android.AndroidMkData{}
-	if a.apexTypes.image() {
-		writers = append(writers, a.androidMkForType(imageApex))
-	}
-	if a.apexTypes.zip() {
-		writers = append(writers, a.androidMkForType(zipApex))
-	}
+	writers = append(writers, a.androidMkForType())
 	return android.AndroidMkData{
 		Custom: func(w io.Writer, name, prefix, moduleDir string, data android.AndroidMkData) {
 			for _, data := range writers {
@@ -1641,36 +1623,35 @@
 		}}
 }
 
-func (a *apexBundle) androidMkForFiles(w io.Writer, apexName, moduleDir string, apexType apexPackaging) []string {
+func (a *apexBundle) androidMkForFiles(w io.Writer, apexName, moduleDir string) []string {
 	moduleNames := []string{}
+	apexType := a.properties.ApexType
+	// To avoid creating duplicate build rules, run this function only when primaryApexType is true
+	// to install symbol files in $(PRODUCT_OUT}/apex.
+	// And if apexType is flattened, run this function to install files in $(PRODUCT_OUT}/system/apex.
+	if !a.primaryApexType && apexType != flattenedApex {
+		return moduleNames
+	}
 
 	for _, fi := range a.filesInfo {
 		if cc, ok := fi.module.(*cc.Module); ok && cc.Properties.HideFromMake {
 			continue
 		}
-		if a.properties.Flattened && !apexType.image() {
-			continue
-		}
-
-		var suffix string
-		if a.isFlattenedVariant() {
-			suffix = ".flattened"
-		}
 
 		if !android.InList(fi.moduleName, moduleNames) {
-			moduleNames = append(moduleNames, fi.moduleName+suffix)
+			moduleNames = append(moduleNames, fi.moduleName)
 		}
 
 		fmt.Fprintln(w, "\ninclude $(CLEAR_VARS)")
 		fmt.Fprintln(w, "LOCAL_PATH :=", moduleDir)
-		fmt.Fprintln(w, "LOCAL_MODULE :=", fi.moduleName+suffix)
+		fmt.Fprintln(w, "LOCAL_MODULE :=", fi.moduleName)
 		// /apex/<apex_name>/{lib|framework|...}
 		pathWhenActivated := filepath.Join("$(PRODUCT_OUT)", "apex", apexName, fi.installDir)
-		if a.properties.Flattened && apexType.image() {
+		if apexType == flattenedApex {
 			// /system/apex/<name>/{lib|framework|...}
 			fmt.Fprintln(w, "LOCAL_MODULE_PATH :=", filepath.Join(a.installDir.ToMakePath().String(),
 				apexName, fi.installDir))
-			if !a.isFlattenedVariant() {
+			if a.primaryApexType {
 				fmt.Fprintln(w, "LOCAL_SOONG_SYMBOL_PATH :=", pathWhenActivated)
 			}
 			if len(fi.symlinks) > 0 {
@@ -1739,7 +1720,7 @@
 		} else {
 			fmt.Fprintln(w, "LOCAL_MODULE_STEM :=", fi.builtFile.Base())
 			// For flattened apexes, compat symlinks are attached to apex_manifest.json which is guaranteed for every apex
-			if !a.isFlattenedVariant() && fi.builtFile.Base() == "apex_manifest.json" && len(a.compatSymlinks) > 0 {
+			if a.primaryApexType && fi.builtFile.Base() == "apex_manifest.json" && len(a.compatSymlinks) > 0 {
 				fmt.Fprintln(w, "LOCAL_POST_INSTALL_CMD :=", strings.Join(a.compatSymlinks, " && "))
 			}
 			fmt.Fprintln(w, "include $(BUILD_PREBUILT)")
@@ -1748,41 +1729,33 @@
 	return moduleNames
 }
 
-func (a *apexBundle) androidMkForType(apexType apexPackaging) android.AndroidMkData {
+func (a *apexBundle) androidMkForType() android.AndroidMkData {
 	return android.AndroidMkData{
 		Custom: func(w io.Writer, name, prefix, moduleDir string, data android.AndroidMkData) {
 			moduleNames := []string{}
+			apexType := a.properties.ApexType
 			if a.installable() {
 				apexName := proptools.StringDefault(a.properties.Apex_name, name)
-				moduleNames = a.androidMkForFiles(w, apexName, moduleDir, apexType)
+				moduleNames = a.androidMkForFiles(w, apexName, moduleDir)
 			}
 
-			if a.isFlattenedVariant() {
-				name = name + ".flattened"
-			}
-
-			if a.properties.Flattened && apexType.image() {
+			if apexType == flattenedApex {
 				// Only image APEXes can be flattened.
 				fmt.Fprintln(w, "\ninclude $(CLEAR_VARS)")
 				fmt.Fprintln(w, "LOCAL_PATH :=", moduleDir)
-				fmt.Fprintln(w, "LOCAL_MODULE :=", name)
+				fmt.Fprintln(w, "LOCAL_MODULE :=", name+a.suffix)
 				if len(moduleNames) > 0 {
 					fmt.Fprintln(w, "LOCAL_REQUIRED_MODULES :=", strings.Join(moduleNames, " "))
 				}
 				fmt.Fprintln(w, "include $(BUILD_PHONY_PACKAGE)")
-				fmt.Fprintln(w, "$(LOCAL_INSTALLED_MODULE): .KATI_IMPLICIT_OUTPUTS :=", a.flattenedOutput.String())
+				fmt.Fprintln(w, "$(LOCAL_INSTALLED_MODULE): .KATI_IMPLICIT_OUTPUTS :=", a.outputFile.String())
 
-			} else if !a.isFlattenedVariant() {
-				// zip-apex is the less common type so have the name refer to the image-apex
-				// only and use {name}.zip if you want the zip-apex
-				if apexType == zipApex && a.apexTypes == both {
-					name = name + ".zip"
-				}
+			} else {
 				fmt.Fprintln(w, "\ninclude $(CLEAR_VARS)")
 				fmt.Fprintln(w, "LOCAL_PATH :=", moduleDir)
-				fmt.Fprintln(w, "LOCAL_MODULE :=", name)
+				fmt.Fprintln(w, "LOCAL_MODULE :=", name+a.suffix)
 				fmt.Fprintln(w, "LOCAL_MODULE_CLASS := ETC") // do we need a new class?
-				fmt.Fprintln(w, "LOCAL_PREBUILT_MODULE_FILE :=", a.outputFiles[apexType].String())
+				fmt.Fprintln(w, "LOCAL_PREBUILT_MODULE_FILE :=", a.outputFile.String())
 				fmt.Fprintln(w, "LOCAL_MODULE_PATH :=", a.installDir.ToMakePath().String())
 				fmt.Fprintln(w, "LOCAL_MODULE_STEM :=", name+apexType.suffix())
 				fmt.Fprintln(w, "LOCAL_UNINSTALLABLE_MODULE :=", !a.installable())
@@ -1812,9 +1785,7 @@
 }
 
 func newApexBundle() *apexBundle {
-	module := &apexBundle{
-		outputFiles: map[apexPackaging]android.WritablePath{},
-	}
+	module := &apexBundle{}
 	module.AddProperties(&module.properties)
 	module.AddProperties(&module.targetProperties)
 	module.Prefer32(func(ctx android.BaseModuleContext, base *android.ModuleBase, class android.OsClass) bool {