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 {