Implement directed vendor snapshot

Vendors can now generate only needed modules by setting the following
Makefile variables:

- DIRECTED_VENDOR_SNAPSHOT: set to true
- VENDOR_SNAPSHOT_MODULES: list of snapshot candidates

e.g.

DIRECTED_VENDOR_SNAPSHOT := true
VENDOR_SNAPSHOT_MODULES := toybox_vendor sh_vendor libbase libcutils ...

Bug: 157967325
Test: m dist vendor-snapshot after setting those in BoardConfig.mk
Change-Id: I6515a43d9358d94483d7c7fa2b066f9dd457f6aa
diff --git a/android/config.go b/android/config.go
index a7e0b67..1f52b6c 100644
--- a/android/config.go
+++ b/android/config.go
@@ -1401,6 +1401,14 @@
 	return c.config.productVariables.BoardReqdMaskPolicy
 }
 
+func (c *deviceConfig) DirectedVendorSnapshot() bool {
+	return c.config.productVariables.DirectedVendorSnapshot
+}
+
+func (c *deviceConfig) VendorSnapshotModules() map[string]bool {
+	return c.config.productVariables.VendorSnapshotModules
+}
+
 // The ConfiguredJarList struct provides methods for handling a list of (apex, jar) pairs.
 // Such lists are used in the build system for things like bootclasspath jars or system server jars.
 // The apex part is either an apex name, or a special names "platform" or "system_ext". Jar is a
diff --git a/android/variable.go b/android/variable.go
index 41cd337..48c4be6 100644
--- a/android/variable.go
+++ b/android/variable.go
@@ -309,6 +309,9 @@
 	VndkUseCoreVariant         *bool `json:",omitempty"`
 	VndkSnapshotBuildArtifacts *bool `json:",omitempty"`
 
+	DirectedVendorSnapshot bool            `json:",omitempty"`
+	VendorSnapshotModules  map[string]bool `json:",omitempty"`
+
 	BoardVendorSepolicyDirs      []string `json:",omitempty"`
 	BoardOdmSepolicyDirs         []string `json:",omitempty"`
 	BoardReqdMaskPolicy          []string `json:",omitempty"`
diff --git a/cc/cc_test.go b/cc/cc_test.go
index f422839..eef13da 100644
--- a/cc/cc_test.go
+++ b/cc/cc_test.go
@@ -1266,6 +1266,95 @@
 	}
 }
 
