Constraints are written to the correct module

Author: Fahad Zubair

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

@@ -107,9 +107,7 @@ class ConstraintViolationSymbolProvider(
         return RustModule.new(
             name = name,
             visibility = visibility,
-            // FZ rebase
-            //parent = ServerRustModule.Model,
-            parent = ServerRustModule.Model,
+            parent = module,
             inline = true,
             documentation = documentation,
         )

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

@@ -131,7 +131,7 @@ private data class InlineModuleWithWriter(val inlineModule : RustModule.LeafModu
 private val crateToInlineModule: ConcurrentHashMap<RustCrate, InnerModule> =
     ConcurrentHashMap()
 
-class InnerModule(private val debugMode : Boolean) {
+class InnerModule(debugMode : Boolean) {
     private val topLevelModuleWriters: MutableSet<RustWriter> = mutableSetOf()
     private val inlineModuleWriters: HashMap<RustWriter, MutableList<InlineModuleWithWriter>> = hashMapOf()
     private val docWriters: HashMap<RustModule.LeafModule, MutableList<DocWriter>> = hashMapOf()
@@ -179,10 +179,10 @@ class InnerModule(private val debugMode : Boolean) {
             rustCrate.withModule(topMost) {
                 var writer = this
                 hierarchy.forEach {
-                    writer = getWriter(writer, it as RustModule.LeafModule)
+                    writer = getWriter(writer, it)
                 }
 
-                withInlineModule(writer, bottomMost as RustModule.LeafModule, docWriter, writable)
+                withInlineModule(writer, bottomMost, docWriter, writable)
             }
         } else {
             check(!bottomMost.isInline()) {
@@ -214,10 +214,10 @@ class InnerModule(private val debugMode : Boolean) {
         // Create an entry in the HashMap for all the descendent modules in the hierarchy.
         var writer = outerWriter
         hierarchy.forEach {
-            writer = getWriter(writer, it as RustModule.LeafModule)
+            writer = getWriter(writer, it)
         }
 
-        withInlineModule(writer, bottomMost as RustModule.LeafModule, docWriter, writable)
+        withInlineModule(writer, bottomMost, docWriter, writable)
     }
 
     /**
@@ -228,7 +228,7 @@ class InnerModule(private val debugMode : Boolean) {
         var hierarchy = listOf<RustModule.LeafModule>()
 
         while (current is RustModule.LeafModule) {
-            hierarchy = listOf(current as RustModule.LeafModule) + hierarchy
+            hierarchy = listOf(current) + hierarchy
             current = current.parent
         }
 
@@ -275,7 +275,6 @@ class InnerModule(private val debugMode : Boolean) {
      * 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)
@@ -301,16 +300,42 @@ class InnerModule(private val debugMode : Boolean) {
         inlineModuleList: MutableList<InlineModuleWithWriter>,
         lookForModule: RustModule.LeafModule
     ): RustWriter {
-        val inlineModule = inlineModuleList.firstOrNull() {
-            it.inlineModule == lookForModule
+        val inlineModuleAndWriter = inlineModuleList.firstOrNull() {
+            it.inlineModule.name == lookForModule.name
         }
-        return if (inlineModule == null) {
+        return if (inlineModuleAndWriter == null) {
             val inlineWriter = createNewInlineModule()
             inlineModuleList.add(InlineModuleWithWriter(lookForModule, inlineWriter))
             inlineWriter
         } else {
-            inlineModule.writer
+            check(inlineModuleAndWriter.inlineModule == lookForModule) {
+                "the two inline modules have the same name but different attributes on them"
+            }
+
+            inlineModuleAndWriter.writer
+        }
+    }
+
+    private fun combine(inlineModuleWithWriter: InlineModuleWithWriter, inlineModule: RustModule.LeafModule) : InlineModuleWithWriter {
+        check(inlineModuleWithWriter.inlineModule.name == inlineModule.name) {
+            "only inline module objects that have the same name can be combined"
+        }
+        check(inlineModuleWithWriter.inlineModule.rustMetadata.visibility == inlineModule.rustMetadata.visibility) {
+            "only inline modules with same visibility can be combined"
         }
+        check(inlineModuleWithWriter.inlineModule.inline && inlineModule.inline) {
+            "both modules need to be inline to be combined together"
+        }
+
+        val newModule = RustModule.new(
+            name = inlineModule.name,
+            parent = inlineModuleWithWriter.inlineModule.parent,
+            documentation = inlineModuleWithWriter.inlineModule.documentation?.plus("\n")?.plus(inlineModule.documentation) ?: inlineModule.documentation,
+            visibility = inlineModule.rustMetadata.visibility,
+            additionalAttributes = inlineModuleWithWriter.inlineModule.rustMetadata.additionalAttributes + inlineModule.rustMetadata.additionalAttributes,
+            inline = inlineModule.inline
+        )
+        return InlineModuleWithWriter(newModule, inlineModuleWithWriter.writer)
     }
 
     private fun writeDocs(innerModule: RustModule.LeafModule) {

codegen-server/src/test/kotlin/software/amazon/smithy/rust/codegen/server/smithy/RustCrateInlineModuleComposingWriterTest.kt

@@ -19,7 +19,6 @@ import software.amazon.smithy.rust.codegen.core.testutil.unitTest
 import software.amazon.smithy.rust.codegen.server.smithy.testutil.serverTestCodegenContext
 import software.amazon.smithy.rust.codegen.server.smithy.testutil.serverTestSymbolProvider
 import java.io.File
-import kotlin.collections.Map.Entry
 
 class RustCrateInlineModuleComposingWriterTest {
     private val rustCrate: RustCrate
@@ -66,11 +65,11 @@ class RustCrateInlineModuleComposingWriterTest {
         rustCrate = RustCrate(context.fileManifest, codegenContext.symbolProvider, settings.codegenConfig)
     }
 
-    private fun createTestInlineModule(parentModule: RustModule, moduleName : String) : RustModule.LeafModule =
+    private fun createTestInlineModule(parentModule: RustModule, moduleName : String, documentation : String? = null) : RustModule.LeafModule =
         RustModule.new(
             moduleName,
             visibility = Visibility.PUBLIC,
-            documentation = moduleName,
+            documentation = documentation ?: moduleName,
             parent = parentModule,
             inline = true,
         )
@@ -98,6 +97,13 @@ class RustCrateInlineModuleComposingWriterTest {
         }
     }
 
+//    private fun extraRustFun(writer: RustWriter, moduleName: String) {
+//        writer.rustBlock("pub fn extra()") {
+//            writer.comment("Module $moduleName")
+//            writer.rust("""println!("extra function defined inside $moduleName");""")
+//        }
+//    }
+
     @Test
     fun `simple inline module works`() {
         val testProject = TestWorkspace.testProject(serverTestSymbolProvider(model))
@@ -145,6 +151,8 @@ class RustCrateInlineModuleComposingWriterTest {
         modules["f"] = createTestInlineModule(ServerRustModule.Output, "f")
         modules["g"] = createTestInlineModule(modules["f"]!!, "g")
         modules["h"] = createTestInlineModule(ServerRustModule.Output, "h")
+        // A different kotlin object but would still go in the right place
+//        val moduleB = createTestInlineModule(ServerRustModule.Model, "b", "A new module")
 
         testProject.withModule(ServerRustModule.Model) {
             testProject.getInlineModuleWriter().withInlineModule(this, modules["a"]!!) {
@@ -165,6 +173,10 @@ class RustCrateInlineModuleComposingWriterTest {
             testProject.getInlineModuleWriter().withInlineModule(this, modules["b"]!!) {
                 byeWorld(this, "b")
             }
+
+//            testProject.getInlineModuleWriter().withInlineModule(this, moduleB) {
+//                extraRustFun(this, "b")
+//            }
         }
 
         // Write directly to an inline module without specifying the immediate parent. crate::model::b::c
@@ -175,7 +187,7 @@ class RustCrateInlineModuleComposingWriterTest {
             }
         }
         // Write to a different top level module to confirm that works.
-        testProject.withModule(ServerRustModule.Model) {
+        testProject.withModule(ServerRustModule.Input) {
             testProject.getInlineModuleWriter().withInlineModuleHierarchy(this, modules["e"]!!) {
                 helloWorld(this, "e")
             }
@@ -226,6 +238,7 @@ class RustCrateInlineModuleComposingWriterTest {
             this.unitTest("test_b") {
                 rust("crate::model::b::hello_world();")
                 rust("crate::model::b::bye_world();")
+//                rust("crate::model::b::extra();")
             }
             this.unitTest("test_someother_writer_wrote") {
                 rust("crate::model::b::some_other_writer_wrote_this();")