Refactored RustCrateInlineWriter to use top level sets

Author: Fahad Zubair

codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/RustCrateInlineModuleComposingWriter.kt

@@ -12,6 +12,21 @@ import java.util.concurrent.ConcurrentHashMap
 typealias DocWriter = () -> Any
 typealias InlineModuleCreator = (Symbol, Writable) -> Unit
 
+
+/**
+ * Initializes RustCrate -> InnerModule data structure.
+ */
+fun RustCrate.initializeInlineModuleWriter(debugMode : Boolean): InnerModule =
+    crateToInlineModule
+        .getOrPut(this) { InnerModule(debugMode) }
+
+/**
+ * Returns the InnerModule for the given RustCrate
+ */
+fun RustCrate.getInlineModuleWriter() : InnerModule {
+    return crateToInlineModule.getOrPut(this) { InnerModule(false) }
+}
+
 /**
  * Returns a function that can be used to create an inline module writer.
  */
@@ -23,6 +38,80 @@ fun RustCrate.createInlineModuleCreator(): InlineModuleCreator {
     }
 }
 
+/**
+ * If the passed in `shape` is a synthetic extracted shape resulting from a constrained struct member,
+ * the `Writable` is called using the structure's builder module. Otherwise the `Writable` is called
+ * using the given `module`.
+ */
+fun RustCrate.withModuleOrWithStructureBuilderModule(
+    module: RustModule,
+    shape: Shape,
+    codegenContext: ServerCodegenContext,
+    codeWritable: Writable,
+) {
+    // All structure constrained-member-shapes code is generated inside the structure builder's module.
+    val parentAndInlineModuleInfo =
+        shape.getParentAndInlineModuleForConstrainedMember(codegenContext.symbolProvider, !codegenContext.settings.codegenConfig.publicConstrainedTypes)
+    if (parentAndInlineModuleInfo == null) {
+        this.withModule(module, codeWritable)
+    } else {
+        val (parent, inline) = parentAndInlineModuleInfo
+        val inlineWriter = this.getInlineModuleWriter()
+
+        inlineWriter.withInlineModuleHierarchyUsingCrate(this, parent) {
+            inlineWriter.withInlineModuleHierarchy(this, inline) {
+                codeWritable(this)
+            }
+        }
+    }
+}
+
+/**
+ * If the passed in `shape` is a synthetic extracted shape resulting from a constrained struct member,
+ * the `Writable` is called using the structure's builder module. Otherwise the `Writable` is called
+ * using shape's `module`.
+ */
+fun RustCrate.useShapeWriterOrUseWithStructureBuilder(
+    shape: Shape,
+    codegenContext: ServerCodegenContext,
+    docWriter: DocWriter? = null,
+    writable: Writable,
+) {
+    // All structure constrained-member-shapes code is generated inside the structure builder's module.
+    val parentAndInlineModuleInfo =
+        shape.getParentAndInlineModuleForConstrainedMember(codegenContext.symbolProvider, !codegenContext.settings.codegenConfig.publicConstrainedTypes)
+    if (parentAndInlineModuleInfo == null) {
+        docWriter?.invoke()
+        this.useShapeWriter(shape, writable)
+    } else {
+        val (parent, inline) = parentAndInlineModuleInfo
+        val inlineWriter = this.getInlineModuleWriter()
+
+        inlineWriter.withInlineModuleHierarchyUsingCrate(this, parent) {
+            inlineWriter.withInlineModuleHierarchy(this, inline) {
+                writable(this)
+            }
+        }
+    }
+}
+
+/**
+ * Given a `RustWriter` calls the `Writable` using a `RustWriter` for the `inlineModule`
+ */
+fun RustCrate.withInMemoryInlineModule(
+    outerWriter: RustWriter,
+    inlineModule: RustModule.LeafModule,
+    docWriter: DocWriter?,
+    codeWritable: Writable,
+) {
+    check(inlineModule.isInline()) {
+        "module has to be an inline module for it to be used with the InlineModuleWriter"
+    }
+    this.getInlineModuleWriter().withInlineModuleHierarchy(outerWriter, inlineModule, docWriter) {
+        codeWritable(this)
+    }
+}
+
 fun RustWriter.createTestInlineModuleCreator(): InlineModuleCreator {
     return { symbol: Symbol, writable: Writable ->
         this.withInlineModule(symbol.module()) {
@@ -31,14 +120,27 @@ fun RustWriter.createTestInlineModuleCreator(): InlineModuleCreator {
     }
 }
 
-
+/**
+ * Maintains the `RustWriter` that has been created for a `RustModule.LeafModule`.
+ */
 private data class InlineModuleWithWriter(val inlineModule : RustModule.LeafModule, val writer : RustWriter)
 
+/**
+ * For each RustCrate a separate mapping of inline-module to `RustWriter` is maintained.
+ */
+private val crateToInlineModule: ConcurrentHashMap<RustCrate, InnerModule> =
+    ConcurrentHashMap()
+
 class InnerModule(private val debugMode : Boolean) {
-    private val topLevelModuleWriters: HashMap<RustWriter, MutableList<InlineModuleWithWriter>> = hashMapOf()
+    private val topLevelModuleWriters: MutableSet<RustWriter> = mutableSetOf()
     private val inlineModuleWriters: HashMap<RustWriter, MutableList<InlineModuleWithWriter>> = hashMapOf()
     private val docWriters: HashMap<RustModule.LeafModule, MutableList<DocWriter>> = hashMapOf()
     private val writerCreator = RustWriter.factory(debugMode)
+    private val emptyLineCount: Int = writerCreator
+        .apply("lines-it-always-writes.rs", "crate")
+        .toString()
+        .split("\n")[0]
+        .length
 
     private fun createNewInlineModule(): RustWriter {
         val writer = writerCreator.apply("unknown-module-would-never-be-written.rs", "crate")
@@ -67,7 +169,9 @@ class InnerModule(private val debugMode : Boolean) {
         val bottomMost = hierarchy.removeLast()
 
         // In case it is a top level module that has been passed (e.g. ModelsModule, OutputsModule) then
-        // register it with the topLevel HashMap and call the writable on it.
+        // register it with the topLevel writers and call the writable on it. Otherwise, go over the
+        // complete hierarchy, registering each of the inner modules and then call the `Writable`
+        // with the bottom most inline module that has been passed.
         if (hierarchy.isNotEmpty()) {
             val topMost = hierarchy.removeFirst()
 
@@ -80,8 +184,7 @@ class InnerModule(private val debugMode : Boolean) {
 
                 withInlineModule(writer, bottomMost as RustModule.LeafModule, docWriter, writable)
             }
-        }
-        else {
+        } else {
             check(!bottomMost.isInline()) {
                 "there is only one module in hierarchy so it has to be non-inlined"
             }
@@ -137,55 +240,63 @@ class InnerModule(private val debugMode : Boolean) {
      * Writes out each inline module's code (`toString`) to the respective top level `RustWriter`.
      */
     fun render() {
-        fun renderDescendents(writer: RustWriter, inMemoryWriter: RustWriter) {
-            val innerModuleCode = inMemoryWriter.toString()
-            writer.writeWithNoFormatting(innerModuleCode)
+        fun writeInlineCode(rustWriter: RustWriter, code: String) {
+            val inlineCode = code.drop(emptyLineCount)
+            rustWriter.writeWithNoFormatting(inlineCode)
+        }
+
+        fun renderDescendents(topLevelWriter: RustWriter, inMemoryWriter: RustWriter) {
+            // Traverse all descendent inline modules and render them.
+            inlineModuleWriters[inMemoryWriter]?.forEach {
+                writeDocs(it.inlineModule)
 
-            inlineModuleWriters.get(inMemoryWriter)?.forEach {
-                writeDocs(it.inlineModule, it.writer)
-                writer.withInlineModule(it.inlineModule) {
+                topLevelWriter.withInlineModule(it.inlineModule) {
+                    writeInlineCode(this, it.writer.toString())
                     renderDescendents(this, it.writer)
                 }
-                it.writer.dependencies.forEach { dep -> writer.addDependency(dep) }
+
+                // Add dependencies introduced by the inline module to the
+                it.writer.dependencies.forEach { dep -> topLevelWriter.addDependency(dep) }
             }
         }
 
         // Go over all the top level modules, create an `inlineModule` on the `RustWriter`
         // and call the descendent hierarchy renderer using the `inlineModule::RustWriter`
-        topLevelModuleWriters.forEach { (outerWriter, list) ->
-            list.forEach {
-                writeDocs(it.inlineModule, it.writer)
-                outerWriter.withInlineModule(it.inlineModule) {
-                    renderDescendents(this, it.writer)
-                }
-                it.writer.dependencies.forEach { dep -> outerWriter.addDependency(dep) }
+        topLevelModuleWriters.forEach {
+            val inlineModuleWithWriter = inlineModuleWriters[it]
+            if (inlineModuleWithWriter != null) {
+                renderDescendents(it, it)
             }
         }
     }
 
+    /**
+     * Given the inline-module returns an existing `RustWriter`, or if that inline module
+     * has never been registered before then a new `RustWriter` is created and returned.
+     */
     private fun getWriter(outerWriter: RustWriter, inlineModule: RustModule.LeafModule): RustWriter {
         // Is this one of our inner writers?
         val nestedModuleWriter = inlineModuleWriters[outerWriter]
         if (nestedModuleWriter != null) {
             return findOrAddToList(nestedModuleWriter, inlineModule)
         }
 
-        val existing = topLevelModuleWriters[outerWriter]
-        return if (existing == null) {
-            val inlineWriter = createNewInlineModule()
-            topLevelModuleWriters[outerWriter] = mutableListOf(InlineModuleWithWriter(inlineModule, inlineWriter))
-            inlineWriter
-        } else {
-            return findOrAddToList(existing, inlineModule)
-        }
+        val inlineWriters = registerTopMostWriter(outerWriter)
+        return findOrAddToList(inlineWriters, inlineModule)
     }
 
-    private fun registerTopMostWriter(outerWriter: RustWriter) {
-        topLevelModuleWriters.getOrPut(outerWriter) {
-            mutableListOf()
-        }
+    /**
+     * Records the root of a dependency graph of inline modules.
+     */
+    private fun registerTopMostWriter(outerWriter: RustWriter) : MutableList<InlineModuleWithWriter> {
+        topLevelModuleWriters.add(outerWriter)
+        return inlineModuleWriters.getOrPut(outerWriter) { mutableListOf() }
     }
 
+    /**
+     * Either gets a new `RustWriter` for the inline module or creates a new one and adds it to
+     * the list of inline modules.
+     */
     private fun findOrAddToList(
         inlineModuleList: MutableList<InlineModuleWithWriter>,
         lookForModule: RustModule.LeafModule
@@ -202,98 +313,9 @@ class InnerModule(private val debugMode : Boolean) {
         }
     }
 
-    private fun writeDocs(innerModule: RustModule.LeafModule, rustWriter: RustWriter) {
+    private fun writeDocs(innerModule: RustModule.LeafModule) {
         docWriters[innerModule]?.forEach{
             it()
         }
     }
 }
-
-// A map of RustCrate to InlineModule
-private val crateToInlineModule: ConcurrentHashMap<RustCrate, InnerModule> =
-    ConcurrentHashMap()
-private val writerToCrateLookup: ConcurrentHashMap<RustWriter, RustCrate> =
-    ConcurrentHashMap()
-
-fun RustCrate.initializeInlineModuleWriter(codegenContext: ServerCodegenContext): InnerModule =
-    crateToInlineModule
-        .getOrPut(this) { InnerModule(codegenContext.settings.codegenConfig.debugMode) }
-
-fun RustCrate.getInlineModuleWriter() : InnerModule {
-    return crateToInlineModule.getOrPut(this) { InnerModule(false) }
-}
-
-
-/**
- * If the passed in `shape` is a synthetic extracted shape resulting from a constrained struct member,
- * the `Writable` is called using the structure's builder module. Otherwise the `Writable` is called
- * using the given `module`.
- */
-fun RustCrate.withModuleOrWithStructureBuilderModule(
-    module: RustModule,
-    shape: Shape,
-    codegenContext: ServerCodegenContext,
-    codeWritable: Writable,
-) {
-    // All structure constrained-member-shapes code is generated inside the structure builder's module.
-    val parentAndInlineModuleInfo =
-        shape.getParentAndInlineModuleForConstrainedMember(codegenContext.symbolProvider, !codegenContext.settings.codegenConfig.publicConstrainedTypes)
-    if (parentAndInlineModuleInfo == null) {
-        this.withModule(module, codeWritable)
-    } else {
-        val (parent, inline) = parentAndInlineModuleInfo
-        val inlineWriter = this.getInlineModuleWriter()
-
-        inlineWriter.withInlineModuleHierarchyUsingCrate(this, parent) {
-            inlineWriter.withInlineModuleHierarchy(this, inline) {
-                codeWritable(this)
-            }
-        }
-    }
-}
-
-/**
- * If the passed in `shape` is a synthetic extracted shape resulting from a constrained struct member,
- * the `Writable` is called using the structure's builder module. Otherwise the `Writable` is called
- * using shape's `module`.
- */
-fun RustCrate.useShapeWriterOrUseWithStructureBuilder(
-    shape: Shape,
-    codegenContext: ServerCodegenContext,
-    docWriter: DocWriter? = null,
-    writable: Writable,
-) {
-    // All structure constrained-member-shapes code is generated inside the structure builder's module.
-    val parentAndInlineModuleInfo =
-        shape.getParentAndInlineModuleForConstrainedMember(codegenContext.symbolProvider, !codegenContext.settings.codegenConfig.publicConstrainedTypes)
-    if (parentAndInlineModuleInfo == null) {
-        docWriter?.invoke()
-        this.useShapeWriter(shape, writable)
-    } else {
-        val (parent, inline) = parentAndInlineModuleInfo
-        val inlineWriter = this.getInlineModuleWriter()
-
-        inlineWriter.withInlineModuleHierarchyUsingCrate(this, parent) {
-            inlineWriter.withInlineModuleHierarchy(this, inline) {
-                writable(this)
-            }
-        }
-    }
-}
-
-/**
- * Given a `RustWriter` calls the `Writable` using a `RustWriter` for the `inlineModule`
- */
-fun RustCrate.withInMemoryInlineModule(
-    outerWriter: RustWriter,
-    inlineModule: RustModule.LeafModule,
-    docWriter: DocWriter?,
-    codeWritable: Writable,
-) {
-    check(inlineModule.isInline()) {
-        "module has to be an inline module for it to be used with the InlineModuleWriter"
-    }
-    this.getInlineModuleWriter().withInlineModuleHierarchy(outerWriter, inlineModule, docWriter) {
-        codeWritable(this)
-    }
-}

codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/ServerCodegenVisitor.kt

@@ -216,7 +216,7 @@ open class ServerCodegenVisitor(
             }
         }
 
-        rustCrate.initializeInlineModuleWriter(codegenContext)
+        rustCrate.initializeInlineModuleWriter(codegenContext.settings.codegenConfig.debugMode)
 
         val serviceShapes = DirectedWalker(model).walkShapes(service)
         serviceShapes.forEach { it.accept(this) }