Rewrite code to copy images to PRODUCT_OUT

Move the logic to android_device. The reasons for this are:
 - As we start to check in android_device and filesystem modules,
   only one set should do the copying. It's easy enough to mark the
   android_device as the one, but harder for the filesystem modules
   given that they may be reused between different devices and now
   that top-down mutators are no longer supported.
 - It's easier to manage this code in one place, especially since it's
   temporary and should be removed eventually.

Having this higher confidence that there's only 1 copy rule set in
the build makes it less of a problem to use actual ninja rules for
the images.

Bug: 376727180
Test: m --soong-only && ls out/target/product/vsoc_x86_64/
Change-Id: If20ea2c55053cc962de9f9770db002edbc194835

9 files changed, 270 insertions, 97 deletions

diff --git a/android/android_info.go b/android/android_info.go
index 225c8f0ae..561fa4c0c 100644
--- a/android/android_info.go
+++ b/android/android_info.go
@@ -82,18 +82,6 @@ func (p *androidInfoModule) GenerateAndroidBuildActions(ctx ModuleContext) {
 
 	ctx.SetOutputFiles(Paths{androidInfoProp}, "")
 	ctx.SetOutputFiles(Paths{androidInfoTxt}, ".txt")
-
-	builder := NewRuleBuilder(pctx, ctx)
-	builder.Command().Text("touch").Output(timestamp)
-	if !ctx.Config().KatiEnabled() {
-		cpPath := PathForModuleInPartitionInstall(ctx, "").Join(ctx, androidInfoTxtName)
-		builder.Command().
-			Text("rsync").
-			Flag("-a").
-			Input(androidInfoTxt).
-			Text(cpPath.String())
-	}
-	builder.Build("copy_android_info", "Copy android-info.txt")
 }
 
 // android_info module generate a file named android-info.txt that contains various information
