From 69b4b5d1c603a59b2ca56d22b3fdaac295d74f55 Mon Sep 17 00:00:00 2001
From: Chris Moran <chris.moran@blockfi.com>
Date: Wed, 23 Feb 2022 00:14:50 -0500
Subject: [PATCH] Support for `-micro_out=module=<module_prefix>` for
 protoc-gen-micro (#2435)

---
 cmd/protoc-gen-micro/generator/generator.go | 30 +++++++++++++++++----
 1 file changed, 25 insertions(+), 5 deletions(-)

diff --git a/cmd/protoc-gen-micro/generator/generator.go b/cmd/protoc-gen-micro/generator/generator.go
index 8530a54f..22015dc3 100644
--- a/cmd/protoc-gen-micro/generator/generator.go
+++ b/cmd/protoc-gen-micro/generator/generator.go
@@ -305,7 +305,7 @@ func (d *FileDescriptor) goPackageOption() (impPath GoImportPath, pkg GoPackageN
 }
 
 // goFileName returns the output name for the generated Go file.
-func (d *FileDescriptor) goFileName(pathType pathType) string {
+func (d *FileDescriptor) goFileName(pathType pathType, moduleRoot string) string {
 	name := *d.Name
 	if ext := path.Ext(name); ext == ".proto" || ext == ".protodevel" {
 		name = name[:len(name)-len(ext)]
@@ -319,9 +319,18 @@ func (d *FileDescriptor) goFileName(pathType pathType) string {
 	// Does the file have a "go_package" option?
 	// If it does, it may override the filename.
 	if impPath, _, ok := d.goPackageOption(); ok && impPath != "" {
-		// Replace the existing dirname with the declared import path.
-		_, name = path.Split(name)
-		name = path.Join(string(impPath), name)
+		if pathType == pathModuleRoot && moduleRoot != "" {
+			root := moduleRoot
+			if !strings.HasSuffix(root, "/") {
+				root = root + "/"
+			}
+			name = strings.TrimPrefix(name, root)
+		} else {
+			// Replace the existing dirname with the declared import path.
+			_, name = path.Split(name)
+			name = path.Join(string(impPath), name)
+		}
+
 		return name
 	}
 
@@ -405,6 +414,7 @@ type Generator struct {
 	PackageImportPath string            // Go import path of the package we're generating code for
 	ImportPrefix      string            // String to prefix to imported package file names.
 	ImportMap         map[string]string // Mapping from .proto file name to import path
+	ModuleRoot        string            // Mapping from the module prefix
 
 	Pkg map[string]string // The names under which we import support packages
 
@@ -429,6 +439,7 @@ type pathType int
 const (
 	pathTypeImport pathType = iota
 	pathTypeSourceRelative
+	pathModuleRoot
 )
 
 // New creates a new generator and allocates the request and response protobufs.
@@ -475,11 +486,20 @@ func (g *Generator) CommandLineParameters(parameter string) {
 			g.ImportPrefix = v
 		case "import_path":
 			g.PackageImportPath = v
+		case "module":
+			if g.pathType == pathTypeSourceRelative {
+				g.Fail(fmt.Sprintf(`Cannot set  module=%q after paths=source_relative`, v))
+			}
+			g.pathType = pathModuleRoot
+			g.ModuleRoot = v
 		case "paths":
 			switch v {
 			case "import":
 				g.pathType = pathTypeImport
 			case "source_relative":
+				if g.pathType == pathModuleRoot {
+					g.Fail("Cannot set paths=source_relative after setting module=<module_root>")
+				}
 				g.pathType = pathTypeSourceRelative
 			default:
 				g.Fail(fmt.Sprintf(`Unknown path type %q: want "import" or "source_relative".`, v))
@@ -1069,7 +1089,7 @@ func (g *Generator) GenerateAllFiles() {
 		if !g.writeOutput {
 			continue
 		}
-		fname := file.goFileName(g.pathType)
+		fname := file.goFileName(g.pathType, g.ModuleRoot)
 		g.Response.File = append(g.Response.File, &plugin.CodeGeneratorResponse_File{
 			Name:    proto.String(fname),
 			Content: proto.String(g.String()),