From b45c1f03d8ee6975b85d5f968a11db36efe1dedb Mon Sep 17 00:00:00 2001
From: Russell Cohen <rcoh@amazon.com>
Date: Wed, 15 Jun 2022 12:20:47 -0400
Subject: [PATCH] Fixes for RustWriter bugs #1465 & #1459 (#1467)

---
 CHANGELOG.next.toml                           |  8 ++++-
 .../rust/codegen/rustlang/RustWriter.kt       |  8 +++--
 .../amazon/smithy/rust/lang/RustWriterTest.kt | 32 +++++++++++++++++++
 3 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.next.toml b/CHANGELOG.next.toml
index 16c847df7f..79ce9f47c1 100644
--- a/CHANGELOG.next.toml
+++ b/CHANGELOG.next.toml
@@ -21,4 +21,10 @@ author = "rcoh"
 message = "Add `Debug` implementation to several types in `aws-config`"
 references = ["smithy-rs#1421"]
 meta = { "breaking" = false, "tada" = false, "bug" = false }
-author = "jdisanti"
\ No newline at end of file
+author = "jdisanti"
+
+[[smithy-rs]]
+message = "Fix RustWriter bugs for `rustTemplate` and `docs` utility methods"
+references = ["smithy-rs#1427", "smithy-rs#1465", "smithy-rs#1459"]
+meta = { "breaking" = false, "tada" = false, "bug" = true }
+author = "rcoh"
diff --git a/codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/rustlang/RustWriter.kt b/codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/rustlang/RustWriter.kt
index 9e8dd8b061..af20e97b99 100644
--- a/codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/rustlang/RustWriter.kt
+++ b/codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/rustlang/RustWriter.kt
@@ -136,8 +136,9 @@ fun <T : AbstractCodeWriter<T>> T.rust(
 /* rewrite #{foo} to #{foo:T} (the smithy template format) */
 private fun transformTemplate(template: String, scope: Array<out Pair<String, Any>>): String {
     check(scope.distinctBy { it.first.lowercase() }.size == scope.size) { "Duplicate cased keys not supported" }
-    return template.replace(Regex("""#\{([a-zA-Z_0-9]+)\}""")) { matchResult ->
+    return template.replace(Regex("""#\{([a-zA-Z_0-9]+)(:\w)?\}""")) { matchResult ->
         val keyName = matchResult.groupValues[1]
+        val templateType = matchResult.groupValues[2].ifEmpty { ":T" }
         if (!scope.toMap().keys.contains(keyName)) {
             throw CodegenException(
                 "Rust block template expected `$keyName` but was not present in template.\n  hint: Template contains: ${
@@ -145,7 +146,7 @@ private fun transformTemplate(template: String, scope: Array<out Pair<String, An
                 }"
             )
         }
-        "#{${keyName.lowercase()}:T}"
+        "#{${keyName.lowercase()}$templateType}"
     }.trim()
 }
 
@@ -259,6 +260,9 @@ fun <T : AbstractCodeWriter<T>> T.docs(text: String, vararg args: Any, newlinePr
             // Rustdoc warns on tabs in documentation
             it.trimStart().replace("\t", "  ")
         }
+    // Because writing docs relies on the newline prefix, ensure that there was a new line written
+    // before we write the docs
+    this.ensureNewline()
     write(cleaned, *args)
     popState()
     return this
diff --git a/codegen/src/test/kotlin/software/amazon/smithy/rust/lang/RustWriterTest.kt b/codegen/src/test/kotlin/software/amazon/smithy/rust/lang/RustWriterTest.kt
index fed41b1399..aee26d475c 100644
--- a/codegen/src/test/kotlin/software/amazon/smithy/rust/lang/RustWriterTest.kt
+++ b/codegen/src/test/kotlin/software/amazon/smithy/rust/lang/RustWriterTest.kt
@@ -15,12 +15,16 @@ import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.SetShape
 import software.amazon.smithy.model.shapes.StringShape
 import software.amazon.smithy.model.shapes.StructureShape
+import software.amazon.smithy.rust.codegen.rustlang.Attribute
 import software.amazon.smithy.rust.codegen.rustlang.CargoDependency
 import software.amazon.smithy.rust.codegen.rustlang.RustType
 import software.amazon.smithy.rust.codegen.rustlang.RustWriter
+import software.amazon.smithy.rust.codegen.rustlang.asType
 import software.amazon.smithy.rust.codegen.rustlang.docs
 import software.amazon.smithy.rust.codegen.rustlang.isEmpty
+import software.amazon.smithy.rust.codegen.rustlang.rust
 import software.amazon.smithy.rust.codegen.rustlang.rustBlock
+import software.amazon.smithy.rust.codegen.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.rustlang.writable
 import software.amazon.smithy.rust.codegen.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.testutil.asSmithyModel
@@ -111,4 +115,32 @@ class RustWriterTest {
         val w = writable {}
         w.isEmpty() shouldBe true
     }
+
+    @Test
+    fun `attributes with comments when using rust`() {
+        val sut = RustWriter.forModule("lib")
+        Attribute.Custom("foo").render(sut)
+        sut.rust("// heres an attribute")
+        sut.toString().shouldContain("#[foo]// heres an attribute")
+    }
+
+    @Test
+    fun `attributes with comments when using docs`() {
+        val sut = RustWriter.forModule("lib")
+        Attribute.Custom("foo").render(sut)
+        sut.docs("heres an attribute")
+        sut.toString().shouldContain("#[foo]\n/// heres an attribute")
+    }
+
+    @Test
+    fun `template writables with upper case names`() {
+        val inner = writable { rust("hello") }
+        val sut = RustWriter.forModule("lib")
+        sut.rustTemplate(
+            "inner: #{Inner:W}, regular: #{http}",
+            "Inner" to inner,
+            "http" to CargoDependency.Http.asType().member("foo")
+        )
+        sut.toString().shouldContain("inner: hello, regular: http::foo")
+    }
 }