From d00d08a85f6f97478e5b62ee78eca25b98425666 Mon Sep 17 00:00:00 2001
From: Matt Morrissette <yinzara@gmail.com>
Date: Wed, 1 Apr 2020 15:24:59 -0700
Subject: [PATCH] Add "CodeGeneratorRequest" access in BuildContext

---
 ast.go           |  7 +++++++
 build_context.go | 16 ++++++++++++++--
 workflow.go      |  2 +-
 3 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/ast.go b/ast.go
index 73035d2..a41fa67 100644
--- a/ast.go
+++ b/ast.go
@@ -21,10 +21,14 @@ type AST interface {
 	// (FQN). The FQN uses dot notation of the form ".{package}.{entity}", or the
 	// input path for Files.
 	Lookup(name string) (Entity, bool)
+
+	// The original CodeGeneratorRequest from protoc
+    CodeGeneratorRequest() *plugin_go.CodeGeneratorRequest
 }
 
 type graph struct {
 	d Debugger
+    req *plugin_go.CodeGeneratorRequest
 
 	targets    map[string]File
 	packages   map[string]Package
@@ -41,6 +45,8 @@ func (g *graph) Lookup(name string) (Entity, bool) {
 	return e, ok
 }
 
+func (g *graph) CodeGeneratorRequest() *plugin_go.CodeGeneratorRequest { return g.req }
+
 // ProcessDescriptors is deprecated; use ProcessCodeGeneratorRequest instead
 func ProcessDescriptors(debug Debugger, req *plugin_go.CodeGeneratorRequest) AST {
 	return ProcessCodeGeneratorRequest(debug, req)
@@ -51,6 +57,7 @@ func ProcessDescriptors(debug Debugger, req *plugin_go.CodeGeneratorRequest) AST
 func ProcessCodeGeneratorRequest(debug Debugger, req *plugin_go.CodeGeneratorRequest) AST {
 	g := &graph{
 		d:          debug,
+		req:        req,
 		targets:    make(map[string]File, len(req.GetFileToGenerate())),
 		packages:   make(map[string]Package),
 		entities:   make(map[string]Entity),
diff --git a/build_context.go b/build_context.go
index b48d0b7..fbdb90d 100644
--- a/build_context.go
+++ b/build_context.go
@@ -1,6 +1,9 @@
 package pgs
 
-import "path/filepath"
+import (
+    plugin_go "github.com/golang/protobuf/protoc-gen-go/plugin"
+    "path/filepath"
+)
 
 // BuildContext tracks code generation relative to an output path. By default,
 // BuildContext's path is relative to the output location specified when
@@ -40,17 +43,21 @@ type BuildContext interface {
 	// Parameters returns the command line parameters passed in from protoc,
 	// mutated with any provided ParamMutators via InitOptions.
 	Parameters() Parameters
+
+    // The original CodeGeneratorRequest from protoc
+    CodeGeneratorRequest() *plugin_go.CodeGeneratorRequest
 }
 
 // Context creates a new BuildContext with the provided debugger and initial
 // output path. For protoc-gen-go plugins, output is typically ".", while
 // Module's may use a custom path.
-func Context(d Debugger, params Parameters, output string) BuildContext {
+func Context(d Debugger, params Parameters, output string, req *plugin_go.CodeGeneratorRequest) BuildContext {
 	return rootContext{
 		dirContext: dirContext{
 			prefixContext: prefixContext{parent: nil, d: d},
 			p:             filepath.Clean(output),
 		},
+		req:    req,
 		params: params,
 	}
 }
@@ -80,6 +87,8 @@ func (c prefixContext) Push(prefix string) BuildContext { return initPrefixConte
 func (c prefixContext) Pop() BuildContext               { return c.parent }
 func (c prefixContext) PopDir() BuildContext            { return c.parent.PopDir() }
 
+func (c prefixContext) CodeGeneratorRequest() *plugin_go.CodeGeneratorRequest { return c.parent.CodeGeneratorRequest() }
+
 type dirContext struct {
 	prefixContext
 	p string
@@ -117,6 +126,7 @@ type prefixContext struct {
 type rootContext struct {
 	dirContext
 	params Parameters
+    req *plugin_go.CodeGeneratorRequest
 }
 
 func (c rootContext) OutputPath() string              { return c.p }
@@ -129,6 +139,8 @@ func (c rootContext) Pop() BuildContext {
 	return nil
 }
 
+func (c rootContext) CodeGeneratorRequest() *plugin_go.CodeGeneratorRequest { return c.req }
+
 func (c rootContext) JoinPath(name ...string) string {
 	return filepath.Join(append([]string{c.OutputPath()}, name...)...)
 }
diff --git a/workflow.go b/workflow.go
index 1744426..4138496 100644
--- a/workflow.go
+++ b/workflow.go
@@ -49,7 +49,7 @@ func (wf *standardWorkflow) Init(g *Generator) AST {
 }
 
 func (wf *standardWorkflow) Run(ast AST) (arts []Artifact) {
-	ctx := Context(wf.Debugger, wf.params, wf.params.OutputPath())
+	ctx := Context(wf.Debugger, wf.params, wf.params.OutputPath(), ast.CodeGeneratorRequest())
 
 	wf.Debug("initializing modules")
 	for _, m := range wf.mods {