Copy boot*img files to IMAGES/ in target_files.zip

The following files will be added to the soong-built zip file after this
CL
1. boot.img
2. init_boot.img
3. vendor_boot.img

Implementation details
- Refactor the IMAGES/ dir generation to `copyImagesToTargetZip`
- Use VisitDirectDepsProxyWithTag(filesystemDepTag) to visit filesystem
  deps (e.g. system) and bootimg deps.
- Set the Stem property in the autogenerated soong bootimg modules

Test: Built the zip file locally and verified that boot*img files are
copied (although they are not bit identical to the make equivalents)
Bug: 388635097

Change-Id: Ia7138761146e93a268eecb9dc9bf4d579445980b
diff --git a/filesystem/android_device.go b/filesystem/android_device.go
index 86f853c..fe30a26 100644
--- a/filesystem/android_device.go
+++ b/filesystem/android_device.go
@@ -208,17 +208,6 @@
 		toCopy = append(toCopy, targetFilesZipCopy{a.partitionProps.Recovery_partition_name, "VENDOR_BOOT/RAMDISK"})
 	}
 
-	// Create an IMAGES/ subdirectory
-	builder.Command().Textf("mkdir -p %s/IMAGES/", targetFilesDir.String())
-	if a.deviceProps.Bootloader != nil {
-		builder.Command().Textf("cp ").Input(android.PathForModuleSrc(ctx, proptools.String(a.deviceProps.Bootloader))).Textf(" %s/IMAGES/bootloader", targetFilesDir.String())
-	}
-	// Copy the vbmeta img files to IMAGES/
-	for _, vbmetaPartition := range a.partitionProps.Vbmeta_partitions {
-		vbmetaInfo, _ := android.OtherModuleProvider(ctx, ctx.GetDirectDepWithTag(vbmetaPartition, filesystemDepTag), vbmetaPartitionProvider)
-		builder.Command().Textf("cp ").Input(vbmetaInfo.Output).Textf(" %s/IMAGES/", targetFilesDir.String())
-	}
-
 	for _, zipCopy := range toCopy {
 		if zipCopy.srcModule == nil {
 			continue
@@ -239,7 +228,6 @@
 			// Create the ROOT partition in target_files.zip
 			builder.Command().Textf("rsync --links --exclude=system/* %s/ -r %s/ROOT", fsInfo.RootDir, targetFilesDir.String())
 		}
-		builder.Command().Textf("cp %s %s/IMAGES/", fsInfo.Output, targetFilesDir.String())
 	}
 	// Copy cmdline, kernel etc. files of boot images
 	if a.partitionProps.Vendor_boot_partition_name != nil {
@@ -276,6 +264,8 @@
 		builder.Command().Textf("cp %s %s/OTA/android-info.txt", android.PathForModuleSrc(ctx, proptools.String(a.deviceProps.Android_info)), targetFilesDir)
 	}
 
+	a.copyImagesToTargetZip(ctx, builder, targetFilesDir)
+
 	builder.Command().
 		BuiltTool("soong_zip").
 		Text("-d").
@@ -286,6 +276,26 @@
 	builder.Build("target_files_"+ctx.ModuleName(), "Build target_files.zip")
 }
 
+func (a *androidDevice) copyImagesToTargetZip(ctx android.ModuleContext, builder *android.RuleBuilder, targetFilesDir android.WritablePath) {
+	// Create an IMAGES/ subdirectory
+	builder.Command().Textf("mkdir -p %s/IMAGES", targetFilesDir.String())
+	if a.deviceProps.Bootloader != nil {
+		builder.Command().Textf("cp ").Input(android.PathForModuleSrc(ctx, proptools.String(a.deviceProps.Bootloader))).Textf(" %s/IMAGES/bootloader", targetFilesDir.String())
+	}
+	// Copy the filesystem ,boot and vbmeta img files to IMAGES/
+	ctx.VisitDirectDepsProxyWithTag(filesystemDepTag, func(child android.ModuleProxy) {
+		if info, ok := android.OtherModuleProvider(ctx, child, FilesystemProvider); ok {
+			builder.Command().Textf("cp ").Input(info.Output).Textf(" %s/IMAGES/", targetFilesDir.String())
+		} else if info, ok := android.OtherModuleProvider(ctx, child, BootimgInfoProvider); ok {
+			builder.Command().Textf("cp ").Input(info.Output).Textf(" %s/IMAGES/", targetFilesDir.String())
+		} else if info, ok := android.OtherModuleProvider(ctx, child, vbmetaPartitionProvider); ok {
+			builder.Command().Textf("cp ").Input(info.Output).Textf(" %s/IMAGES/", targetFilesDir.String())
+		} else {
+			ctx.ModuleErrorf("Module %s does not provide an .img file output for target_files.zip", child.Name())
+		}
+	})
+}
+
 func (a *androidDevice) getFilesystemInfo(ctx android.ModuleContext, depName string) FilesystemInfo {
 	fsMod := ctx.GetDirectDepWithTag(depName, filesystemDepTag)
 	fsInfo, ok := android.OtherModuleProvider(ctx, fsMod, FilesystemProvider)