make godelegator extend writerdelegator<gowriter> and add writeAdditi…

…onalFiles(context) (#520)

codegen/smithy-go-codegen/src/main/java/software/amazon/smithy/go/codegen/CodegenVisitor.java

@@ -32,7 +32,6 @@
 import software.amazon.smithy.codegen.core.Symbol;
 import software.amazon.smithy.codegen.core.SymbolDependency;
 import software.amazon.smithy.codegen.core.SymbolProvider;
-import software.amazon.smithy.codegen.core.WriterDelegator;
 import software.amazon.smithy.go.codegen.GoSettings.ArtifactType;
 import software.amazon.smithy.go.codegen.integration.GoIntegration;
 import software.amazon.smithy.go.codegen.integration.ProtocolGenerator;
@@ -157,14 +156,7 @@ final class CodegenVisitor extends ShapeVisitor.Default<Void> {
 
         this.eventStreamGenerator = new EventStreamGenerator(settings, model, writers, symbolProvider, service);
 
-        this.ctx = new GoCodegenContext(
-                model,
-                settings,
-                symbolProvider,
-                fileManifest,
-                // FUTURE: GoDelegator should satisfy this interface
-                new WriterDelegator<>(fileManifest, symbolProvider, (filename, namespace) -> new GoWriter(namespace)),
-                integrations);
+        this.ctx = new GoCodegenContext(model, settings, symbolProvider, fileManifest, writers, integrations);
     }
 
     private static ProtocolGenerator resolveProtocolGenerator(
@@ -220,6 +212,7 @@ void execute() {
         for (GoIntegration integration : integrations) {
             integration.writeAdditionalFiles(settings, model, symbolProvider, writers::useFileWriter);
             integration.writeAdditionalFiles(settings, model, symbolProvider, writers);
+            integration.writeAdditionalFiles(ctx);
         }
 
         eventStreamGenerator.generateEventStreamInterfaces();

codegen/smithy-go-codegen/src/main/java/software/amazon/smithy/go/codegen/GoDelegator.java

@@ -17,9 +17,8 @@
 
 import java.util.function.Consumer;
 import software.amazon.smithy.build.FileManifest;
-import software.amazon.smithy.codegen.core.Symbol;
 import software.amazon.smithy.codegen.core.SymbolProvider;
-import software.amazon.smithy.codegen.core.SymbolReference;
+import software.amazon.smithy.codegen.core.WriterDelegator;
 import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.utils.SmithyInternalApi;
 
@@ -28,85 +27,43 @@
  * for getting shape specific GoWriters.
  */
 @SmithyInternalApi
-public final class GoDelegator extends GoWriterDelegator {
+public final class GoDelegator extends WriterDelegator<GoWriter> {
     private final SymbolProvider symbolProvider;
 
     public GoDelegator(FileManifest fileManifest, SymbolProvider symbolProvider) {
-        super(fileManifest);
+        super(fileManifest, symbolProvider, (filename, namespace) -> new GoWriter(namespace));
 
         this.symbolProvider = symbolProvider;
     }
 
     /**
-     * Gets a previously created writer or creates a new one if needed.
-     *
-     * @param shape          Shape to create the writer for.
-     * @param writerConsumer Consumer that accepts and works with the file.
-     */
-    public void useShapeWriter(Shape shape, Consumer<GoWriter> writerConsumer) {
-        Symbol symbol = symbolProvider.toSymbol(shape);
-        useShapeWriter(symbol, writerConsumer);
-    }
-
-    /**
-     * Gets a previously created writer or creates a new one for the a Go test file for the associated shape.
+     * Gets a previously created writer or creates a new one for the Go test file for the associated shape.
      *
      * @param shape          Shape to create the writer for.
      * @param writerConsumer Consumer that accepts and works with the file.
      */
     public void useShapeTestWriter(Shape shape, Consumer<GoWriter> writerConsumer) {
-        Symbol symbol = symbolProvider.toSymbol(shape);
-        String filename = symbol.getDefinitionFile();
-
-        StringBuilder b = new StringBuilder(filename);
-        b.insert(filename.lastIndexOf(".go"), "_test");
-        filename = b.toString();
-
-        symbol = symbol.toBuilder()
-                .definitionFile(filename)
-                .build();
-
-        useShapeWriter(symbol, writerConsumer);
+        var symbol = symbolProvider.toSymbol(shape);
+        var filename = symbol.getDefinitionFile();
+        var testFilename = new StringBuilder(filename)
+                .insert(filename.lastIndexOf(".go"), "_test")
+                .toString();
+        useFileWriter(testFilename, symbol.getNamespace(), writerConsumer);
     }
 
     /**
-     * Gets a previously created writer or creates a new one for the a Go public package test file for the associated
+     * Gets a previously created writer or creates a new one for the Go public package test file for the associated
      * shape.
      *
      * @param shape          Shape to create the writer for.
      * @param writerConsumer Consumer that accepts and works with the file.
      */
     public void useShapeExportedTestWriter(Shape shape, Consumer<GoWriter> writerConsumer) {
-        Symbol symbol = symbolProvider.toSymbol(shape);
-        String filename = symbol.getDefinitionFile();
-
-        StringBuilder b = new StringBuilder(filename);
-        b.insert(filename.lastIndexOf(".go"), "_exported_test");
-        filename = b.toString();
-
-        symbol = symbol.toBuilder()
-                .definitionFile(filename)
-                .namespace(symbol.getNamespace() + "_test", symbol.getNamespaceDelimiter())
-                .build();
-
-        useShapeWriter(symbol, writerConsumer);
-    }
-
-    /**
-     * Gets a previously created writer or creates a new one if needed.
-     *
-     * @param symbol         symbol to create the writer for.
-     * @param writerConsumer Consumer that accepts and works with the file.
-     */
-    private void useShapeWriter(Symbol symbol, Consumer<GoWriter> writerConsumer) {
-        GoWriter writer = checkoutWriter(symbol.getDefinitionFile(), symbol.getNamespace());
-
-        // Add any needed DECLARE symbols.
-        writer.addImportReferences(symbol, SymbolReference.ContextOption.DECLARE);
-        symbol.getDependencies().forEach(writer::addDependency);
-
-        writer.pushState();
-        writerConsumer.accept(writer);
-        writer.popState();
+        var symbol = symbolProvider.toSymbol(shape);
+        var filename = symbol.getDefinitionFile();
+        var testFilename = new StringBuilder(filename)
+                .insert(filename.lastIndexOf(".go"), "_exported_test")
+                .toString();
+        useFileWriter(testFilename, symbol.getNamespace() + "_test", writerConsumer);
     }
 }

codegen/smithy-go-codegen/src/main/java/software/amazon/smithy/go/codegen/integration/GoIntegration.java

@@ -127,6 +127,7 @@ default void onShapeWriterUse(
      * @param writerFactory A factory function that takes the name of a file
      *   to write and a {@code Consumer} that receives a
      *   {@link GoSettings} to perform the actual writing to the file.
+     * @deprecated use {@link #writeAdditionalFiles(GoCodegenContext)}.
      */
     default void writeAdditionalFiles(
             GoSettings settings,
@@ -144,6 +145,7 @@ default void writeAdditionalFiles(
      * @param model Model to generate from.
      * @param symbolProvider Symbol provider used for codegen.
      * @param goDelegator GoDelegator used to manage writer for the file.
+     * @deprecated use {@link #writeAdditionalFiles(GoCodegenContext)}.
      */
     default void writeAdditionalFiles(
             GoSettings settings,
@@ -154,6 +156,14 @@ default void writeAdditionalFiles(
         // pass
     }
 
+    /**
+     * Writes additional files.
+     * @param ctx The codegen context.
+     */
+    default void writeAdditionalFiles(GoCodegenContext ctx) {
+        // pass
+    }
+
     /**
      * Gets a list of protocol generators to register.
      *

codegen/smithy-go-codegen/src/main/java/software/amazon/smithy/go/codegen/integration/Waiters.java

@@ -26,8 +26,7 @@
 import software.amazon.smithy.codegen.core.Symbol;
 import software.amazon.smithy.codegen.core.SymbolProvider;
 import software.amazon.smithy.go.codegen.ClientOptions;
-import software.amazon.smithy.go.codegen.GoDelegator;
-import software.amazon.smithy.go.codegen.GoSettings;
+import software.amazon.smithy.go.codegen.GoCodegenContext;
 import software.amazon.smithy.go.codegen.GoWriter;
 import software.amazon.smithy.go.codegen.SmithyGoDependency;
 import software.amazon.smithy.go.codegen.SymbolUtils;
@@ -60,58 +59,36 @@ public Set<Symbol> getAdditionalClientOptions() {
     }
 
     @Override
-    public void writeAdditionalFiles(
-            GoSettings settings,
-            Model model,
-            SymbolProvider symbolProvider,
-            GoDelegator goDelegator
-    ) {
-        ServiceShape serviceShape = settings.getService(model);
-        TopDownIndex topDownIndex = TopDownIndex.of(model);
+    public void writeAdditionalFiles(GoCodegenContext ctx) {
+        var service = ctx.settings().getService(ctx.model());
 
-        topDownIndex.getContainedOperations(serviceShape).stream()
+        TopDownIndex.of(ctx.model()).getContainedOperations(service).stream()
                 .forEach(operation -> {
                     if (!operation.hasTrait(WaitableTrait.ID)) {
                         return;
                     }
 
                     Map<String, Waiter> waiters = operation.expectTrait(WaitableTrait.class).getWaiters();
-
-                    goDelegator.useShapeWriter(operation, writer -> {
-                        generateOperationWaiter(model, symbolProvider, writer, serviceShape, operation, waiters);
-                    });
+                    generateOperationWaiter(ctx, operation, waiters);
                 });
     }
 
 
     /**
      * Generates all waiter components used for the operation.
      */
-    private void generateOperationWaiter(
-            Model model,
-            SymbolProvider symbolProvider,
-            GoWriter writer,
-            ServiceShape service,
-            OperationShape operation,
-            Map<String, Waiter> waiters
-    ) {
-        // generate waiter function
-        waiters.forEach((name, waiter) -> {
-            // write waiter options
-            generateWaiterOptions(model, symbolProvider, writer, operation, name, waiter);
-
-            // write waiter client
-            generateWaiterClient(model, symbolProvider, writer, operation, name, waiter);
-
-            // write waiter specific invoker
-            generateWaiterInvoker(model, symbolProvider, writer, operation, name, waiter);
-
-            // write waiter specific invoker with output
-            generateWaiterInvokerWithOutput(model, symbolProvider, writer, operation, name, waiter);
-
-            // write waiter state mutator for each waiter
-            generateRetryable(model, symbolProvider, writer, service, operation, name, waiter);
-
+    private void generateOperationWaiter(GoCodegenContext ctx, OperationShape operation, Map<String, Waiter> waiters) {
+        var model = ctx.model();
+        var symbolProvider = ctx.symbolProvider();
+        var service = ctx.settings().getService(model);
+        ctx.writerDelegator().useShapeWriter(operation, writer -> {
+            waiters.forEach((name, waiter) -> {
+                generateWaiterOptions(model, symbolProvider, writer, operation, name, waiter);
+                generateWaiterClient(model, symbolProvider, writer, operation, name, waiter);
+                generateWaiterInvoker(model, symbolProvider, writer, operation, name, waiter);
+                generateWaiterInvokerWithOutput(model, symbolProvider, writer, operation, name, waiter);
+                generateRetryable(model, symbolProvider, writer, service, operation, name, waiter);
+            });
         });
     }
 

codegen/smithy-go-codegen/src/test/java/software/amazon/smithy/go/codegen/GenerateStandaloneGoModuleTest.java

@@ -25,8 +25,9 @@ public void testGenerateGoModule() throws Exception {
         var testPath = Files.createTempDirectory(getClass().getName());
         LOGGER.warning("generating test suites into " + testPath);
 
+        // symbol provider not needed here since only useFileWriter is called
         var fileManifest = FileManifest.create(testPath);
-        var writers = new GoWriterDelegator(fileManifest);
+        var writers = new GoDelegator(fileManifest, null);
 
         writers.useFileWriter("test-directory/package-name/gofile.go",
                 "github.com/aws/smithy-go/internal/testmodule/packagename",