Enable BiDi graph building (#70)

* Enable BiDi graph building

init_option.go

@@ -41,3 +41,9 @@ func MutateParams(pm ...ParamMutator) InitOption {
 // disk. By default, the OS's file system is used. This option currently only
 // impacts CustomFile and CustomTemplateFile artifacts generated by modules.
 func FileSystem(fs afero.Fs) InitOption { return func(g *Generator) { g.persister.SetFS(fs) } }
+
+// BiDirectional instructs the Generator to build the AST graph in both
+// directions (ie, accessing dependents of an entity, not just dependencies).
+func BiDirectional() InitOption {
+	return func(g *Generator) { g.workflow = &onceWorkflow{workflow: &standardWorkflow{BiDi: true}} }
+}

init_option_test.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/spf13/afero"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestDebugMode(t *testing.T) {
@@ -70,3 +71,21 @@ func TestProtocOutput(t *testing.T) {
 	ProtocOutput(b)(g)
 	assert.Equal(t, b, g.out)
 }
+
+func TestBiDirectional(t *testing.T) {
+	t.Parallel()
+
+	g := &Generator{}
+	assert.Nil(t, g.workflow)
+
+	BiDirectional()(g)
+	wf := g.workflow
+
+	require.IsType(t, &onceWorkflow{}, wf)
+	once := wf.(*onceWorkflow)
+
+	require.IsType(t, &standardWorkflow{}, once.workflow)
+	std := once.workflow.(*standardWorkflow)
+
+	assert.True(t, std.BiDi)
+}

workflow.go

@@ -15,9 +15,12 @@ type workflow interface {
 }
 
 // standardWorkflow describes a typical protoc-plugin flow, with the only
-// exception being the behavior of the persistor directly writing custom file
+// exception being the behavior of the persister directly writing custom file
 // artifacts to disk (instead of via the plugin's output to protoc).
-type standardWorkflow struct{ *Generator }
+type standardWorkflow struct {
+	*Generator
+	BiDi bool
+}
 
 func (wf *standardWorkflow) Init(g *Generator) AST {
 	wf.Generator = g
@@ -38,6 +41,10 @@ func (wf *standardWorkflow) Init(g *Generator) AST {
 		pm(wf.params)
 	}
 
+	if wf.BiDi {
+		return ProcessCodeGeneratorRequestBidirectional(g, req)
+	}
+
 	return ProcessCodeGeneratorRequest(g, req)
 }
 

workflow_test.go

@@ -23,6 +23,14 @@ func TestStandardWorkflow_Init(t *testing.T) {
 	g.workflow.Init(g)
 
 	assert.True(t, mutated)
+
+	t.Run("bidi", func(t *testing.T) {
+		mutated = false
+		g = Init(ProtocInput(bytes.NewReader(b)), BiDirectional(), MutateParams(func(p Parameters) { mutated = true }))
+		g.workflow.Init(g)
+
+		assert.True(t, mutated)
+	})
 }
 
 func TestStandardWorkflow_Run(t *testing.T) {