Prebuilt replacing source should not change partition

This reveals unintended mistake (setting a wrong target partition) at
build-time instead of runtime, which is much harder to debug.

Bug: 280368661
Test: m nothing (soong test)
Change-Id: Ic5e5e97ba918e24f7a59aceb405c2b105e28cccc

android/override_module.go

@@ -46,8 +46,8 @@ type OverrideModule interface {
 
 	// Internal funcs to handle interoperability between override modules and prebuilts.
 	// i.e. cases where an overriding module, too, is overridden by a prebuilt module.
-	setOverriddenByPrebuilt(overridden bool)
-	getOverriddenByPrebuilt() bool
+	setOverriddenByPrebuilt(prebuilt Module)
+	getOverriddenByPrebuilt() Module
 
 	// Directory containing the Blueprint definition of the overriding module
 	setModuleDir(string)
@@ -60,7 +60,7 @@ type OverrideModuleBase struct {
 
 	overridingProperties []interface{}
 
-	overriddenByPrebuilt bool
+	overriddenByPrebuilt Module
 
 	moduleDir string
 }
@@ -96,11 +96,11 @@ func (o *OverrideModuleBase) GetOverriddenModuleName() string {
 	return proptools.String(o.moduleProperties.Base)
 }
 
-func (o *OverrideModuleBase) setOverriddenByPrebuilt(overridden bool) {
-	o.overriddenByPrebuilt = overridden
+func (o *OverrideModuleBase) setOverriddenByPrebuilt(prebuilt Module) {
+	o.overriddenByPrebuilt = prebuilt
 }
 
-func (o *OverrideModuleBase) getOverriddenByPrebuilt() bool {
+func (o *OverrideModuleBase) getOverriddenByPrebuilt() Module {
 	return o.overriddenByPrebuilt
 }
 
@@ -281,7 +281,7 @@ func overrideModuleDepsMutator(ctx BottomUpMutatorContext) {
 				panic("PrebuiltDepTag leads to a non-prebuilt module " + dep.Name())
 			}
 			if prebuilt.UsePrebuilt() {
-				module.setOverriddenByPrebuilt(true)
+				module.setOverriddenByPrebuilt(dep)
 				return
 			}
 		})
@@ -314,8 +314,10 @@ func performOverrideMutator(ctx BottomUpMutatorContext) {
 		ctx.AliasVariation(variants[0])
 		for i, o := range overrides {
 			mods[i+1].(OverridableModule).override(ctx, mods[i+1], o)
-			if o.getOverriddenByPrebuilt() {
+			if prebuilt := o.getOverriddenByPrebuilt(); prebuilt != nil {
 				// The overriding module itself, too, is overridden by a prebuilt.
+				// Perform the same check for replacement
+				checkInvariantsForSourceAndPrebuilt(ctx, mods[i+1], prebuilt)
 				// Copy the flag and hide it in make
 				mods[i+1].ReplacedByPrebuilt()
 			}

android/prebuilt.go

@@ -419,6 +419,20 @@ func PrebuiltSourceDepsMutator(ctx BottomUpMutatorContext) {
 	}
 }
 
+// checkInvariantsForSourceAndPrebuilt checks if invariants are kept when replacing
+// source with prebuilt. Note that the current module for the context is the source module.
+func checkInvariantsForSourceAndPrebuilt(ctx BaseModuleContext, s, p Module) {
+	if _, ok := s.(OverrideModule); ok {
+		// skip the check when the source module is `override_X` because it's only a placeholder
+		// for the actual source module. The check will be invoked for the actual module.
+		return
+	}
+	if sourcePartition, prebuiltPartition := s.PartitionTag(ctx.DeviceConfig()), p.PartitionTag(ctx.DeviceConfig()); sourcePartition != prebuiltPartition {
+		ctx.OtherModuleErrorf(p, "partition is different: %s(%s) != %s(%s)",
+			sourcePartition, ctx.ModuleName(), prebuiltPartition, ctx.OtherModuleName(p))
+	}
+}
+
 // PrebuiltSelectModuleMutator marks prebuilts that are used, either overriding source modules or
 // because the source module doesn't exist.  It also disables installing overridden source modules.
 func PrebuiltSelectModuleMutator(ctx TopDownMutatorContext) {
@@ -434,6 +448,8 @@ func PrebuiltSelectModuleMutator(ctx TopDownMutatorContext) {
 		ctx.VisitDirectDepsWithTag(PrebuiltDepTag, func(prebuiltModule Module) {
 			p := GetEmbeddedPrebuilt(prebuiltModule)
 			if p.usePrebuilt(ctx, s, prebuiltModule) {
+				checkInvariantsForSourceAndPrebuilt(ctx, s, prebuiltModule)
+
 				p.properties.UsePrebuilt = true
 				s.ReplacedByPrebuilt()
 			}

android/prebuilt_test.go

@@ -497,6 +497,52 @@ func TestPrebuilts(t *testing.T) {
 	}
 }
 
+func testPrebuiltError(t *testing.T, expectedError, bp string) {
+	t.Helper()
+	fs := MockFS{
+		"prebuilt_file": nil,
+	}
+	GroupFixturePreparers(
+		PrepareForTestWithArchMutator,
+		PrepareForTestWithPrebuilts,
+		PrepareForTestWithOverrides,
+		fs.AddToFixture(),
+		FixtureRegisterWithContext(registerTestPrebuiltModules),
+	).
+		ExtendWithErrorHandler(FixtureExpectsAtLeastOneErrorMatchingPattern(expectedError)).
+		RunTestWithBp(t, bp)
+}
+
+func TestPrebuiltShouldNotChangePartition(t *testing.T) {
+	testPrebuiltError(t, `partition is different`, `
+		source {
+			name: "foo",
+			vendor: true,
+		}
+		prebuilt {
+			name: "foo",
+			prefer: true,
+			srcs: ["prebuilt_file"],
+		}`)
+}
+
+func TestPrebuiltShouldNotChangePartition_WithOverride(t *testing.T) {
+	testPrebuiltError(t, `partition is different`, `
+		source {
+			name: "foo",
+			vendor: true,
+		}
+		override_source {
+			name: "bar",
+			base: "foo",
+		}
+		prebuilt {
+			name: "bar",
+			prefer: true,
+			srcs: ["prebuilt_file"],
+		}`)
+}
+
 func registerTestPrebuiltBuildComponents(ctx RegistrationContext) {
 	registerTestPrebuiltModules(ctx)