From 8999cb9c0ba1873847a4d83b766952aaf11e2dd3 Mon Sep 17 00:00:00 2001
From: Miles Ziemer <45497130+milesziemer@users.noreply.github.com>
Date: Thu, 13 Feb 2025 17:44:30 -0500
Subject: [PATCH] Fix trait application and empty string parsing (#197)

https://github.com/smithy-lang/smithy-language-server/pull/190 fixed
trait application parsing for traits like `@foo()`, making it so the end
of the trait doesn't include trailing whitespace. This commit fixes the
case for non-empty traits, like `@foo(bar: "")`.

It also fixes parsing of empty strings. Previously the character _after_
the end of the string would be skipped.
---
 .../amazon/smithy/lsp/syntax/Parser.java      |  2 --
 .../amazon/smithy/lsp/TextWithPositions.java  |  8 +++++
 .../smithy/lsp/syntax/IdlParserTest.java      | 30 +++++++++++++++++--
 3 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/src/main/java/software/amazon/smithy/lsp/syntax/Parser.java b/src/main/java/software/amazon/smithy/lsp/syntax/Parser.java
index 27eb7087..a9da011d 100644
--- a/src/main/java/software/amazon/smithy/lsp/syntax/Parser.java
+++ b/src/main/java/software/amazon/smithy/lsp/syntax/Parser.java
@@ -193,7 +193,6 @@ private Syntax.Node traitValueKvps(int from) {
         while (!eof()) {
             if (is(')')) {
                 setEnd(kvps);
-                skip();
                 return kvps;
             }
 
@@ -398,7 +397,6 @@ private Syntax.Node str() {
             }
 
             // Empty string
-            skip();
             int strEnd = position();
             return new Syntax.Node.Str(start, strEnd, "");
         }
diff --git a/src/test/java/software/amazon/smithy/lsp/TextWithPositions.java b/src/test/java/software/amazon/smithy/lsp/TextWithPositions.java
index 6cb2b594..60bd2872 100644
--- a/src/test/java/software/amazon/smithy/lsp/TextWithPositions.java
+++ b/src/test/java/software/amazon/smithy/lsp/TextWithPositions.java
@@ -38,6 +38,8 @@ public record TextWithPositions(String text, Position... positions) {
     public static TextWithPositions from(String raw) {
         Document document = Document.of(safeString(raw));
         List<Position> positions = new ArrayList<>();
+
+        Position lastPosition = null;
         int i = 0;
         while (true) {
             int next = document.nextIndexOf(POSITION_MARKER, i);
@@ -45,6 +47,12 @@ public static TextWithPositions from(String raw) {
                 break;
             }
             Position position = document.positionAtIndex(next);
+            if (lastPosition != null && position.getLine() == lastPosition.getLine()) {
+                // If there's two or more markers on the same line, any markers after the
+                // first will be off by one when we do the replacement.
+                position.setCharacter(position.getCharacter() - 1);
+            }
+            lastPosition = position;
             positions.add(position);
             i = next + 1;
         }
diff --git a/src/test/java/software/amazon/smithy/lsp/syntax/IdlParserTest.java b/src/test/java/software/amazon/smithy/lsp/syntax/IdlParserTest.java
index 004cceb0..d3607b1e 100644
--- a/src/test/java/software/amazon/smithy/lsp/syntax/IdlParserTest.java
+++ b/src/test/java/software/amazon/smithy/lsp/syntax/IdlParserTest.java
@@ -18,6 +18,7 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
+import software.amazon.smithy.lsp.TextWithPositions;
 import software.amazon.smithy.lsp.document.Document;
 
 public class IdlParserTest {
@@ -304,6 +305,27 @@ public void stringKeysInTraits() {
                 Syntax.Node.Type.Str));
     }
 
+    @Test
+    public void traitApplicationsDontContainTrailingWhitespace() {
+        var twp = TextWithPositions.from("""
+                %@foo(foo: "")%
+                structure Foo {
+                    foo: Foo
+                }
+                """);
+        Document document = Document.of(twp.text());
+        Syntax.IdlParseResult parse = Syntax.parseIdl(document);
+
+        assertThat(getTypes(parse), contains(
+                Syntax.Statement.Type.TraitApplication,
+                Syntax.Statement.Type.ShapeDef,
+                Syntax.Statement.Type.MemberDef));
+
+        Syntax.Statement traitApplication = parse.statements().get(0);
+        assertThat(document.positionAtIndex(traitApplication.start()), equalTo(twp.positions()[0]));
+        assertThat(document.positionAtIndex(traitApplication.end()), equalTo(twp.positions()[1]));
+    }
+
     @ParameterizedTest
     @MethodSource("brokenProvider")
     public void broken(String desc, String text, List<String> expectedErrorMessages, List<Syntax.Statement.Type> expectedTypes) {
@@ -460,10 +482,14 @@ private static Stream<Arguments> brokenProvider() {
 
     private static void assertTypesEqual(String text, Syntax.Statement.Type... types) {
         Syntax.IdlParseResult parse = Syntax.parseIdl(Document.of(text));
-        List<Syntax.Statement.Type> actualTypes = parse.statements().stream()
+        var actualTypes = getTypes(parse);
+        assertThat(actualTypes, contains(types));
+    }
+
+    private static List<Syntax.Statement.Type> getTypes(Syntax.IdlParseResult parse) {
+        return parse.statements().stream()
                 .map(Syntax.Statement::type)
                 .filter(type -> type != Syntax.Statement.Type.Block)
                 .toList();
-        assertThat(actualTypes, contains(types));
     }
 }