3 files changed, 50 insertions, 4 deletions
diff --git a/android/image.go b/android/image.go
index 012267a85..e8b6352fd 100644
--- a/android/image.go
+++ b/android/image.go
@@ -98,16 +98,47 @@ const (
 	DebugRamdiskVariation string = "debug_ramdisk"
 )
 
+type imageInterfaceContextAdapter struct {
+	IncomingTransitionContext
+	kind moduleKind
+}
+
+var _ ImageInterfaceContext = (*imageInterfaceContextAdapter)(nil)
+
+func (e *imageInterfaceContextAdapter) Platform() bool {
+	return e.kind == platformModule
+}
+
+func (e *imageInterfaceContextAdapter) DeviceSpecific() bool {
+	return e.kind == deviceSpecificModule
+}
+
+func (e *imageInterfaceContextAdapter) SocSpecific() bool {
+	return e.kind == socSpecificModule
+}
+
+func (e *imageInterfaceContextAdapter) ProductSpecific() bool {
+	return e.kind == productSpecificModule
+}
+
+func (e *imageInterfaceContextAdapter) SystemExtSpecific() bool {
+	return e.kind == systemExtSpecificModule
+}
+
+func imageMutatorBeginMutator(ctx BottomUpMutatorContext) {
+	if m, ok := ctx.Module().(ImageInterface); ok && ctx.Os() == Android {
+		m.ImageMutatorBegin(ctx)
+	}
+}
+
 // imageTransitionMutator creates variants for modules that implement the ImageInterface that
 // allow them to build differently for each partition (recovery, core, vendor, etc.).
 type imageTransitionMutator struct{}
 
-func (imageTransitionMutator) Split(ctx BaseModuleContext) []string {
+func getImageVariations(ctx ImageInterfaceContext) []string {
 	var variations []string
 
 	if m, ok := ctx.Module().(ImageInterface); ctx.Os() == Android && ok {
-		m.ImageMutatorBegin(ctx)
-
 		if m.CoreVariantNeeded(ctx) {
 			variations = append(variations, CoreVariation)
 		}
@@ -141,6 +172,10 @@ func (imageTransitionMutator) Split(ctx BaseModuleContext) []string {
 	return variations
 }
 
+func (imageTransitionMutator) Split(ctx BaseModuleContext) []string {
+	return getImageVariations(ctx)
+}
+
 func (imageTransitionMutator) OutgoingTransition(ctx OutgoingTransitionContext, sourceVariation string) string {
 	return sourceVariation
 }
@@ -149,6 +184,16 @@ func (imageTransitionMutator) IncomingTransition(ctx IncomingTransitionContext,
 	if _, ok := ctx.Module().(ImageInterface); ctx.Os() != Android || !ok {
 		return CoreVariation
 	}
+	variations := getImageVariations(&imageInterfaceContextAdapter{
+		IncomingTransitionContext: ctx,
+		kind:                      determineModuleKind(ctx.Module().base(), ctx),
+	})
+	// If there's only 1 possible variation, use that. This is a holdover from when blueprint,
+	// when adding dependencies, would use the only variant of a module regardless of its variations
+	// if only 1 variant existed.
+	if len(variations) == 1 {
+		return variations[0]
+	}
 	return incomingVariation
 }
 
diff --git a/android/module.go b/android/module.go
index bc7f5e63a..6ef5c6ae5 100644
--- a/android/module.go
+++ b/android/module.go
@@ -1667,7 +1667,7 @@ func (m *ModuleBase) generateModuleTarget(ctx *moduleContext) {
 	}
 }
 
-func determineModuleKind(m *ModuleBase, ctx blueprint.EarlyModuleContext) moduleKind {
+func determineModuleKind(m *ModuleBase, ctx ModuleErrorContext) moduleKind {
 	var socSpecific = Bool(m.commonProperties.Vendor) || Bool(m.commonProperties.Proprietary) || Bool(m.commonProperties.Soc_specific)
 	var deviceSpecific = Bool(m.commonProperties.Device_specific)
 	var productSpecific = Bool(m.commonProperties.Product_specific)
diff --git a/android/mutator.go b/android/mutator.go
index b73679641..4ddc606b1 100644
--- a/android/mutator.go
+++ b/android/mutator.go
@@ -151,6 +151,7 @@ var preArch = []RegisterMutatorFunc{
 
 func registerArchMutator(ctx RegisterMutatorsContext) {
 	ctx.Transition("os", &osTransitionMutator{})
+	ctx.BottomUp("image_begin", imageMutatorBeginMutator)
 	ctx.Transition("image", &imageTransitionMutator{})
 	ctx.Transition("arch", &archTransitionMutator{})
 }