Add shapes generation ordering in CodegenDirector

smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/ShapesGenerationOrdering.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.smithy.codegen.core;
+
+/**
+ * Shapes ordering for code generation.
+ *
+ * <p>CodegenDirector order the shapes appropriately before feeding them to the code generators. See {@link
+ * software.amazon.smithy.codegen.core.directed.CodegenDirector#setShapesGenerationOrdering(ShapesGenerationOrdering)}
+ */
+public enum ShapesGenerationOrdering {
+    /**
+     * Shapes ordered in reverse-topological ordering. Also see {@link TopologicalIndex}
+     */
+    TOPOLOGICAL,
+
+    /**
+     * Shapes ordered alphabetically by their names.
+     */
+    ALPHABETICAL,
+
+    /**
+     * Shapes without ordering.
+     */
+    NONE
+}

smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/CodegenDirector.java

@@ -16,6 +16,7 @@
 package software.amazon.smithy.codegen.core.directed;
 
 import java.util.ArrayList;
+import java.util.Comparator;
 import java.util.List;
 import java.util.Objects;
 import java.util.ServiceLoader;
@@ -26,6 +27,7 @@
 import software.amazon.smithy.build.FileManifest;
 import software.amazon.smithy.codegen.core.CodegenContext;
 import software.amazon.smithy.codegen.core.ImportContainer;
+import software.amazon.smithy.codegen.core.ShapesGenerationOrdering;
 import software.amazon.smithy.codegen.core.SmithyIntegration;
 import software.amazon.smithy.codegen.core.SymbolProvider;
 import software.amazon.smithy.codegen.core.SymbolWriter;
@@ -75,6 +77,7 @@ public final class CodegenDirector<
     private Supplier<Iterable<I>> integrationFinder;
     private DirectedCodegen<C, S, I> directedCodegen;
     private final List<BiFunction<Model, ModelTransformer, Model>> transforms = new ArrayList<>();
