Do not use getModelWithoutTraitShapes (#547)

Author: gosar

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

@@ -50,7 +50,6 @@
 import software.amazon.smithy.model.shapes.UnionShape;
 import software.amazon.smithy.model.traits.EnumTrait;
 import software.amazon.smithy.model.traits.PaginatedTrait;
-import software.amazon.smithy.model.transform.ModelTransformer;
 import software.amazon.smithy.model.validation.ValidationEvent;
 import software.amazon.smithy.typescript.codegen.TypeScriptSettings.ArtifactType;
 import software.amazon.smithy.typescript.codegen.integration.ProtocolGenerator;
@@ -85,7 +84,6 @@ class CodegenVisitor extends ShapeVisitor.Default<Void> {
     private final ServiceShape service;
     private final FileManifest fileManifest;
     private final SymbolProvider symbolProvider;
-    private final Model nonTraits;
     private final TypeScriptDelegator writers;
     private final List<TypeScriptIntegration> integrations = new ArrayList<>();
     private final List<RuntimeClientPlugin> runtimePlugins = new ArrayList<>();
@@ -115,7 +113,6 @@ class CodegenVisitor extends ShapeVisitor.Default<Void> {
             modifiedModel = integration.preprocessModel(modifiedModel, settings);
         }
         model = modifiedModel;
-        nonTraits = ModelTransformer.create().getModelWithoutTraitShapes(model);
 
         service = settings.getService(model);
 
@@ -179,19 +176,18 @@ void execute() {
 
         // Generate models that are connected to the service being generated.
         LOGGER.fine("Walking shapes from " + service.getId() + " to find shapes to generate");
-        // Walk the tree.
-        Collection<Shape> shapeSet = new Walker(nonTraits).walkShapes(service);
-
-        Model prunedModel = Model.builder().addShapes(shapeSet).build();
-
-        // Generate models from shapes.
-        for (Shape shape : TopologicalIndex.of(prunedModel).getOrderedShapes()) {
-            shape.accept(this);
+        Collection<Shape> shapes = new Walker(model).walkShapes(service);
+        for (Shape shape : TopologicalIndex.of(model).getOrderedShapes()) {
+            if (shapes.contains(shape)) {
+                shape.accept(this);
+            }
         }
-        for (Shape shape : TopologicalIndex.of(prunedModel).getRecursiveShapes()) {
-            shape.accept(this);
+        for (Shape shape : TopologicalIndex.of(model).getRecursiveShapes()) {
+            if (shapes.contains(shape)) {
+                shape.accept(this);
+            }
         }
-        SymbolVisitor.writeModelIndex(prunedModel, symbolProvider, fileManifest);
+        SymbolVisitor.writeModelIndex(shapes, symbolProvider, fileManifest);
 
         // Generate the client Node and Browser configuration files. These
         // files are switched between in package.json based on the targeted

smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/SymbolVisitor.java

@@ -18,6 +18,7 @@
 import static java.lang.String.format;
 
 import java.nio.file.Paths;
+import java.util.Collection;
 import java.util.Comparator;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -118,8 +119,8 @@ final class SymbolVisitor implements SymbolProvider, ShapeVisitor<Symbol> {
         moduleNameDelegator = new ModuleNameDelegator(shapeChunkSize);
     }
 
-    static void writeModelIndex(Model model, SymbolProvider symbolProvider, FileManifest fileManifest) {
-        ModuleNameDelegator.writeModelIndex(model, symbolProvider, fileManifest);
+    static void writeModelIndex(Collection<Shape> shapes, SymbolProvider symbolProvider, FileManifest fileManifest) {
+        ModuleNameDelegator.writeModelIndex(shapes, symbolProvider, fileManifest);
     }
 
     @Override
@@ -442,19 +443,21 @@ public String formatModuleName(Shape shape, String name) {
             return path;
         }
 
-        static void writeModelIndex(Model model, SymbolProvider symbolProvider, FileManifest fileManifest) {
+        static void writeModelIndex(Collection<Shape> shapes, SymbolProvider symbolProvider,
+                                    FileManifest fileManifest) {
             TypeScriptWriter writer = new TypeScriptWriter("");
             String modelPrefix = Paths.get(".", CodegenUtils.SOURCE_FOLDER, SHAPE_NAMESPACE_PREFIX).toString();
-            model.shapes()
+            shapes.stream()
                     .map(shape -> symbolProvider.toSymbol(shape).getNamespace())
                     .filter(namespace -> namespace.startsWith(modelPrefix))
                     .distinct()
                     .sorted(Comparator.naturalOrder())
                     .forEach(namespace -> writer.write(
-                        "export * from $S;", namespace.replaceFirst(modelPrefix, ".")));
+                            "export * from $S;", namespace.replaceFirst(modelPrefix, ".")));
             fileManifest.writeFile(
-                Paths.get(CodegenUtils.SOURCE_FOLDER, SHAPE_NAMESPACE_PREFIX, "index.ts").toString(),
-                writer.toString());
+                    Paths.get(CodegenUtils.SOURCE_FOLDER, SHAPE_NAMESPACE_PREFIX, "index.ts").toString(),
+                    writer.toString());
         }
+
     }
 }

smithy-typescript-codegen/src/test/java/software/amazon/smithy/typescript/codegen/SymbolProviderTest.java

@@ -4,6 +4,7 @@
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
 
+import java.util.Arrays;
 import org.junit.jupiter.api.Test;
 import software.amazon.smithy.build.MockManifest;
 import software.amazon.smithy.codegen.core.Symbol;
@@ -57,7 +58,7 @@ public void createsSymbolsIntoTargetNamespace() {
         Symbol symbol1 = provider.toSymbol(shape1);
         Symbol symbol2 = provider.toSymbol(shape2);
         MockManifest manifest = new MockManifest();
-        SymbolVisitor.writeModelIndex(model, provider, manifest);
+        SymbolVisitor.writeModelIndex(Arrays.asList(shape1, shape2), provider, manifest);
 
         assertThat(symbol1.getName(), equalTo("Hello"));
         assertThat(symbol1.getNamespace(), equalTo("./" + CodegenUtils.SOURCE_FOLDER + "/models/models_0"));