toposort models and chunk big model files (#210)

This change exports shapes in topological order, and chunk them into models/models_0.ts, models/models_1.ts, models/models_2.ts..., models/models_x.ts. A models/index.ts will re-export shapes from each chunks.

Author: AllanZhengYP

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

@@ -34,6 +34,7 @@
 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.TopologicalIndex;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.knowledge.TopDownIndex;
 import software.amazon.smithy.model.neighbor.Walker;
@@ -152,15 +153,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");
-        Set<Shape> serviceShapes = new TreeSet<>(new Walker(nonTraits).walkShapes(service));
-
-        // Condense duplicate shapes
-        Map<String, Shape> shapeMap = condenseShapes(serviceShapes);
+        // Walk the tree and condense duplicate shapes
+        Collection<Shape> shapeSet = condenseShapes(new Walker(nonTraits).walkShapes(service));
+        Model prunedModel = Model.builder().addShapes(shapeSet).build();
 
         // Generate models from condensed shapes
-        for (Shape shape : shapeMap.values()) {
+        for (Shape shape : TopologicalIndex.of(prunedModel).getOrderedShapes()) {
+            shape.accept(this);
+        }
+        for (Shape shape : TopologicalIndex.of(prunedModel).getRecursiveShapes()) {
             shape.accept(this);
         }
+        SymbolVisitor.writeModelIndex(prunedModel, symbolProvider, fileManifest);
 
         // Generate the client Node and Browser configuration files. These
         // files are switched between in package.json based on the targeted
@@ -325,7 +329,7 @@ public Void serviceShape(ServiceShape shape) {
         return null;
     }
 
-    private Map<String, Shape> condenseShapes(Set<Shape> shapes) {
+    private Collection<Shape> condenseShapes(Set<Shape> shapes) {
         Map<String, Shape> shapeMap = new LinkedHashMap<>();
 
         // Check for colliding shapes and prune non-unique shapes
@@ -342,7 +346,7 @@ private Map<String, Shape> condenseShapes(Set<Shape> shapes) {
             }
         }
 
-        return shapeMap;
+        return shapeMap.values();
     }
 
     private boolean isShapeCollision(Shape shapeA, Shape shapeB) {

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

@@ -17,10 +17,15 @@
 
 import static java.lang.String.format;
 
+import java.util.Comparator;
+import java.util.HashMap;
 import java.util.HashSet;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.logging.Logger;
+
+import software.amazon.smithy.build.FileManifest;
 import software.amazon.smithy.codegen.core.CodegenException;
 import software.amazon.smithy.codegen.core.ReservedWordSymbolProvider;
 import software.amazon.smithy.codegen.core.ReservedWords;
@@ -75,8 +80,13 @@ final class SymbolVisitor implements SymbolProvider, ShapeVisitor<Symbol> {
     private final Model model;
     private final ReservedWordSymbolProvider.Escaper escaper;
     private final Set<StructureShape> errorShapes = new HashSet<>();
+    private final ModuleNameDelegator moduleNameDelegator;
 
     SymbolVisitor(Model model) {
+        this(model, ModuleNameDelegator.DEFAULT_CHUNK_SIZE);
+    }
+
+    SymbolVisitor(Model model, int shapeChunkSize) {
         this.model = model;
 
         // Load reserved words from a new-line delimited file.
@@ -96,6 +106,12 @@ final class SymbolVisitor implements SymbolProvider, ShapeVisitor<Symbol> {
         model.shapes(OperationShape.class).forEach(operationShape -> {
             errorShapes.addAll(operationIndex.getErrors(operationShape));
         });
+
+        moduleNameDelegator = new ModuleNameDelegator(shapeChunkSize);
+    }
+
+    static void writeModelIndex(Model model, SymbolProvider symbolProvider, FileManifest fileManifest) {
+        ModuleNameDelegator.writeModelIndex(model, symbolProvider, fileManifest);
     }
 
     @Override
@@ -225,7 +241,7 @@ public Symbol documentShape(DocumentShape shape) {
     @Override
     public Symbol operationShape(OperationShape shape) {
         String commandName = flattenShapeName(shape) + "Command";
-        String moduleName = formatModuleName(shape.getType(), commandName);
+        String moduleName = moduleNameDelegator.formatModuleName(shape, commandName);
         Symbol intermediate = createGeneratedSymbolBuilder(shape, commandName, moduleName).build();
         Symbol.Builder builder = intermediate.toBuilder();
         // Add input and output type symbols (XCommandInput / XCommandOutput).
@@ -271,7 +287,7 @@ public Symbol resourceShape(ResourceShape shape) {
     @Override
     public Symbol serviceShape(ServiceShape shape) {
         String name = StringUtils.capitalize(shape.getId().getName()) + "Client";
-        String moduleName = formatModuleName(shape.getType(), name);
+        String moduleName = moduleNameDelegator.formatModuleName(shape, name);
         return createGeneratedSymbolBuilder(shape, name, moduleName).build();
     }
 
@@ -358,7 +374,7 @@ private String flattenShapeName(ToShapeId id) {
 
     private Symbol.Builder createObjectSymbolBuilder(Shape shape) {
         String name = flattenShapeName(shape);
-        String moduleName = formatModuleName(shape.getType(), name);
+        String moduleName = moduleNameDelegator.formatModuleName(shape, name);
         return createGeneratedSymbolBuilder(shape, name, moduleName);
     }
 
@@ -378,18 +394,57 @@ private Symbol.Builder createGeneratedSymbolBuilder(Shape shape, String typeName
                 .definitionFile(toFilename(namespace));
     }
 
-    private String formatModuleName(ShapeType shapeType, String name) {
-        // All shapes except for the service and operations are stored in models.
-        if (shapeType == ShapeType.SERVICE) {
-            return "./" + name;
-        } else if (shapeType == ShapeType.OPERATION) {
-            return "./commands/" + name;
-        } else {
-            return "./models/index";
-        }
-    }
-
     private String toFilename(String namespace) {
         return namespace + ".ts";
     }
+
+    /**
+     * Utility class to locate which path should the symbol be generated into.
+     * It will break the models into multiple files to prevent it getting too big.
+     */
+    static final class ModuleNameDelegator {
+        static final int DEFAULT_CHUNK_SIZE = 300;
+        static final String SHAPE_NAMESPACE_PREFIX = "./models/";
+
+        private final Map<Shape, String> visitedModels = new HashMap<>();
+        private int bucketCount = 0;
+        private int currentBucketSize = 0;
+        private final int chunkSize;
+
+        ModuleNameDelegator(int shapeChunkSize) {
+            chunkSize = shapeChunkSize;
+        }
+
+        public String formatModuleName(Shape shape, String name) {
+            // All shapes except for the service and operations are stored in models.
+            if (shape.getType() == ShapeType.SERVICE) {
+                return "./" + name;
+            } else if (shape.getType() == ShapeType.OPERATION) {
+                return "./commands/" + name;
+            } else if (visitedModels.containsKey(shape)) {
+                return visitedModels.get(shape);
+            }
+            // Add models into buckets no bigger than chunk size.
+            String path = SHAPE_NAMESPACE_PREFIX + "models_" + bucketCount;
+            visitedModels.put(shape, path);
+            currentBucketSize++;
+            if (currentBucketSize == chunkSize) {
+                bucketCount++;
+                currentBucketSize = 0;
+            }
+            return path;
+        }
+
+        static void writeModelIndex(Model model, SymbolProvider symbolProvider, FileManifest fileManifest) {
+            TypeScriptWriter writer = new TypeScriptWriter("");
+            model.shapes()
+                    .map(shape -> symbolProvider.toSymbol(shape).getNamespace())
+                    .filter(namespace -> namespace.startsWith(SHAPE_NAMESPACE_PREFIX))
+                    .distinct()
+                    .sorted(Comparator.naturalOrder())
+                    .forEach(namespace -> writer.write(
+                        "export * from $S;", namespace.replaceFirst(SHAPE_NAMESPACE_PREFIX, "./")));
+            fileManifest.writeFile("models/index.ts", writer.toString());
+        }
+    }
 }

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

@@ -141,7 +141,10 @@ public void generatesServiceClients() {
         new TypeScriptCodegenPlugin().execute(context);
 
         Assertions.assertTrue(manifest.hasFile("models/index.ts"));
+        Assertions.assertTrue(manifest.hasFile("models/models_0.ts"));
         assertThat(manifest.getFileString("models/index.ts").get(),
+                containsString("export * from \"./models_0\";"));
+        assertThat(manifest.getFileString("models/models_0.ts").get(),
                 containsString("export interface Bar {\n" +
                         "  baz: string | undefined;\n" +
                         "}"));

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

@@ -471,7 +471,7 @@ private String testStructureCodegen(String file, String expectedType) {
                 .build();
 
         new TypeScriptCodegenPlugin().execute(context);
-        String contents = manifest.getFileString("/models/index.ts").get();
+        String contents = manifest.getFileString("/models/models_0.ts").get();
 
         assertThat(contents, containsString(expectedType));
         return contents;

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

@@ -1,10 +1,12 @@
 package software.amazon.smithy.typescript.codegen;
 
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
 
 import org.junit.jupiter.api.Test;
+import software.amazon.smithy.build.MockManifest;
 import software.amazon.smithy.codegen.core.Symbol;
 import software.amazon.smithy.codegen.core.SymbolProvider;
 import software.amazon.smithy.model.Model;
@@ -25,9 +27,9 @@ public void createsSymbols() {
         Symbol symbol = provider.toSymbol(shape);
 
         assertThat(symbol.getName(), equalTo("Hello"));
-        assertThat(symbol.getNamespace(), equalTo("./models/index"));
+        assertThat(symbol.getNamespace(), equalTo("./models/models_0"));
         assertThat(symbol.getNamespaceDelimiter(), equalTo("/"));
-        assertThat(symbol.getDefinitionFile(), equalTo("./models/index.ts"));
+        assertThat(symbol.getDefinitionFile(), equalTo("./models/models_0.ts"));
     }
 
     @Test
@@ -38,16 +40,20 @@ public void createsSymbolsIntoTargetNamespace() {
         SymbolProvider provider = TypeScriptCodegenPlugin.createSymbolProvider(model);
         Symbol symbol1 = provider.toSymbol(shape1);
         Symbol symbol2 = provider.toSymbol(shape2);
+        MockManifest manifest = new MockManifest();
+        SymbolVisitor.writeModelIndex(model, provider, manifest);
 
         assertThat(symbol1.getName(), equalTo("Hello"));
-        assertThat(symbol1.getNamespace(), equalTo("./models/index"));
+        assertThat(symbol1.getNamespace(), equalTo("./models/models_0"));
         assertThat(symbol1.getNamespaceDelimiter(), equalTo("/"));
-        assertThat(symbol1.getDefinitionFile(), equalTo("./models/index.ts"));
+        assertThat(symbol1.getDefinitionFile(), equalTo("./models/models_0.ts"));
 
         assertThat(symbol2.getName(), equalTo("Hello"));
-        assertThat(symbol2.getNamespace(), equalTo("./models/index"));
+        assertThat(symbol2.getNamespace(), equalTo("./models/models_0"));
         assertThat(symbol2.getNamespaceDelimiter(), equalTo("/"));
-        assertThat(symbol2.getDefinitionFile(), equalTo("./models/index.ts"));
+        assertThat(symbol2.getDefinitionFile(), equalTo("./models/models_0.ts"));
+        assertThat(manifest.getFileString("models/index.ts").get(),
+                containsString("export * from \"./models_0\";"));
     }
 
     @Test
@@ -78,7 +84,7 @@ public void doesNotEscapeBuiltins() {
 
         // Normal structure with escaping.
         assertThat(structSymbol.getName(), equalTo("_Object"));
-        assertThat(structSymbol.getNamespace(), equalTo("./models/index"));
+        assertThat(structSymbol.getNamespace(), equalTo("./models/models_0"));
 
         // Reference to built-in type with no escaping.
         assertThat(memberSymbol.getName(), equalTo("string"));