+    private ShapesGenerationOrdering shapesGenerationOrdering = ShapesGenerationOrdering.TOPOLOGICAL;
 
     /**
      * Simplifies a Smithy model for code generation of a single service.
@@ -250,6 +253,22 @@ public void changeStringEnumsToEnumShapes(boolean synthesizeEnumNames) {
         });
     }
 
+    /**
+     * Sets the shapes ordering for code generation.
+     *
+     * <p>CodegenDirector order the shapes appropriately before passing them to the code generators.
+     * The default ordering is topological, and can be overridden with this method
+     *
+     * @param ordering the ordering to use for the code generation process.
+     */
+    public void setShapesGenerationOrdering(ShapesGenerationOrdering ordering) {
+        if (Objects.isNull(ordering)) {
+            this.shapesGenerationOrdering = ShapesGenerationOrdering.NONE;
+        } else {
+            this.shapesGenerationOrdering = ordering;
+        }
+    }
+
     /**
      * Sorts all members of the model prior to codegen.
      *
@@ -390,16 +409,36 @@ private void registerInterceptors(C context, List<I> integrations) {
     }
 
     private void generateShapesInService(C context, ServiceShape serviceShape) {
-        LOGGER.fine(() -> "Generating shapes for " + directedCodegen.getClass().getName());
+        LOGGER.fine(() -> String.format("Generating shapes for %s in %s ordering",
+                directedCodegen.getClass().getName(), this.shapesGenerationOrdering.name()));
         Set<Shape> shapes = new Walker(context.model()).walkShapes(serviceShape);
-        TopologicalIndex topologicalIndex = TopologicalIndex.of(context.model());
         ShapeGenerator<W, C, S> generator = new ShapeGenerator<>(context, serviceShape, directedCodegen);
-        for (Shape shape : topologicalIndex.getOrderedShapes()) {
-            if (shapes.contains(shape)) {
-                shape.accept(generator);
-            }
+        List<Shape> orderedShapes = new ArrayList<>();
+
+        switch (this.shapesGenerationOrdering) {
+            case ALPHABETICAL:
+                orderedShapes.addAll(shapes);
+                orderedShapes.sort(Comparator.comparing(s -> s.getId().getName(serviceShape)));
+                break;
+            case NONE:
+                orderedShapes.addAll(shapes);
+                break;
+            case TOPOLOGICAL:
+            default:
+                TopologicalIndex topologicalIndex = TopologicalIndex.of(context.model());
+                for (Shape shape : topologicalIndex.getOrderedShapes()) {
+                    if (shapes.contains(shape)) {
+                        orderedShapes.add(shape);
+                    }
+                }
+                for (Shape shape : topologicalIndex.getRecursiveShapes()) {
+                    if (shapes.contains(shape)) {
+                        orderedShapes.add(shape);
+                    }
+                }
         }
-        for (Shape shape : topologicalIndex.getRecursiveShapes()) {
+
+        for (Shape shape : orderedShapes) {
             if (shapes.contains(shape)) {
                 shape.accept(generator);
             }

smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/CodegenDirectorTest.java

@@ -26,6 +26,7 @@
 import org.junit.jupiter.api.Test;
 import software.amazon.smithy.build.FileManifest;
 import software.amazon.smithy.build.MockManifest;
+import software.amazon.smithy.codegen.core.ShapesGenerationOrdering;
 import software.amazon.smithy.codegen.core.Symbol;
 import software.amazon.smithy.codegen.core.SymbolProvider;
 import software.amazon.smithy.codegen.core.WriterDelegator;
@@ -227,7 +228,7 @@ public void performsCodegenWithStringEnumsChangedToEnumShapes() {
     }
 
     @Test
-    public void sortsShapes() {
+    public void sortsShapesWithDefaultTopologicalOrdering() {
         TestDirected testDirected = new TestDirected();
         CodegenDirector<TestWriter, TestIntegration, TestContext, TestSettings> runner
                 = new CodegenDirector<>();
@@ -260,4 +261,72 @@ public void sortsShapes() {
                 ShapeId.from("smithy.example#Foo")
         ));
     }
+
+    @Test
+    public void testShapesGenerationByAlphabeticalOrdering() {
+        TestDirected testDirected = new TestDirected();
+        CodegenDirector<TestWriter, TestIntegration, TestContext, TestSettings> runner
+                = new CodegenDirector<>();
+        FileManifest manifest = new MockManifest();
+        Model model = Model.assembler()
+                .addImport(getClass().getResource("needs-sorting.smithy"))
+                .assemble()
+                .unwrap();
+
+        runner.settings(new TestSettings());
+        runner.directedCodegen(testDirected);
+        runner.fileManifest(manifest);
+        runner.service(ShapeId.from("smithy.example#Foo"));
+        runner.model(model);
+        runner.integrationClass(TestIntegration.class);
+        runner.performDefaultCodegenTransforms();
+        runner.setShapesGenerationOrdering(ShapesGenerationOrdering.ALPHABETICAL);
+        runner.run();
+
+        assertThat(testDirected.generatedShapes, contains(
+                ShapeId.from("smithy.example#A"),
+                ShapeId.from("smithy.example#B"),
+                ShapeId.from("smithy.example#C"),
+                ShapeId.from("smithy.example#D"),
+                ShapeId.from("smithy.example#FooOperationInput"),
+                ShapeId.from("smithy.example#FooOperationOutput"),
+                ShapeId.from("smithy.example#RecursiveA"),
+                ShapeId.from("smithy.example#RecursiveB"),
+                ShapeId.from("smithy.example#Foo")
+        ));
+    }
+
+    @Test
+    public void testShapesGenerationWithoutOrdering() {
+        TestDirected testDirected = new TestDirected();
+        CodegenDirector<TestWriter, TestIntegration, TestContext, TestSettings> runner
+                = new CodegenDirector<>();
+        FileManifest manifest = new MockManifest();
+        Model model = Model.assembler()
+                .addImport(getClass().getResource("needs-sorting.smithy"))
+                .assemble()
+                .unwrap();
+
+        runner.settings(new TestSettings());
+        runner.directedCodegen(testDirected);
+        runner.fileManifest(manifest);
+        runner.service(ShapeId.from("smithy.example#Foo"));
+        runner.model(model);
+        runner.integrationClass(TestIntegration.class);
+        runner.performDefaultCodegenTransforms();
+        runner.setShapesGenerationOrdering(ShapesGenerationOrdering.NONE);
+        runner.run();
+
+        assertThat(testDirected.generatedShapes, contains(
+                ShapeId.from("smithy.example#FooOperationOutput"),
+                ShapeId.from("smithy.example#A"),
+                ShapeId.from("smithy.example#B"),
+                ShapeId.from("smithy.example#C"),
+                ShapeId.from("smithy.example#D"),
+                ShapeId.from("smithy.example#FooOperationInput"),
+                ShapeId.from("smithy.example#RecursiveA"),
+                ShapeId.from("smithy.example#RecursiveB"),
+                ShapeId.from("smithy.example#Foo")
+        ));
+    }
 }