From c039f6f07116f22599ab489174355977e7a368b4 Mon Sep 17 00:00:00 2001
From: Miles Ziemer <mziemer@amazon.com>
Date: Wed, 8 May 2024 13:40:39 -0400
Subject: [PATCH] Fix trailing doc comment conversion, update test case

---
 .../amazon/smithy/syntax/FixBadDocComments.java   | 10 ++++++++++
 .../fixes-invalid-doc-comments.formatted.smithy   | 15 +++++++++++++--
 .../formatter/fixes-invalid-doc-comments.smithy   | 14 ++++++++++++--
 3 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/smithy-syntax/src/main/java/software/amazon/smithy/syntax/FixBadDocComments.java b/smithy-syntax/src/main/java/software/amazon/smithy/syntax/FixBadDocComments.java
index a358a3f0f1d..d4e180f6cb0 100644
--- a/smithy-syntax/src/main/java/software/amazon/smithy/syntax/FixBadDocComments.java
+++ b/smithy-syntax/src/main/java/software/amazon/smithy/syntax/FixBadDocComments.java
@@ -82,6 +82,16 @@ public TokenTree apply(TokenTree tree) {
                     updateNestedChildren(br);
                 }
             }
+            // Fix any trailing doc comments in shape bodies
+            for (TreeCursor members : shapeStatements.findChildrenByType(TreeType.SHAPE_MEMBERS)) {
+                TreeCursor closeBrace = members.getLastChild();
+                if (closeBrace != null) {
+                    TreeCursor possibleTrailingWs = closeBrace.getPreviousSibling();
+                    if (possibleTrailingWs != null && possibleTrailingWs.getTree().getType() == TreeType.WS) {
+                        updateDirectChildren(possibleTrailingWs);
+                    }
+                }
+            }
         }
 
         return tree;
diff --git a/smithy-syntax/src/test/resources/software/amazon/smithy/syntax/formatter/fixes-invalid-doc-comments.formatted.smithy b/smithy-syntax/src/test/resources/software/amazon/smithy/syntax/formatter/fixes-invalid-doc-comments.formatted.smithy
index d9537ec43ed..e7e8d0d29b7 100644
--- a/smithy-syntax/src/test/resources/software/amazon/smithy/syntax/formatter/fixes-invalid-doc-comments.formatted.smithy
+++ b/smithy-syntax/src/test/resources/software/amazon/smithy/syntax/formatter/fixes-invalid-doc-comments.formatted.smithy
@@ -35,7 +35,12 @@ structure Foo {
     // 17 (change)
     bar: String
 
-    /// 18 (TODO: handle this case somehow)
+    /// 17a (keep)
+    @length(min: 1)
+    baz: String = ""
+
+    /// 18 (TODO: Fix trailing comment after VALUE_ASSIGNMENT)
+
 }
 
 // 19 (change)
@@ -45,4 +50,10 @@ apply Foo @tags(["a"])
 list Baz {
     member: Integer
 }
-// 21 (change)
+
+structure Foo2 {
+    foo2: String
+
+    // 21 (change)
+}
+// 22 (change)
diff --git a/smithy-syntax/src/test/resources/software/amazon/smithy/syntax/formatter/fixes-invalid-doc-comments.smithy b/smithy-syntax/src/test/resources/software/amazon/smithy/syntax/formatter/fixes-invalid-doc-comments.smithy
index 22dccaa6945..5aeb298412a 100644
--- a/smithy-syntax/src/test/resources/software/amazon/smithy/syntax/formatter/fixes-invalid-doc-comments.smithy
+++ b/smithy-syntax/src/test/resources/software/amazon/smithy/syntax/formatter/fixes-invalid-doc-comments.smithy
@@ -35,7 +35,11 @@ structure Foo {
     /// 17 (change)
     bar: String
 
-    /// 18 (TODO: handle this case somehow)
+    /// 17a (keep)
+    @length(min: 1)
+    baz: String = ""
+
+    /// 18 (TODO: Fix trailing comment after VALUE_ASSIGNMENT)
 }
 
 /// 19 (change)
@@ -46,4 +50,10 @@ list Baz {
     member: Integer
 }
 
-/// 21 (change)
+structure Foo2 {
+    foo2: String
+
+    /// 21 (change)
+}
+
+/// 22 (change)