+func TestVendorSnapshotDirected(t *testing.T) {
+	bp := `
+	cc_library_shared {
+		name: "libvendor",
+		vendor: true,
+		nocrt: true,
+	}
+
+	cc_library_shared {
+		name: "libvendor_available",
+		vendor_available: true,
+		nocrt: true,
+	}
+
+	genrule {
+		name: "libfoo_gen",
+		cmd: "",
+		out: ["libfoo.so"],
+	}
+
+	cc_prebuilt_library_shared {
+		name: "libfoo",
+		vendor: true,
+		prefer: true,
+		srcs: [":libfoo_gen"],
+	}
+
+	cc_library_shared {
+		name: "libfoo",
+		vendor: true,
+		nocrt: true,
+	}
+`
+	config := TestConfig(buildDir, android.Android, nil, bp, nil)
+	config.TestProductVariables.DeviceVndkVersion = StringPtr("current")
+	config.TestProductVariables.Platform_vndk_version = StringPtr("VER")
+	config.TestProductVariables.DirectedVendorSnapshot = true
+	config.TestProductVariables.VendorSnapshotModules = make(map[string]bool)
+	config.TestProductVariables.VendorSnapshotModules["libvendor"] = true
+	config.TestProductVariables.VendorSnapshotModules["libfoo"] = true
+	ctx := testCcWithConfig(t, config)
+
+	// Check Vendor snapshot output.
+
+	snapshotDir := "vendor-snapshot"
+	snapshotVariantPath := filepath.Join(buildDir, snapshotDir, "arm64")
+	snapshotSingleton := ctx.SingletonForTests("vendor-snapshot")
+
+	var includeJsonFiles []string
+	var excludeJsonFiles []string
+
+	for _, arch := range [][]string{
+		[]string{"arm64", "armv8-a"},
+		[]string{"arm", "armv7-a-neon"},
+	} {
+		archType := arch[0]
+		archVariant := arch[1]
+		archDir := fmt.Sprintf("arch-%s-%s", archType, archVariant)
+
+		sharedVariant := fmt.Sprintf("android_vendor.VER_%s_%s_shared", archType, archVariant)
+		sharedDir := filepath.Join(snapshotVariantPath, archDir, "shared")
+
+		// Included modules
+		checkSnapshot(t, ctx, snapshotSingleton, "libvendor", "libvendor.so", sharedDir, sharedVariant)
+		includeJsonFiles = append(includeJsonFiles, filepath.Join(sharedDir, "libvendor.so.json"))
+		// Check that snapshot captures "prefer: true" prebuilt
+		checkSnapshot(t, ctx, snapshotSingleton, "prebuilt_libfoo", "libfoo.so", sharedDir, sharedVariant)
+		includeJsonFiles = append(includeJsonFiles, filepath.Join(sharedDir, "libfoo.so.json"))
+
+		// Excluded modules
+		checkSnapshotExclude(t, ctx, snapshotSingleton, "libvendor_available", "libvendor_available.so", sharedDir, sharedVariant)
+		excludeJsonFiles = append(excludeJsonFiles, filepath.Join(sharedDir, "libvendor_available.so.json"))
+	}
+
+	// Verify that each json file for an included module has a rule.
+	for _, jsonFile := range includeJsonFiles {
+		if snapshotSingleton.MaybeOutput(jsonFile).Rule == nil {
+			t.Errorf("include json file %q not found", jsonFile)
+		}
+	}
+
+	// Verify that each json file for an excluded module has no rule.
+	for _, jsonFile := range excludeJsonFiles {
+		if snapshotSingleton.MaybeOutput(jsonFile).Rule != nil {
+			t.Errorf("exclude json file %q found", jsonFile)
+		}
+	}
+}
+
 func TestVendorSnapshotUse(t *testing.T) {
 	frameworkBp := `
 	cc_library {
diff --git a/cc/snapshot_prebuilt.go b/cc/snapshot_prebuilt.go
index 8979846..d9c46ea 100644
--- a/cc/snapshot_prebuilt.go
+++ b/cc/snapshot_prebuilt.go
@@ -82,6 +82,11 @@
 
 	// Whether to skip the source mutator for a given module.
 	skipSourceMutator(ctx android.BottomUpMutatorContext) bool
+
+	// Whether to exclude a given module from the directed snapshot or not.
+	// If the makefile variable DIRECTED_{IMAGE}_SNAPSHOT is true, directed snapshot is turned on,
+	// and only modules listed in {IMAGE}_SNAPSHOT_MODULES will be captured.
+	excludeFromDirectedSnapshot(cfg android.DeviceConfig, name string) bool
 }
 
 type vendorSnapshotImage struct{}
@@ -193,6 +198,16 @@
 	return false
 }
 
+// returns true iff a given module SHOULD BE EXCLUDED, false if included
+func (vendorSnapshotImage) excludeFromDirectedSnapshot(cfg android.DeviceConfig, name string) bool {
+	// If we're using full snapshot, not directed snapshot, capture every module
+	if !cfg.DirectedVendorSnapshot() {
+		return false
+	}
+	// Else, checks if name is in VENDOR_SNAPSHOT_MODULES.
+	return !cfg.VendorSnapshotModules()[name]
+}
+
 func (recoverySnapshotImage) init() {
 	android.RegisterSingletonType("recovery-snapshot", RecoverySnapshotSingleton)
 	android.RegisterModuleType("recovery_snapshot_shared", RecoverySnapshotSharedFactory)
@@ -275,6 +290,11 @@
 	return !ok || !module.InRecovery()
 }
 
+func (recoverySnapshotImage) excludeFromDirectedSnapshot(cfg android.DeviceConfig, name string) bool {
+	// directed recovery snapshot is not implemented yet
+	return false
+}
+
 var vendorSnapshotImageSingleton vendorSnapshotImage
 var recoverySnapshotImageSingleton recoverySnapshotImage
 
diff --git a/cc/snapshot_utils.go b/cc/snapshot_utils.go
index 3e6444b..c50ef45 100644
--- a/cc/snapshot_utils.go
+++ b/cc/snapshot_utils.go
@@ -80,7 +80,7 @@
 	}
 
 	for _, image := range []snapshotImage{vendorSnapshotImageSingleton, recoverySnapshotImageSingleton} {
-		if isSnapshotAware(m, image.isProprietaryPath(ctx.ModuleDir()), apexInfo, image) {
+		if isSnapshotAware(ctx.DeviceConfig(), m, image.isProprietaryPath(ctx.ModuleDir()), apexInfo, image) {
 			return true
 		}
 	}
diff --git a/cc/vendor_snapshot.go b/cc/vendor_snapshot.go
index 0a89e47..6bd095f 100644
--- a/cc/vendor_snapshot.go
+++ b/cc/vendor_snapshot.go
@@ -198,7 +198,7 @@
 }
 
 // Determines if the module is a candidate for snapshot.
-func isSnapshotAware(m *Module, inProprietaryPath bool, apexInfo android.ApexInfo, image snapshotImage) bool {
+func isSnapshotAware(cfg android.DeviceConfig, m *Module, inProprietaryPath bool, apexInfo android.ApexInfo, image snapshotImage) bool {
 	if !m.Enabled() || m.Properties.HideFromMake {
 		return false
 	}
@@ -241,6 +241,10 @@
 	if _, ok := m.linker.(*llndkHeadersDecorator); ok {
 		return false
 	}
+	// If we are using directed snapshot AND we have to exclude this module, skip this
+	if image.excludeFromDirectedSnapshot(cfg, m.BaseModuleName()) {
+		return false
+	}
 
 	// Libraries
 	if l, ok := m.linker.(snapshotLibraryInterface); ok {
@@ -535,7 +539,7 @@
 			}
 		}
 
-		if !isSnapshotAware(m, inProprietaryPath, apexInfo, c.image) {
+		if !isSnapshotAware(ctx.DeviceConfig(), m, inProprietaryPath, apexInfo, c.image) {
 			return
 		}