Determine end of definition location

.gitignore

@@ -25,4 +25,5 @@ bin
 
 .java-version
 *.smithy
+!/src/test/resources/**/*.smithy
 .ammonite

src/main/java/software/amazon/smithy/lsp/ext/SmithyProject.java

@@ -17,12 +17,14 @@
 
 import java.io.File;
 import java.io.IOException;
+import java.io.Serializable;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -35,17 +37,20 @@
 import software.amazon.smithy.lsp.SmithyInterface;
 import software.amazon.smithy.lsp.Utils;
 import software.amazon.smithy.lsp.ext.model.SmithyBuildExtensions;
+import software.amazon.smithy.model.FromSourceLocation;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.SourceLocation;
+import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.model.shapes.ShapeType;
+import software.amazon.smithy.model.traits.Trait;
 import software.amazon.smithy.model.validation.ValidatedResult;
 
 public final class SmithyProject {
     private final List<Path> imports;
     private final List<File> smithyFiles;
     private final List<File> externalJars;
     private Map<String, List<Location>> locations = Collections.emptyMap();
-    private ValidatedResult<Model> model;
+    private final ValidatedResult<Model> model;
     private final File root;
 
     private SmithyProject(List<Path> imports, List<File> smithyFiles, List<File> externalJars, File root,
@@ -55,9 +60,7 @@ private SmithyProject(List<Path> imports, List<File> smithyFiles, List<File> ext
         this.model = model;
         this.smithyFiles = smithyFiles;
         this.externalJars = externalJars;
-        model.getResult().ifPresent(m -> {
-            this.locations = collectLocations(m);
-        });
+        model.getResult().ifPresent(m -> this.locations = collectLocations(m));
     }
 
     /**
@@ -71,7 +74,7 @@ private SmithyProject(List<Path> imports, List<File> smithyFiles, List<File> ext
      * @return either an error, or a loaded project
      */
     public Either<Exception, SmithyProject> recompile(File changed, File exclude) {
-        List<File> newFiles = new ArrayList<File>();
+        List<File> newFiles = new ArrayList<>();
 
         for (File existing : onlyExistingFiles(this.smithyFiles)) {
             if (exclude != null && !existing.equals(exclude)) {
@@ -106,10 +109,6 @@ public Map<String, List<Location>> getLocations() {
         return this.locations;
     }
 
-    public Either<Exception, SmithyProject> reload(SmithyBuildExtensions config) {
-        return load(config, this.root);
-    }
-
     /**
      * Load the project using a SmithyBuildExtensions configuration and workspace
      * root.
@@ -145,7 +144,7 @@ private static Either<Exception, SmithyProject> load(List<Path> imports, List<Fi
         if (model.isLeft()) {
             return Either.forLeft(model.getLeft());
         } else {
-            model.getRight().getValidationEvents().forEach(event -> LspLog.println(event));
+            model.getRight().getValidationEvents().forEach(LspLog::println);
             return Either.forRight(new SmithyProject(imports, smithyFiles, externalJars, root, model.getRight()));
         }
     }
@@ -161,35 +160,97 @@ public File getRoot() {
 
     private static Map<String, List<Location>> collectLocations(Model model) {
         Map<String, List<Location>> locations = new HashMap<>();
-        model.shapes().forEach(shape -> {
-            SourceLocation sourceLocation = shape.getSourceLocation();
-            String fileName = sourceLocation.getFilename();
-            String uri = Utils.isJarFile(fileName)
-                ? Utils.toSmithyJarFile(fileName)
-                : !fileName.startsWith("file:") ? "file:" + fileName
-                : fileName;
+        List<String> modelFiles = model.shapes()
+                .map(shape -> shape.getSourceLocation().getFilename())
+                .distinct()
+                .collect(Collectors.toList());
+        for (String modelFile : modelFiles) {
+            Path modelPath = Paths.get(modelFile);
+            List<String> lines = getFileLines(modelPath);
+            LspLog.println(lines);
+            int endMarker = 0;
+            try {
+                endMarker = (int) Files.lines(modelPath).count();
+            } catch (IOException e) {
+                LspLog.println("Could not read model file length");
+            }
 
-            Position pos = new Position(sourceLocation.getLine() - 1, sourceLocation.getColumn() - 1);
-            Location location = new Location(uri, new Range(pos, pos));
+            // Get shapes reverse-sorted by source location to work from bottom of file to top.
+            List<Shape> shapes = model.shapes()
+                    .filter(shape -> shape.getSourceLocation().getFilename().equals(modelFile))
+                    .sorted(new SourceLocationSorter().reversed())
+                    .collect(Collectors.toList());
 
-            String shapeName = shape.getId().getName();
-            // Members get the same shapeName as their parent structure
-            // so we ignore them, to avoil producing a location per-member
-            // TODO: index members somehow as well?
-            if (shape.getType() != ShapeType.MEMBER) {
-                if (locations.containsKey(shapeName)) {
-                    locations.get(shapeName).add(location);
+
+            for (Shape shape : shapes) {
+                SourceLocation sourceLocation = shape.getSourceLocation();
+                Position startPosition = new Position(sourceLocation.getLine() - 1, sourceLocation.getColumn() - 1);
+                Position endPosition;
+                if (endMarker < sourceLocation.getLine()) {
+                    endPosition = new Position(sourceLocation.getLine() - 1, sourceLocation.getColumn() - 1);
                 } else {
-                    List<Location> locList = new ArrayList<Location>();
-                    locList.add(location);
-                    locations.put(shapeName, locList);
+                    endPosition = new Position(endMarker, 0);
                 }
-            }
-        });
 
+                // If a shape is not a member, move the end marker for setting the next shape location.
+                if (shape.getType() != ShapeType.MEMBER) {
+                    endMarker = startPosition.getLine();
+                    List<Trait> traits = new ArrayList<>(shape.getAllTraits().values());
+                    // If the shape has traits, advance the end marker again.
+                    if (!traits.isEmpty()) {
+                        traits.sort(new TraitSorter());
+                        endMarker = traits.get(0).getSourceLocation().getLine() - 1;
+                    }
+                    // Move the end marker when encountering line comments or empty lines.
+                    if (lines.size() > endMarker) {
+                        while (lines.get(endMarker - 1).trim().startsWith("//")
+                                || lines.get(endMarker - 1).trim().equals("")) {
+                            endMarker = endMarker - 1;
+                        }
+                    }
+                }
+                Location location =  new Location(getUri(modelFile), new Range(startPosition, endPosition));
+                addLocationToMap(shape, locations, location);
+            }
+        }
         return locations;
     }
 
+    private static List<String> getFileLines(Path path) {
+        try {
+            if (Utils.isSmithyJarFile(path.toString())) {
+                return Utils.jarFileContents(path.toString());
+            }
+            return Files.readAllLines(path);
+        } catch (IOException e) {
+            LspLog.println("Path " + path + " was found in temporary buffer");
+        }
+        return Collections.emptyList();
+    }
+
+    private static void addLocationToMap(Shape shape, Map<String, List<Location>> locations, Location location) {
+        String shapeName = shape.getId().getName();
+        // Members get the same shapeName as their parent structure
+        // so we ignore them, to avoid producing a location per-member
+        // TODO: index members somehow as well?
+        if (shape.getType() != ShapeType.MEMBER) {
+            if (locations.containsKey(shapeName)) {
+                locations.get(shapeName).add(location);
+            } else {
+                List<Location> locList = new ArrayList<>();
+                locList.add(location);
+                locations.put(shapeName, locList);
+            }
+        }
+    }
+
+    private static String getUri(String fileName) {
+        return Utils.isJarFile(fileName)
+                ? Utils.toSmithyJarFile(fileName)
+                : !fileName.startsWith("file:") ? "file:" + fileName
+                : fileName;
+    }
+
     private static Boolean isValidSmithyFile(Path file) {
         String fName = file.getFileName().toString();
         return fName.endsWith(Constants.SMITHY_EXTENSION);
@@ -232,4 +293,30 @@ private static List<File> downloadExternalDependencies(SmithyBuildExtensions ext
     private static List<File> onlyExistingFiles(Collection<File> files) {
         return files.stream().filter(File::isFile).collect(Collectors.toList());
     }
+
+    private static class SourceLocationSorter implements Comparator<FromSourceLocation>, Serializable {
+        @Override
+        public int compare(FromSourceLocation s1, FromSourceLocation s2) {
+            SourceLocation sourceLocation = s1.getSourceLocation();
+            SourceLocation otherSourceLocation = s2.getSourceLocation();
+
+            if (!sourceLocation.getFilename().equals(otherSourceLocation.getFilename())) {
+                return sourceLocation.getFilename().compareTo(otherSourceLocation.getFilename());
+            }
+
+            int lineComparison = Integer.compare(sourceLocation.getLine(), otherSourceLocation.getLine());
+            if (lineComparison != 0) {
+                return lineComparison;
+            }
+
+            return Integer.compare(sourceLocation.getColumn(), otherSourceLocation.getColumn());
+        }
+    }
+
+    private static class TraitSorter implements Comparator<Trait>, Serializable {
+        @Override
+        public int compare(Trait o1, Trait o2) {
+            return new SourceLocationSorter().compare(o1.getSourceLocation(), o2.getSourceLocation());
+        }
+    }
 }

src/test/java/software/amazon/smithy/lsp/ext/Harness.java

@@ -16,16 +16,17 @@
 package software.amazon.smithy.lsp.ext;
 
 import com.google.common.io.Files;
-import org.eclipse.lsp4j.jsonrpc.messages.Either;
-
 import java.io.File;
 import java.io.FileWriter;
 import java.nio.charset.Charset;
+import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
+import org.eclipse.lsp4j.jsonrpc.messages.Either;
 import software.amazon.smithy.lsp.ext.model.SmithyBuildExtensions;
+import software.amazon.smithy.utils.IoUtils;
 
 public class Harness implements AutoCloseable {
   private File root;
@@ -95,12 +96,7 @@ public static Harness create(SmithyBuildExtensions ext) throws Exception {
     // TODO: How to make this safe?
     File hs = Files.createTempDir();
     File tmp = Files.createTempDir();
-
-    Either<Exception, SmithyProject> loaded = SmithyProject.load(ext, hs);
-    if (loaded.isRight())
-      return new Harness(hs, tmp, loaded.getRight(), ext);
-    else
-      throw loaded.getLeft();
+    return loadHarness(ext, hs, tmp);
   }
 
   public static Harness create(SmithyBuildExtensions ext, Map<String, String> files) throws Exception {
@@ -110,11 +106,23 @@ public static Harness create(SmithyBuildExtensions ext, Map<String, String> file
     for (Entry<String, String> entry : files.entrySet()) {
       safeCreateFile(entry.getKey(), entry.getValue(), hs);
     }
+    return loadHarness(ext, hs, tmp);
+  }
+
+  public static Harness create(SmithyBuildExtensions ext, List<Path> files) throws Exception {
+    File hs = Files.createTempDir();
+    File tmp = Files.createTempDir();
+    for (Path path : files) {
+      safeCreateFile(path.getFileName().toString(), IoUtils.readUtf8File(path), hs);
+    }
+    return loadHarness(ext, hs, tmp);
+  }
+
+  private static Harness loadHarness(SmithyBuildExtensions ext, File hs, File tmp) throws Exception {
     Either<Exception, SmithyProject> loaded = SmithyProject.load(ext, hs);
     if (loaded.isRight())
       return new Harness(hs, tmp, loaded.getRight(), ext);
     else
       throw loaded.getLeft();
   }
-
 }

src/test/java/software/amazon/smithy/lsp/ext/SmithyProjectTest.java

@@ -15,20 +15,24 @@
 
 package software.amazon.smithy.lsp.ext;
 
-import org.junit.Test;
-import software.amazon.smithy.lsp.ext.model.SmithyBuildExtensions;
-import software.amazon.smithy.utils.ListUtils;
-import software.amazon.smithy.utils.MapUtils;
+import static org.junit.Assert.assertEquals;
 
+import com.google.common.collect.ImmutableList;
 import java.io.File;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.nio.file.Paths;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
-
-import static org.junit.Assert.assertEquals;
+import org.eclipse.lsp4j.Location;
+import org.eclipse.lsp4j.Position;
+import org.eclipse.lsp4j.Range;
+import org.junit.Test;
+import software.amazon.smithy.lsp.ext.model.SmithyBuildExtensions;
+import software.amazon.smithy.utils.ListUtils;
+import software.amazon.smithy.utils.MapUtils;
 
 public class SmithyProjectTest {
 
@@ -71,4 +75,29 @@ public void respectingEmptyConfig() throws Exception {
 
     }
 
+    @Test
+    public void definitionLocations() throws Exception {
+        Path baseDir = Paths.get(SmithyProjectTest.class.getResource("models").toURI());
+        Path modelMain = baseDir.resolve("main.smithy");
+        Path modelTest = baseDir.resolve("test.smithy");
+        List<Path> modelFiles = ImmutableList.of(modelMain, modelTest);
+
+        try (Harness hs = Harness.create(SmithyBuildExtensions.builder().build(), modelFiles)) {
+            Map<String, List<Location>> locations = hs.getProject().getLocations();
+
+            correctLocation(locations, "SingleLine", 4, 5);
+            correctLocation(locations, "MultiLine", 6, 11);
+            correctLocation(locations, "SingleTrait", 13, 14);
+            correctLocation(locations, "MultiTrait", 17, 20);
+            correctLocation(locations, "MultiTraitAndLineComments", 33, 36);
+            correctLocation(locations,"MultiTraitAndDocComments", 41, 44);
+            correctLocation(locations, "OtherStructure", 4, 9);
+        }
+    }
+
+    private void correctLocation(Map<String, List<Location>> locations, String shapeName, int start, int end) {
+        Location location = locations.get(shapeName).get(0);
+        Range range = new Range(new Position(start, 0), new Position(end, 0));
+        assertEquals(range, location.getRange());
+    }
 }

src/test/resources/software/amazon/smithy/lsp/ext/models/main.smithy

@@ -0,0 +1,44 @@
+$version: "1.0"
+
+namespace example.foo
+
+structure SingleLine {}
+
+structure MultiLine {
+    a: String,
+    b: String,
+    c: String
+}
+
+@pattern("^[A-Za-z0-9 ]+$")
+string SingleTrait
+
+@input
+@error("client")
+structure MultiTrait {
+    a: String
+}
+
+// Line comments
+// comments
+@input
+// comments
+@error("client")
+@references(
+    [
+        {
+            resource: City
+        }
+    ]
+)
+structure MultiTraitAndLineComments {
+    a: String
+}
+
+/// Doc comments
+/// comments
+@input
+@error("client")
+structure MultiTraitAndDocComments {
+    a: String
+}