rust: Add support for multiple protos per module.

This further emulates the rust-protobuf gen_mod_rs flag by providing
support for generating a single module containing multiple protobuf
definitions.

Bug: 171361369
Test: New Soong tests.
Test: Example module containing multiple protos works.
Change-Id: I815f9628a8289ae512758073dac49bc4535abf01

5 files changed, 73 insertions, 35 deletions

diff --git a/rust/library.go b/rust/library.go
index ae33f0fdd..971588d8a 100644
--- a/rust/library.go
+++ b/rust/library.go
@@ -428,6 +428,7 @@ func (library *libraryDecorator) compile(ctx ModuleContext, flags Flags, deps Pa
 	var srcPath android.Path
 
 	if library.sourceProvider != nil {
+		// Assume the first source from the source provider is the library entry point.
 		srcPath = library.sourceProvider.Srcs()[0]
 	} else {
 		srcPath, _ = srcPathFromModuleSrcs(ctx, library.baseCompiler.Properties.Srcs)
diff --git a/rust/protobuf.go b/rust/protobuf.go
index 7d6e1fd5c..235b4ad2b 100644
--- a/rust/protobuf.go
+++ b/rust/protobuf.go
@@ -46,8 +46,8 @@ func init() {
 var _ SourceProvider = (*protobufDecorator)(nil)
 
 type ProtobufProperties struct {
-	// Path to the proto file that will be used to generate the source
-	Proto *string `android:"path,arch_variant"`
+	// List of realtive paths to proto files that will be used to generate the source
+	Protos []string `android:"path,arch_variant"`
 
 	// List of additional flags to pass to aprotoc
 	Proto_flags []string `android:"arch_variant"`
@@ -66,6 +66,7 @@ type protobufDecorator struct {
 func (proto *protobufDecorator) GenerateSource(ctx ModuleContext, deps PathDeps) android.Path {
 	var protoFlags android.ProtoFlags
 	var pluginPaths android.Paths
+	var protoNames []string
 
 	protoFlags.OutTypeFlag = "--rust_out"
 	outDir := android.PathForModuleOut(ctx)
@@ -77,10 +78,7 @@ func (proto *protobufDecorator) GenerateSource(ctx ModuleContext, deps PathDeps)
 
 	protoFlags.Deps = append(protoFlags.Deps, pluginPaths...)
 
-	protoFile := android.OptionalPathForModuleSrc(ctx, proto.Properties.Proto)
-	if !protoFile.Valid() {
-		ctx.PropertyErrorf("proto", "invalid path to proto file")
-	}
+	protoFiles := android.PathsForModuleSrc(ctx, proto.Properties.Protos)
 
 	// Add exported dependency include paths
 	for _, include := range deps.depIncludePaths {
@@ -88,35 +86,58 @@ func (proto *protobufDecorator) GenerateSource(ctx ModuleContext, deps PathDeps)
 	}
 
 	stem := proto.BaseSourceProvider.getStem(ctx)
-	// rust protobuf-codegen output <stem>.rs
-	stemFile := android.PathForModuleOut(ctx, stem+".rs")
-	// add mod_<stem>.rs to import <stem>.rs
-	modFile := android.PathForModuleOut(ctx, "mod_"+stem+".rs")
-	// mod_<stem>.rs is the main/first output file to be included/compiled
-	outputs := android.WritablePaths{modFile, stemFile}
-	if proto.plugin == Grpc {
-		outputs = append(outputs, android.PathForModuleOut(ctx, stem+grpcSuffix+".rs"))
-	}
-	depFile := android.PathForModuleOut(ctx, "mod_"+stem+".d")
+
+	// The mod_stem.rs file is used to avoid collisions if this is not included as a crate.
+	stemFile := android.PathForModuleOut(ctx, "mod_"+stem+".rs")
+
+	// stemFile must be first here as the first path in BaseSourceProvider.OutputFiles is the library entry-point.
+	outputs := android.WritablePaths{stemFile}
 
 	rule := android.NewRuleBuilder()
-	android.ProtoRule(ctx, rule, protoFile.Path(), protoFlags, protoFlags.Deps, outDir, depFile, outputs)
-	rule.Command().Text("printf '" + proto.getModFileContents(ctx) + "' >").Output(modFile)
-	rule.Build(pctx, ctx, "protoc_"+protoFile.Path().Rel(), "protoc "+protoFile.Path().Rel())
+	for _, protoFile := range protoFiles {
+		protoName := strings.TrimSuffix(protoFile.Base(), ".proto")
+		protoNames = append(protoNames, protoName)
+
+		protoOut := android.PathForModuleOut(ctx, protoName+".rs")
+		ruleOutputs := android.WritablePaths{android.WritablePath(protoOut)}
+
+		if proto.plugin == Grpc {
+			grpcOut := android.PathForModuleOut(ctx, protoName+grpcSuffix+".rs")
+			ruleOutputs = append(ruleOutputs, android.WritablePath(grpcOut))
+		}
+
+		depFile := android.PathForModuleOut(ctx, protoName+".d")
 
-	proto.BaseSourceProvider.OutputFiles = android.Paths{modFile, stemFile}
-	return modFile
+		android.ProtoRule(ctx, rule, protoFile, protoFlags, protoFlags.Deps, outDir, depFile, ruleOutputs)
+		outputs = append(outputs, ruleOutputs...)
+	}
+
+	rule.Command().
+		Implicits(outputs.Paths()).
+		Text("printf '" + proto.genModFileContents(ctx, protoNames) + "' >").
+		Output(stemFile)
+
+	rule.Build(pctx, ctx, "protoc_"+ctx.ModuleName(), "protoc "+ctx.ModuleName())
+
+	proto.BaseSourceProvider.OutputFiles = outputs.Paths()
+
+	// mod_stem.rs is the entry-point for our library modules, so this is what we return.
+	return stemFile
 }
 
-func (proto *protobufDecorator) getModFileContents(ctx ModuleContext) string {
-	stem := proto.BaseSourceProvider.getStem(ctx)
+func (proto *protobufDecorator) genModFileContents(ctx ModuleContext, protoNames []string) string {
 	lines := []string{
-		"// @generated",
-		fmt.Sprintf("pub mod %s;", stem),
+		"// @Soong generated Source",
+	}
+	for _, protoName := range protoNames {
+		lines = append(lines, fmt.Sprintf("pub mod %s;", protoName))
+
+		if proto.plugin == Grpc {
+			lines = append(lines, fmt.Sprintf("pub mod %s%s;", protoName, grpcSuffix))
+		}
 	}
 
 	if proto.plugin == Grpc {
-		lines = append(lines, fmt.Sprintf("pub mod %s%s;", stem, grpcSuffix))
 		lines = append(
 			lines,
 			"pub mod empty {",
diff --git a/rust/protobuf_test.go b/rust/protobuf_test.go
index 845911f75..608a4e819 100644
--- a/rust/protobuf_test.go
+++ b/rust/protobuf_test.go
@@ -25,7 +25,7 @@ func TestRustProtobuf(t *testing.T) {
 	ctx := testRust(t, `
 		rust_protobuf {
 			name: "librust_proto",
-			proto: "buf.proto",
+			protos: ["buf.proto", "proto.proto"],
 			crate_name: "rust_proto",
 			source_stem: "buf",
 			shared_libs: ["libfoo_shared"],
@@ -60,13 +60,20 @@ func TestRustProtobuf(t *testing.T) {
 	if w := "-Istatic_include"; !strings.Contains(cmd, w) {
 		t.Errorf("expected %q in %q", w, cmd)
 	}
+
+	// Check proto.rs, the second protobuf, is listed as an output
+	librust_proto_outputs := ctx.ModuleForTests("librust_proto", "android_arm64_armv8-a_source").AllOutputs()
+	if android.InList("proto.rs", librust_proto_outputs) {
+		t.Errorf("rust_protobuf is not producing multiple outputs; expected 'proto.rs' in list, got: %#v ",
+			librust_proto_outputs)
+	}
 }
 
 func TestRustGrpcio(t *testing.T) {
 	ctx := testRust(t, `
 		rust_grpcio {
 			name: "librust_grpcio",
-			proto: "buf.proto",
+			protos: ["buf.proto", "proto.proto"],
 			crate_name: "rust_grpcio",
 			source_stem: "buf",
 			shared_libs: ["libfoo_shared"],
@@ -117,4 +124,11 @@ func TestRustGrpcio(t *testing.T) {
 	if w := "-Ilibprotobuf-cpp-full-includes"; !strings.Contains(cmd, w) {
 		t.Errorf("expected %q in %q", w, cmd)
 	}
+
+	// Check proto.rs, the second protobuf, is listed as an output
+	librust_grpcio_outputs := ctx.ModuleForTests("librust_grpcio", "android_arm64_armv8-a_source").AllOutputs()
+	if android.InList("proto_grpc.rs", librust_grpcio_outputs) {
+		t.Errorf("rust_protobuf is not producing multiple outputs; expected 'proto_grpc.rs' in list, got: %#v ",
+			librust_grpcio_outputs)
+	}
 }
diff --git a/rust/rust_test.go b/rust/rust_test.go
index 14bbd0b9a..81ad5ce77 100644
--- a/rust/rust_test.go
+++ b/rust/rust_test.go
@@ -106,13 +106,14 @@ func newTestRustCtx(t *testing.T, bp string) *testRustCtx {
 // useMockedFs setup a default mocked filesystem for the test environment.
 func (tctx *testRustCtx) useMockedFs() {
 	tctx.fs = map[string][]byte{
-		"foo.rs":     nil,
-		"foo.c":      nil,
-		"src/bar.rs": nil,
-		"src/any.h":  nil,
-		"buf.proto":  nil,
-		"liby.so":    nil,
-		"libz.so":    nil,
+		"foo.rs":      nil,
+		"foo.c":       nil,
+		"src/bar.rs":  nil,
+		"src/any.h":   nil,
+		"proto.proto": nil,
+		"buf.proto":   nil,
+		"liby.so":     nil,
+		"libz.so":     nil,
 	}
 }
 
diff --git a/rust/source_provider.go b/rust/source_provider.go
index 436518c70..7719611d5 100644
--- a/rust/source_provider.go
+++ b/rust/source_provider.go
@@ -30,6 +30,7 @@ type SourceProviderProperties struct {
 type BaseSourceProvider struct {
 	Properties SourceProviderProperties
 
+	// The first file in OutputFiles must be the library entry point.
 	OutputFiles      android.Paths
 	subAndroidMkOnce map[SubAndroidMkProvider]bool
 	subName          string