Modify autogen vbmeta chain and include partitions logic

This change modifies the logic in selecting the chain and include
partitions for the top level vbmeta partition generation to follow that
in make.

Implementation details:
- Introduce separate functions for selecting the include and chain
  partitions
- Set vbmetaPartitionProvider in boot images and filesystem modules
  given that boot images can be selected as chained partitions.

Test: compare soong and make generated vbmeta partitions command
Bug: 390204058
Change-Id: I7940629212c4624e88d25b162755205503cfa469
diff --git a/filesystem/bootimg.go b/filesystem/bootimg.go
index 98af479..6d6c15c 100644
--- a/filesystem/bootimg.go
+++ b/filesystem/bootimg.go
@@ -99,6 +99,9 @@
 	// The index used to prevent rollback of the image on device.
 	Avb_rollback_index *int64
 
+	// Rollback index location of this image. Must be 0, 1, 2, etc.
+	Avb_rollback_index_location *int64
+
 	// The security patch passed to as the com.android.build.<type>.security_patch avb property.
 	// Replacement for the make variables BOOT_SECURITY_PATCH / INIT_BOOT_SECURITY_PATCH.
 	Security_patch *string
@@ -237,6 +240,27 @@
 		Bootconfig: b.getBootconfigPath(ctx),
 		Output:     output,
 	})
+
+	extractedPublicKey := android.PathForModuleOut(ctx, b.partitionName()+".avbpubkey")
+	if b.properties.Avb_private_key != nil {
+		key := android.PathForModuleSrc(ctx, proptools.String(b.properties.Avb_private_key))
+		ctx.Build(pctx, android.BuildParams{
+			Rule:   extractPublicKeyRule,
+			Input:  key,
+			Output: extractedPublicKey,
+		})
+	}
+	var ril int
+	if b.properties.Avb_rollback_index_location != nil {
+		ril = proptools.Int(b.properties.Avb_rollback_index_location)
+	}
+
+	android.SetProvider(ctx, vbmetaPartitionProvider, vbmetaPartitionInfo{
+		Name:                  b.bootImageType.String(),
+		RollbackIndexLocation: ril,
+		PublicKey:             extractedPublicKey,
+		Output:                output,
+	})
 }
 
 var BootimgInfoProvider = blueprint.NewProvider[BootimgInfo]()
diff --git a/filesystem/filesystem.go b/filesystem/filesystem.go
index 68cbee9..357ec32 100644
--- a/filesystem/filesystem.go
+++ b/filesystem/filesystem.go
@@ -131,6 +131,9 @@
 	// The index used to prevent rollback of the image. Only used if use_avb is true.
 	Rollback_index *int64
 
+	// Rollback index location of this image. Must be 0, 1, 2, etc.
+	Rollback_index_location *int64
+
 	// Name of the partition stored in vbmeta desc. Defaults to the name of this module.
 	Partition_name *string
 
@@ -537,6 +540,33 @@
 	if proptools.Bool(f.properties.Unchecked_module) {
 		ctx.UncheckedModule()
 	}
+
+	f.setVbmetaPartitionProvider(ctx)
+}
+
+func (f *filesystem) setVbmetaPartitionProvider(ctx android.ModuleContext) {
+	var extractedPublicKey android.ModuleOutPath
+	if f.properties.Avb_private_key != nil {
+		key := android.PathForModuleSrc(ctx, proptools.String(f.properties.Avb_private_key))
+		extractedPublicKey = android.PathForModuleOut(ctx, f.partitionName()+".avbpubkey")
+		ctx.Build(pctx, android.BuildParams{
+			Rule:   extractPublicKeyRule,
+			Input:  key,
+			Output: extractedPublicKey,
+		})
+	}
+
+	var ril int
+	if f.properties.Rollback_index_location != nil {
+		ril = proptools.Int(f.properties.Rollback_index_location)
+	}
+
+	android.SetProvider(ctx, vbmetaPartitionProvider, vbmetaPartitionInfo{
+		Name:                  f.partitionName(),
+		RollbackIndexLocation: ril,
+		PublicKey:             extractedPublicKey,
+		Output:                f.output,
+	})
 }
 
 func (f *filesystem) getMapFile(ctx android.ModuleContext) android.WritablePath {
diff --git a/filesystem/vbmeta.go b/filesystem/vbmeta.go
index b7f312c..91826b2 100644
--- a/filesystem/vbmeta.go
+++ b/filesystem/vbmeta.go
@@ -148,8 +148,8 @@
 var vbmetaChainedPartitionDep = chainedPartitionDep{}
 
 func (v *vbmeta) DepsMutator(ctx android.BottomUpMutatorContext) {
-	ctx.AddDependency(ctx.Module(), vbmetaPartitionDep, v.properties.Partitions.GetOrDefault(ctx, nil)...)
-	ctx.AddDependency(ctx.Module(), vbmetaChainedPartitionDep, v.properties.Chained_partitions...)
+	ctx.AddVariationDependencies(ctx.Config().AndroidFirstDeviceTarget.Variations(), vbmetaPartitionDep, v.properties.Partitions.GetOrDefault(ctx, nil)...)
+	ctx.AddVariationDependencies(ctx.Config().AndroidFirstDeviceTarget.Variations(), vbmetaChainedPartitionDep, v.properties.Chained_partitions...)
 }
 
 func (v *vbmeta) installFileName() string {