Improve Location for Diagnostics

What is changed?
* Updated SmithyLanguageServer.java to improve location for diagnostics.
* Added a new method to find contiguousRange for non-whitespace characters given a source location.

Why is it necessary?
* To improve location for diagnostics.
* See [Issue #76](#76)

How was the change tested?
* Updated unit tests in SmithyLanguageServerTest.java and DocumentParserTest.java
* Ran the full test suite to ensure no regressions
* Manually tested the language server with various Smithy documents to verify improvements

src/main/java/software/amazon/smithy/lsp/SmithyLanguageServer.java

@@ -806,30 +806,41 @@ List<Diagnostic> getFileDiagnostics(String uri) {
     private static Diagnostic toDiagnostic(ValidationEvent validationEvent, SmithyFile smithyFile) {
         DiagnosticSeverity severity = toDiagnosticSeverity(validationEvent.getSeverity());
         SourceLocation sourceLocation = validationEvent.getSourceLocation();
+        Range range = determineRange(validationEvent, sourceLocation, smithyFile);
+        String message = validationEvent.getId() + ": " + validationEvent.getMessage();
+        return new Diagnostic(range, message, severity, "Smithy");
+    }
+
+    private static Range determineRange(ValidationEvent validationEvent,
+                                        SourceLocation sourceLocation,
+                                        SmithyFile smithyFile) {
+        final Range defaultRange = LspAdapter.lineOffset(LspAdapter.toPosition(sourceLocation));
+
+        if (smithyFile == null) {
+            return defaultRange;
+        }
 
-        // TODO: Improve location of diagnostics
-        Range range = LspAdapter.lineOffset(LspAdapter.toPosition(sourceLocation));
+        DocumentParser parser = DocumentParser.forDocument(smithyFile.document());
+
+        // Case where we have shapes present
         if (validationEvent.getShapeId().isPresent()) {
             // Event is (probably) on a member target
             if (validationEvent.containsId("Target")) {
                 DocumentShape documentShape = smithyFile.documentShapesByStartPosition()
                         .get(LspAdapter.toPosition(sourceLocation));
                 if (documentShape != null && documentShape.hasMemberTarget()) {
-                    range = documentShape.targetReference().range();
+                    return documentShape.targetReference().range();
                 }
-            }  else {
+            } else {
                 // Check if the event location is on a trait application
-                Range traitRange = DocumentParser.forDocument(smithyFile.document()).traitIdRange(sourceLocation);
-                if (traitRange != null) {
-                    range = traitRange;
-                }
+                return Objects.requireNonNullElse(parser.traitIdRange(sourceLocation), defaultRange);
             }
         }
 
-        String message = validationEvent.getId() + ": " + validationEvent.getMessage();
-        return new Diagnostic(range, message, severity, "Smithy");
+        return  Objects.requireNonNullElse(parser.findContiguousRange(sourceLocation), defaultRange);
     }
 
+
     private static DiagnosticSeverity toDiagnosticSeverity(Severity severity) {
         return switch (severity) {
             case ERROR, DANGER -> DiagnosticSeverity.Error;

src/main/java/software/amazon/smithy/lsp/document/DocumentParser.java

@@ -704,4 +704,57 @@ private void nextNonWsNonComment() {
     private void reset() {
         rewind(0, 1, 1);
     }
+
+    /**
+     * Finds a contiguous range of non-whitespace characters starting from the given SourceLocation.
+     * If the sourceLocation happens to be a whitespace character, it returns a Range representing that column.
+     *
+     * Here is how it works:
+     * 1. We first jump to sourceLocation. If we can't, we return null.
+     * 2. We then check if the sourceLocation is a whitespace character. If it is, we return that column.
+     * 3. We then find the start of the contiguous range by traversing backwards until a whitespace character is found.
+     * 4. We then find the end of the contiguous range by traversing forwards until a whitespace character is found.
+     *
+     * @param sourceLocation The starting location to search from.
+     * @return A Range object representing the contiguous non-whitespace characters,
+     *         or null if not found.
+     */
+    public Range findContiguousRange(SourceLocation sourceLocation) {
+        if (!jumpToSource(sourceLocation)) {
+            return null;
+        }
+
+        int startLine = sourceLocation.getLine() - 1;
+        int startColumn = sourceLocation.getColumn() - 1;
+
+        if (jumpToSource(sourceLocation) && isWs()) {
+            return new Range(
+                        new Position(startLine, startColumn),
+                        // As per LSP docs the end postion is exclusive,
+                        // so adding `+1` makes it highlight the startColumn.
+                        new Position(startLine, startColumn + 1)
+                    );
+        }
+
+        // The column offset is NOT the position, but an offset from the sourceLocation column.
+        // This is required as the `isWs` uses offset, and not position to determine whether the token at the offset
+        // is whitespace or not.
+        int startColumnOffset = 0;
+        // Find the start of the contiguous range by traversing backwards until a whitespace.
+        while (startColumn + startColumnOffset > 0 && !isWs(startColumnOffset - 1)) {
+            startColumnOffset--;
+        }
+
+        int endColumn = startColumn;
+        // Find the end of the contiguous range
+        while (!isEof() && !isWs()) {
+            endColumn++;
+            skip();
+        }
+
+        // We add one to the column as it helps us shift it to correct character.
+        return new Range(
+                new Position(startLine, startColumn + startColumnOffset),
+                new Position(startLine, endColumn));
+    }
 }

src/test/java/software/amazon/smithy/lsp/SmithyLanguageServerTest.java

@@ -58,6 +58,7 @@
 import org.eclipse.lsp4j.Location;
 import org.eclipse.lsp4j.Position;
 import org.eclipse.lsp4j.PublishDiagnosticsParams;
+import org.eclipse.lsp4j.Range;
 import org.eclipse.lsp4j.SymbolInformation;
 import org.eclipse.lsp4j.TextDocumentIdentifier;
 import org.eclipse.lsp4j.TextEdit;
@@ -801,6 +802,56 @@ public void diagnosticsOnMemberTarget() {
         assertThat(diagnostic.getRange(), hasText(document, equalTo("Bar")));
     }
 
+    @Test
+    public void diagnosticsOnInvalidStructureMember() {
+        String model = safeString("""
+                $version: "2"
+                namespace com.foo
+                
+                structure Foo {
+                    abc
+                }
+                """);
+        TestWorkspace workspace = TestWorkspace.singleModel(model);
+        SmithyLanguageServer server = initFromWorkspace(workspace);
+        String uri = workspace.getUri("main.smithy");
+
+        List<Diagnostic> diagnostics = server.getFileDiagnostics(uri);
+        assertThat(diagnostics, hasSize(1));
+
+        Diagnostic diagnostic = diagnostics.getFirst();
+        Document document = server.getFirstProject().getDocument(uri);
+
+        assertThat(diagnostic.getRange(), equalTo(
+                new Range(
+                        new Position(4, 7),
+                        new Position(4, 8)
+                    )
+                )
+        );
+    }
+
+    @Test
+    public void diagnosticsOnUse() {
+        String model = safeString("""
+                $version: "2"
+                namespace com.foo
+                
+                use mything#SomeUnknownThing
+                """);
+        TestWorkspace workspace = TestWorkspace.singleModel(model);
+        SmithyLanguageServer server = initFromWorkspace(workspace);
+        String uri = workspace.getUri("main.smithy");
+
+        List<Diagnostic> diagnostics = server.getFileDiagnostics(uri);
+
+        Diagnostic diagnostic = diagnostics.getFirst();
+        Document document = server.getFirstProject().getDocument(uri);
+
+        assertThat(diagnostic.getRange(), hasText(document, equalTo("mything#SomeUnknownThing")));
+
+    }
+
     @Test
     public void diagnosticOnTrait() {
         String model = safeString("""

src/test/java/software/amazon/smithy/lsp/document/DocumentParserTest.java

@@ -11,22 +11,29 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
 import static software.amazon.smithy.lsp.document.DocumentTest.safeIndex;
 import static software.amazon.smithy.lsp.document.DocumentTest.safeString;
 import static software.amazon.smithy.lsp.document.DocumentTest.string;
 
 import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
 import org.eclipse.lsp4j.Position;
 import org.eclipse.lsp4j.Range;
 import org.junit.jupiter.api.Test;
+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.protocol.LspAdapter;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.SourceLocation;
 import software.amazon.smithy.model.shapes.Shape;
 
 public class DocumentParserTest {
+
     @Test
     public void jumpsToLines() {
         String text = """
@@ -67,6 +74,8 @@ public void jumpsToLines() {
         assertEquals(1, parser.column());
     }
 
+
+
     @Test
     public void jumpsToSource() {
         String text = "abc\ndef\nghi\n";
@@ -319,4 +328,63 @@ enum Baz {
         assertThat(getInputA.kind(), equalTo(DocumentShape.Kind.DefinedMember));
         assertThat(getInputA.shapeName(), string("a"));
     }
+
+    @ParameterizedTest
+    @MethodSource("contiguousRangeTestCases")
+    public void findsContiguousRange(SourceLocation input, Range expected)  {
+        String text = """
+            abc def
+              ghi jkl
+            mno  pqr
+            """;
+        DocumentParser parser = DocumentParser.of(safeString(text));
+
+        Range actual = parser.findContiguousRange(input);
+
+        if (expected == null) {
+            assertNull(actual);
+        } else {
+            assertEquals(expected, actual);
+        }
+    }
+
+    private static Stream<Arguments> contiguousRangeTestCases() {
+        return Stream.of(
+                // Middle of a word
+                Arguments.of(
+                        new SourceLocation("test.smithy", 1, 2),
+                        new Range(new Position(0, 0), new Position(0, 3))
+                ),
+                // Start of a word
+                Arguments.of(
+                        new SourceLocation("test.smithy", 1, 5),
+                        new Range(new Position(0, 4), new Position(0, 7))
+                ),
+                // End of a word
+                Arguments.of(
+                        new SourceLocation("test.smithy", 1, 7),
+                        new Range(new Position(0, 4), new Position(0, 7))
+                ),
+                // Start of line
+                Arguments.of(
+                        new SourceLocation("test.smithy", 3, 1),
+                        new Range(new Position(2, 0), new Position(2, 3))
+                ),
+                // End of line
+                Arguments.of(
+                        new SourceLocation("test.smithy", 3, 6),
+                        new Range(new Position(2, 5), new Position(2, 8))
+                ),
+// Invalid location
+                Arguments.of(
+                        new SourceLocation("test.smithy", 10, 1),
+                        null
+                ),
+                // At whitespace between words
+                Arguments.of(
+                        new SourceLocation("test.smithy", 1, 4),
+                        new Range(new Position(0, 3), new Position(0, 4))
+                )
+        );
+    }
 }