diff --git a/filesystem/Android.bp b/filesystem/Android.bp
index 127faa75f..986b72edb 100644
--- a/filesystem/Android.bp
+++ b/filesystem/Android.bp
@@ -17,6 +17,7 @@ bootstrap_go_package {
     srcs: [
         "aconfig_files.go",
         "android_device.go",
+        "android_device_product_out.go",
         "avb_add_hash_footer.go",
         "avb_gen_vbmeta_image.go",
         "bootimg.go",
diff --git a/filesystem/android_device.go b/filesystem/android_device.go
index 080e6b747..e340602cd 100644
--- a/filesystem/android_device.go
+++ b/filesystem/android_device.go
@@ -62,6 +62,12 @@ type DeviceProperties struct {
 	Bootloader *string `android:"path"`
 	// Path to android-info.txt file containing board specific info.
 	Android_info *string `android:"path"`
+	// If this is the "main" android_device target for the build, i.e. the one that gets built
+	// when running a plain `m` command. Currently, this is the autogenerated android_device module
+	// in soong-only builds, but in the future when we check in android_device modules, the main
+	// one will be determined based on the lunch product. TODO: Figure out how to make this
+	// blueprint:"mutated" and still set it from filesystem_creator
+	Main_device *bool
 }
 
 type androidDevice struct {
@@ -70,9 +76,6 @@ type androidDevice struct {
 	partitionProps PartitionNameProperties
 
 	deviceProps DeviceProperties
-
-	// copyToProductOutTimestamp for copying necessary files to PRODUCT_OUT
-	copyToProductOutTimestamp android.WritablePath
 }
 
 func AndroidDeviceFactory() android.Module {
@@ -82,7 +85,7 @@ func AndroidDeviceFactory() android.Module {
 	return module
 }
 
-var numAutogeneratedAndroidDevicesOnceKey android.OnceKey = android.NewOnceKey("num_auto_generated_anroid_devices")
+var numMainAndroidDevicesOnceKey android.OnceKey = android.NewOnceKey("num_auto_generated_anroid_devices")
 
 type partitionDepTagType struct {
 	blueprint.BaseDependencyTag
@@ -123,25 +126,19 @@ func (a *androidDevice) DepsMutator(ctx android.BottomUpMutatorContext) {
 	}
 }
 
-func (a *androidDevice) copyToProductOut(ctx android.ModuleContext, builder *android.RuleBuilder, src android.Path, dest string) {
-	destPath := android.PathForModuleInPartitionInstall(ctx, "").Join(ctx, dest)
-	builder.Command().Text("rsync").Flag("-a").Flag("--checksum").Input(src).Text(destPath.String())
-}
-
-func (a *androidDevice) copyFilesToProductOut(ctx android.ModuleContext) {
-	a.copyToProductOutTimestamp = android.PathForModuleOut(ctx, "timestamp")
-	builder := android.NewRuleBuilder(pctx, ctx)
-	builder.Command().Text("touch").Output(a.copyToProductOutTimestamp)
-
-	// List all individual files to be copied to PRODUCT_OUT here
-	if a.deviceProps.Bootloader != nil {
-		a.copyToProductOut(ctx, builder, android.PathForModuleSrc(ctx, proptools.String(a.deviceProps.Bootloader)), "bootloader")
+func (a *androidDevice) GenerateAndroidBuildActions(ctx android.ModuleContext) {
+	if proptools.Bool(a.deviceProps.Main_device) {
+		numMainAndroidDevices := ctx.Config().Once(numMainAndroidDevicesOnceKey, func() interface{} {
+			return &atomic.Int32{}
+		}).(*atomic.Int32)
+		total := numMainAndroidDevices.Add(1)
+		if total > 1 {
+			// There should only be 1 main android_device module. That one will be
+			// made the default thing to build in soong-only builds.
+			ctx.ModuleErrorf("There cannot be more than 1 main android_device module")
+		}
 	}
 
-	builder.Build("copy_to_product_out", "Copy files to PRODUCT_OUT")
-}
-
-func (a *androidDevice) GenerateAndroidBuildActions(ctx android.ModuleContext) {
 	a.buildTargetFilesZip(ctx)
 	var deps []android.Path
 	if proptools.String(a.partitionProps.Super_partition_name) != "" {
@@ -193,36 +190,26 @@ func (a *androidDevice) GenerateAndroidBuildActions(ctx android.ModuleContext) {
 		deps = append(deps, imageOutput.DefaultOutputFiles[0])
 	})
 
-	a.copyFilesToProductOut(ctx)
-
 	allImagesStamp := android.PathForModuleOut(ctx, "all_images_stamp")
-	ctx.Build(pctx, android.BuildParams{
-		Rule:       android.Touch,
-		Output:     allImagesStamp,
-		Implicits:  deps,
-		Validation: a.copyToProductOutTimestamp,
-	})
-	ctx.SetOutputFiles(android.Paths{allImagesStamp}, "")
-	ctx.CheckbuildFile(allImagesStamp)
-
-	if ctx.OtherModuleIsAutoGenerated(ctx.Module()) {
-		numAutogeneratedAndroidDevices := ctx.Config().Once(numAutogeneratedAndroidDevicesOnceKey, func() interface{} {
-			return &atomic.Int32{}
-		}).(*atomic.Int32)
-		total := numAutogeneratedAndroidDevices.Add(1)
-		if total > 1 {
-			// There should only be 1 autogenerated android_device module. That one will be
-			// made the default thing to build in soong-only builds.
-			ctx.ModuleErrorf("There cannot be more than 1 autogenerated android_device module")
-		}
-	}
-
-	if !ctx.Config().KatiEnabled() && ctx.OtherModuleIsAutoGenerated(ctx.Module()) {
+	var validations android.Paths
+	if !ctx.Config().KatiEnabled() && proptools.Bool(a.deviceProps.Main_device) {
 		// In soong-only builds, build this module by default.
 		// This is the analogue to this make code:
 		// https://cs.android.com/android/platform/superproject/main/+/main:build/make/core/main.mk;l=1396;drc=6595459cdd8164a6008335f6372c9f97b9094060
 		ctx.Phony("droidcore-unbundled", allImagesStamp)
+
+		validations = append(validations, a.copyFilesToProductOutForSoongOnly(ctx))
 	}
+
+	ctx.Build(pctx, android.BuildParams{
+		Rule:        android.Touch,
+		Output:      allImagesStamp,
+		Implicits:   deps,
+		Validations: validations,
+	})
+
+	// Checkbuilding it causes soong to make a phony, so you can say `m <module name>`
+	ctx.CheckbuildFile(allImagesStamp)
 }
 
 // Helper structs for target_files.zip creation
@@ -335,7 +322,7 @@ func (a *androidDevice) buildTargetFilesZip(ctx android.ModuleContext) {
 
 	if a.deviceProps.Android_info != nil {
 		builder.Command().Textf("mkdir -p %s/OTA", targetFilesDir)
-		builder.Command().Textf("cp ").Input(android.PathForModuleSrc(ctx, proptools.String(a.deviceProps.Android_info))).Textf(" %s/OTA/android-info.txt", targetFilesDir)
+		builder.Command().Textf("cp ").Input(android.PathForModuleSrc(ctx, *a.deviceProps.Android_info)).Textf(" %s/OTA/android-info.txt", targetFilesDir)
 	}
 
 	a.copyImagesToTargetZip(ctx, builder, targetFilesDir)
diff --git a/filesystem/android_device_product_out.go b/filesystem/android_device_product_out.go
new file mode 100644
index 000000000..916c45adc
--- /dev/null
+++ b/filesystem/android_device_product_out.go
@@ -0,0 +1,214 @@
+// Copyright (C) 2025 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package filesystem
+
+import (
+	"android/soong/android"
+
+	"github.com/google/blueprint"
+	"github.com/google/blueprint/proptools"
+)
+
+var (
+	copyStagingDirRule = pctx.AndroidStaticRule("copy_staging_dir", blueprint.RuleParams{
+		Command: "rsync -a --checksum $dir/ $dest && touch $out",
+	}, "dir", "dest")
+)
+
+func (a *androidDevice) copyToProductOut(ctx android.ModuleContext, builder *android.RuleBuilder, src android.Path, dest string) {
+	destPath := android.PathForModuleInPartitionInstall(ctx, "").Join(ctx, dest)
+	builder.Command().Text("rsync").Flag("-a").Flag("--checksum").Input(src).Text(destPath.String())
+}
+
+func (a *androidDevice) copyFilesToProductOutForSoongOnly(ctx android.ModuleContext) android.Path {
+	filesystemInfos := a.getFsInfos(ctx)
+
+	// The current logic to copy the staging directories to PRODUCT_OUT isn't very sound.
+	// We only track dependencies on the image file, so if the image file wasn't changed, the
+	// staging directory won't be re-copied. If you do an installclean, it would remove the copied
+	// staging directories but not affect the intermediates path image file, so the next build
+	// wouldn't re-copy them. As a hack, create a presence detector that would be deleted on
+	// an installclean to use as a dep for the staging dir copies.
+	productOutPresenceDetector := android.PathForModuleInPartitionInstall(ctx, "", "product_out_presence_detector.txt")
+	ctx.Build(pctx, android.BuildParams{
+		Rule:   android.Touch,
+		Output: productOutPresenceDetector,
+	})
+
+	var deps android.Paths
+
+	for _, partition := range android.SortedKeys(filesystemInfos) {
+		info := filesystemInfos[partition]
+		imgInstallPath := android.PathForModuleInPartitionInstall(ctx, "", partition+".img")
+		ctx.Build(pctx, android.BuildParams{
+			Rule:   android.Cp,
+			Input:  info.Output,
+			Output: imgInstallPath,
+		})
+		dirStamp := android.PathForModuleOut(ctx, partition+"_staging_dir_copy_stamp.txt")
+		dirInstallPath := android.PathForModuleInPartitionInstall(ctx, "", partition)
+		ctx.Build(pctx, android.BuildParams{
+			Rule:   copyStagingDirRule,
+			Output: dirStamp,
+			Implicits: []android.Path{
+				info.Output,
+				productOutPresenceDetector,
+			},
+			Args: map[string]string{
+				"dir":  info.RebasedDir.String(),
+				"dest": dirInstallPath.String(),
+			},
+		})
+
+		// Make it so doing `m <moduleName>` or `m <partitionType>` will copy the files to
+		// PRODUCT_OUT
+		ctx.Phony(info.ModuleName, dirStamp, imgInstallPath)
+		ctx.Phony(partition, dirStamp, imgInstallPath)
+
+		deps = append(deps, imgInstallPath, dirStamp)
+	}
+
+	// List all individual files to be copied to PRODUCT_OUT here
+	if a.deviceProps.Bootloader != nil {
+		bootloaderInstallPath := android.PathForModuleInPartitionInstall(ctx, "", "bootloader")
+		ctx.Build(pctx, android.BuildParams{
+			Rule:   android.Cp,
+			Input:  android.PathForModuleSrc(ctx, *a.deviceProps.Bootloader),
+			Output: bootloaderInstallPath,
+		})
+		deps = append(deps, bootloaderInstallPath)
+	}
+
+	copyBootImg := func(prop *string, type_ string) {
+		if proptools.String(prop) != "" {
+			partition := ctx.GetDirectDepWithTag(*prop, filesystemDepTag)
+			if info, ok := android.OtherModuleProvider(ctx, partition, BootimgInfoProvider); ok {
+				installPath := android.PathForModuleInPartitionInstall(ctx, "", type_+".img")
+				ctx.Build(pctx, android.BuildParams{
+					Rule:   android.Cp,
+					Input:  info.Output,
+					Output: installPath,
+				})
+				deps = append(deps, installPath)
+			} else {
+				ctx.ModuleErrorf("%s does not set BootimgInfo\n", *prop)
+			}
+		}
+	}
+
+	copyBootImg(a.partitionProps.Init_boot_partition_name, "init_boot")
+	copyBootImg(a.partitionProps.Boot_partition_name, "boot")
+	copyBootImg(a.partitionProps.Vendor_boot_partition_name, "vendor_boot")
+
+	for _, vbmetaModName := range a.partitionProps.Vbmeta_partitions {
+		partition := ctx.GetDirectDepWithTag(vbmetaModName, filesystemDepTag)
+		if info, ok := android.OtherModuleProvider(ctx, partition, vbmetaPartitionProvider); ok {
+			installPath := android.PathForModuleInPartitionInstall(ctx, "", info.Name+".img")
+			ctx.Build(pctx, android.BuildParams{
+				Rule:   android.Cp,
+				Input:  info.Output,
+				Output: installPath,
+			})
+			deps = append(deps, installPath)
+		} else {
+			ctx.ModuleErrorf("%s does not set vbmetaPartitionProvider\n", vbmetaModName)
+		}
+	}
+
+	if proptools.String(a.partitionProps.Super_partition_name) != "" {
+		partition := ctx.GetDirectDepWithTag(*a.partitionProps.Super_partition_name, superPartitionDepTag)
+		if info, ok := android.OtherModuleProvider(ctx, partition, SuperImageProvider); ok {
+			installPath := android.PathForModuleInPartitionInstall(ctx, "", "super.img")
+			ctx.Build(pctx, android.BuildParams{
+				Rule:   android.Cp,
+				Input:  info.SuperImage,
+				Output: installPath,
+			})
+			deps = append(deps, installPath)
+		} else {
+			ctx.ModuleErrorf("%s does not set SuperImageProvider\n", *a.partitionProps.Super_partition_name)
+		}
+	}
+
+	if proptools.String(a.deviceProps.Android_info) != "" {
+		installPath := android.PathForModuleInPartitionInstall(ctx, "", "android_info.txt")
+		ctx.Build(pctx, android.BuildParams{
+			Rule:   android.Cp,
+			Input:  android.PathForModuleSrc(ctx, *a.deviceProps.Android_info),
+			Output: installPath,
+		})
+		deps = append(deps, installPath)
+	}
+
+	copyToProductOutTimestamp := android.PathForModuleOut(ctx, "product_out_copy_timestamp")
+	ctx.Build(pctx, android.BuildParams{
+		Rule:      android.Touch,
+		Output:    copyToProductOutTimestamp,
+		Implicits: deps,
+	})
+
+	return copyToProductOutTimestamp
+}
+
+// Returns a mapping from partition type -> FilesystemInfo. This includes filesystems that are
+// nested inside of other partitions, such as the partitions inside super.img, or ramdisk inside
+// of boot.
+func (a *androidDevice) getFsInfos(ctx android.ModuleContext) map[string]FilesystemInfo {
+	type propToType struct {
+		prop *string
+		ty   string
+	}
+
+	filesystemInfos := make(map[string]FilesystemInfo)
+
+	partitionDefinitions := []propToType{
+		propToType{a.partitionProps.System_partition_name, "system"},
+		propToType{a.partitionProps.System_ext_partition_name, "system_ext"},
+		propToType{a.partitionProps.Product_partition_name, "product"},
+		propToType{a.partitionProps.Vendor_partition_name, "vendor"},
+		propToType{a.partitionProps.Odm_partition_name, "odm"},
+		propToType{a.partitionProps.Recovery_partition_name, "recovery"},
+		propToType{a.partitionProps.System_dlkm_partition_name, "system_dlkm"},
+		propToType{a.partitionProps.Vendor_dlkm_partition_name, "vendor_dlkm"},
+		propToType{a.partitionProps.Odm_dlkm_partition_name, "odm_dlkm"},
+		propToType{a.partitionProps.Userdata_partition_name, "userdata"},
+		// filesystemInfo from init_boot and vendor_boot actually are re-exports of the ramdisk
+		// images inside of them
+		propToType{a.partitionProps.Init_boot_partition_name, "ramdisk"},
+		propToType{a.partitionProps.Vendor_boot_partition_name, "vendor_ramdisk"},
+	}
+	for _, partitionDefinition := range partitionDefinitions {
+		if proptools.String(partitionDefinition.prop) != "" {
+			partition := ctx.GetDirectDepWithTag(*partitionDefinition.prop, filesystemDepTag)
+			if info, ok := android.OtherModuleProvider(ctx, partition, FilesystemProvider); ok {
+				filesystemInfos[partitionDefinition.ty] = info
+			} else {
+				ctx.ModuleErrorf("Super partition %s does not set FilesystemProvider\n", partition.Name())
+			}
+		}
+	}
+	if a.partitionProps.Super_partition_name != nil {
+		superPartition := ctx.GetDirectDepWithTag(*a.partitionProps.Super_partition_name, superPartitionDepTag)
+		if info, ok := android.OtherModuleProvider(ctx, superPartition, SuperImageProvider); ok {
+			for partition := range info.SubImageInfo {
+				filesystemInfos[partition] = info.SubImageInfo[partition]
+			}
+		} else {
+			ctx.ModuleErrorf("Super partition %s does not set SuperImageProvider\n", superPartition.Name())
+		}
+	}
+
+	return filesystemInfos
+}
diff --git a/filesystem/bootimg.go b/filesystem/bootimg.go
index 803e98130..98af47991 100644
--- a/filesystem/bootimg.go
+++ b/filesystem/bootimg.go
@@ -420,10 +420,6 @@ func (b *bootimg) addAvbFooter(ctx android.ModuleContext, unsignedImage android.
 		cmd.FlagWithArg("--rollback_index ", strconv.FormatInt(*b.properties.Avb_rollback_index, 10))
 	}
 
-	if !ctx.Config().KatiEnabled() {
-		copyImageFileToProductOut(ctx, builder, b.bootImageType.String(), output)
-	}
-
 	builder.Build("add_avb_footer", fmt.Sprintf("Adding avb footer to %s", b.BaseModuleName()))
 	return output
 }
@@ -439,10 +435,6 @@ func (b *bootimg) signImage(ctx android.ModuleContext, unsignedImage android.Pat
 		Implicits(toolDeps).
 		Output(output)
 
-	if !ctx.Config().KatiEnabled() {
-		copyImageFileToProductOut(ctx, builder, b.bootImageType.String(), output)
-	}
-
 	builder.Build("sign_bootimg", fmt.Sprintf("Signing %s", b.BaseModuleName()))
 	return output
 }
diff --git a/filesystem/filesystem.go b/filesystem/filesystem.go
index 37a296536..306710c7c 100644
--- a/filesystem/filesystem.go
+++ b/filesystem/filesystem.go
@@ -373,9 +373,17 @@ type FilesystemInfo struct {
 	// to add a dependency on the Output file, as you cannot add dependencies on directories
 	// in ninja.
 	RootDir android.Path
+	// The rebased staging directory used to build the output filesystem. If consuming this, make
+	// sure to add a dependency on the Output file, as you cannot add dependencies on directories
+	// in ninja. In many cases this is the same as RootDir, only in the system partition is it
+	// different. There, it points to the "system" sub-directory of RootDir.
+	RebasedDir android.Path
 	// A text file with block data of the .img file
 	// This is an implicit output of `build_image`
 	MapFile android.Path
+	// Name of the module that produced this FilesystemInfo origionally. (though it may be
+	// re-exported by super images or boot images)
+	ModuleName string
 }
 
 var FilesystemProvider = blueprint.NewProvider[FilesystemInfo]()
@@ -463,16 +471,17 @@ func (f *filesystem) GenerateAndroidBuildActions(ctx android.ModuleContext) {
 	}
 
 	var rootDir android.OutputPath
+	var rebasedDir android.OutputPath
 	var mapFile android.Path
 	var outputHermetic android.Path
 	switch f.fsType(ctx) {
 	case ext4Type, erofsType, f2fsType:
-		f.output, outputHermetic, rootDir = f.buildImageUsingBuildImage(ctx)
+		f.output, outputHermetic, rootDir, rebasedDir = f.buildImageUsingBuildImage(ctx)
 		mapFile = f.getMapFile(ctx)
 	case compressedCpioType:
-		f.output, rootDir = f.buildCpioImage(ctx, true)
+		f.output, rootDir, rebasedDir = f.buildCpioImage(ctx, true)
 	case cpioType:
-		f.output, rootDir = f.buildCpioImage(ctx, false)
+		f.output, rootDir, rebasedDir = f.buildCpioImage(ctx, false)
 	default:
 		return
 	}
@@ -492,6 +501,8 @@ func (f *filesystem) GenerateAndroidBuildActions(ctx android.ModuleContext) {
 		Output:       f.output,
 		FileListFile: fileListFile,
 		RootDir:      rootDir,
+		RebasedDir:   rebasedDir,
+		ModuleName:   ctx.ModuleName(),
 	}
 	if mapFile != nil {
 		fsInfo.MapFile = mapFile
@@ -640,24 +651,11 @@ func (f *filesystem) copyPackagingSpecs(ctx android.ModuleContext, builder *andr
 	return f.CopySpecsToDirs(ctx, builder, dirsToSpecs)
 }
 
-func (f *filesystem) copyFilesToProductOut(ctx android.ModuleContext, builder *android.RuleBuilder, rebasedDir android.OutputPath) {
-	if !(f.Name() == ctx.Config().SoongDefinedSystemImage() || proptools.Bool(f.properties.Is_auto_generated)) {
-		return
-	}
-	installPath := android.PathForModuleInPartitionInstall(ctx, f.partitionName())
-	builder.Command().Textf("rsync --checksum %s %s", rebasedDir, installPath)
-}
-
-func copyImageFileToProductOut(ctx android.ModuleContext, builder *android.RuleBuilder, partition string, output android.Path) {
-	copyDir := android.PathForModuleInPartitionInstall(ctx, "").Join(ctx, fmt.Sprintf("%s.img", partition))
-	builder.Command().Textf("rsync -a %s %s", output, copyDir)
-}
-
 func (f *filesystem) rootDirString() string {
 	return f.partitionName()
 }
 
-func (f *filesystem) buildImageUsingBuildImage(ctx android.ModuleContext) (android.Path, android.Path, android.OutputPath) {
+func (f *filesystem) buildImageUsingBuildImage(ctx android.ModuleContext) (android.Path, android.Path, android.OutputPath, android.OutputPath) {
 	rootDir := android.PathForModuleOut(ctx, f.rootDirString()).OutputPath
 	rebasedDir := rootDir
 	if f.properties.Base_dir != nil {
@@ -675,7 +673,6 @@ func (f *filesystem) buildImageUsingBuildImage(ctx android.ModuleContext) (andro
 	f.buildEventLogtagsFile(ctx, builder, rebasedDir)
 	f.buildAconfigFlagsFiles(ctx, builder, specs, rebasedDir)
 	f.filesystemBuilder.BuildLinkerConfigFile(ctx, builder, rebasedDir)
-	f.copyFilesToProductOut(ctx, builder, rebasedDir)
 
 	// run host_init_verifier
 	// Ideally we should have a concept of pluggable linters that verify the generated image.
@@ -720,10 +717,6 @@ func (f *filesystem) buildImageUsingBuildImage(ctx android.ModuleContext) (andro
 		Output(outputHermetic).
 		Text(rootDir.String()) // directory where to find fs_config_files|dirs
 
-	if !ctx.Config().KatiEnabled() {
-		copyImageFileToProductOut(ctx, builder, f.partitionName(), output)
-	}
-
 	if f.properties.Partition_size != nil {
 		assertMaxImageSize(builder, output, *f.properties.Partition_size, false)
 	}
@@ -731,7 +724,7 @@ func (f *filesystem) buildImageUsingBuildImage(ctx android.ModuleContext) (andro
 	// rootDir is not deleted. Might be useful for quick inspection.
 	builder.Build("build_filesystem_image", fmt.Sprintf("Creating filesystem %s", f.BaseModuleName()))
 
-	return output, outputHermetic, rootDir
+	return output, outputHermetic, rootDir, rebasedDir
 }
 
 func (f *filesystem) buildFileContexts(ctx android.ModuleContext) android.Path {
@@ -925,7 +918,7 @@ func includeFilesRootDir(ctx android.ModuleContext) (rootDirs android.Paths, par
 	return rootDirs, partitions
 }
 
-func (f *filesystem) buildCpioImage(ctx android.ModuleContext, compressed bool) (android.Path, android.OutputPath) {
+func (f *filesystem) buildCpioImage(ctx android.ModuleContext, compressed bool) (android.Path, android.OutputPath, android.OutputPath) {
 	if proptools.Bool(f.properties.Use_avb) {
 		ctx.PropertyErrorf("use_avb", "signing compresed cpio image using avbtool is not supported."+
 			"Consider adding this to bootimg module and signing the entire boot image.")
@@ -955,7 +948,6 @@ func (f *filesystem) buildCpioImage(ctx android.ModuleContext, compressed bool)
 	f.buildEventLogtagsFile(ctx, builder, rebasedDir)
 	f.buildAconfigFlagsFiles(ctx, builder, specs, rebasedDir)
 	f.filesystemBuilder.BuildLinkerConfigFile(ctx, builder, rebasedDir)
-	f.copyFilesToProductOut(ctx, builder, rebasedDir)
 
 	rootDirs, partitions := includeFilesRootDir(ctx)
 
@@ -986,7 +978,7 @@ func (f *filesystem) buildCpioImage(ctx android.ModuleContext, compressed bool)
 	// rootDir is not deleted. Might be useful for quick inspection.
 	builder.Build("build_cpio_image", fmt.Sprintf("Creating filesystem %s", f.BaseModuleName()))
 
-	return output, rootDir
+	return output, rootDir, rebasedDir
 }
 
 var validPartitions = []string{
diff --git a/filesystem/vbmeta.go b/filesystem/vbmeta.go
index d468186f0..fb766875d 100644
--- a/filesystem/vbmeta.go
+++ b/filesystem/vbmeta.go
@@ -284,10 +284,6 @@ func (v *vbmeta) GenerateAndroidBuildActions(ctx android.ModuleContext) {
 		FlagWithArg("-s ", strconv.Itoa(vbmetaMaxSize)).
 		Output(output)
 
-	if !ctx.Config().KatiEnabled() {
-		copyImageFileToProductOut(ctx, builder, v.partitionName(), output)
-	}
-
 	builder.Build("vbmeta", fmt.Sprintf("vbmeta %s", ctx.ModuleName()))
 
 	v.installDir = android.PathForModuleInstall(ctx, "etc")
diff --git a/fsgen/filesystem_creator.go b/fsgen/filesystem_creator.go
index 9aed460e3..fbe907898 100644
--- a/fsgen/filesystem_creator.go
+++ b/fsgen/filesystem_creator.go
@@ -263,7 +263,9 @@ func (f *filesystemCreator) createDeviceModule(
 	}
 	partitionProps.Vbmeta_partitions = vbmetaPartitions
 
-	deviceProps := &filesystem.DeviceProperties{}
+	deviceProps := &filesystem.DeviceProperties{
+		Main_device: proptools.BoolPtr(true),
+	}
 	if bootloader, ok := f.createBootloaderFilegroup(ctx); ok {
 		deviceProps.Bootloader = proptools.StringPtr(":" + bootloader)
 	}
diff --git a/fsgen/vbmeta_partitions.go b/fsgen/vbmeta_partitions.go
index d146cbb8e..a75f59ce8 100644
--- a/fsgen/vbmeta_partitions.go
+++ b/fsgen/vbmeta_partitions.go
@@ -173,6 +173,7 @@ func createVbmetaPartitions(ctx android.LoadHookContext, generatedPartitionTypes
 			Rollback_index:     ri,
 			Chained_partitions: chainedPartitions,
 			Partitions:         proptools.NewSimpleConfigurable(partitionModules),
+			Partition_name:     proptools.StringPtr("vbmeta"),
 		}, &struct {
 			Name *string
 		}{