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/SmithyTextDocumentService.java

@@ -63,6 +63,10 @@
 import software.amazon.smithy.lsp.ext.model.SmithyBuildExtensions;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.SourceLocation;
+import software.amazon.smithy.model.knowledge.NeighborProviderIndex;
+import software.amazon.smithy.model.neighbor.Walker;
+import software.amazon.smithy.model.shapes.Shape;
+import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.validation.ValidatedResult;
 import software.amazon.smithy.model.validation.ValidationEvent;
 
@@ -202,7 +206,11 @@ private List<String> textBufferContents(String path) throws IOException {
                 LspLog.println("Path " + path + " was found in temporary buffer");
                 contents = Arrays.stream(tempContents.split("\n")).collect(Collectors.toList());
             } else {
-                contents = readAll(new File(URI.create(path)));
+                try {
+                    contents = readAll(new File(URI.create(path)));
+                } catch (IllegalArgumentException e) {
+                    contents = readAll(new File(path));
+                }
             }
 
         }
@@ -254,8 +262,30 @@ private String findToken(String path, Position p) throws IOException {
     public CompletableFuture<Either<List<? extends Location>, List<? extends LocationLink>>> definition(
             DefinitionParams params) {
         try {
+            List<Location> locations;
+            Optional<ShapeId> initialShapeId = project.getShapeIdFromLocation(params.getTextDocument().getUri(),
+                    params.getPosition());
             String found = findToken(params.getTextDocument().getUri(), params.getPosition());
-            return Utils.completableFuture(Either.forLeft(project.getLocations().getOrDefault(found, noLocations)));
+            if (initialShapeId.isPresent()) {
+                Model model = project.getModel().unwrap();
+                Shape initialShape = model.getShape(initialShapeId.get()).get();
+                // Find first neighbor (non-member) with name that matches token.
+                Walker shapeWalker = new Walker(NeighborProviderIndex.of(model).getProvider());
+                Optional<ShapeId> target = shapeWalker.walkShapes(initialShape).stream()
+                        .filter(shape -> !shape.isMemberShape())
+                        .map(shape -> shape.getId())
+                        .filter(shape -> shape.getName().equals(found))
+                        .findFirst();
+                // Use location on target, or else default to initial shape.
+                locations = Collections.singletonList(project.getLocations().get(target.orElse(initialShapeId.get())));
+            } else {
+                // If cursor location doesn't correspond to definition, return locations of all shapes that match token.
+                locations = project.getLocations().entrySet().stream()
+                        .filter(entry -> entry.getKey().getName().equals(found))
+                        .map(entry -> entry.getValue())
+                        .collect(Collectors.toList());
+            }
+            return Utils.completableFuture(Either.forLeft(locations));
         } catch (Exception e) {
             // TODO: handle exception
 
@@ -330,7 +360,11 @@ private File fileUri(TextDocumentItem tdi) {
     }
 
     private File fileFromUri(String uri) {
-        return new File(URI.create(uri));
+        try {
+            return new File(URI.create(uri));
+        } catch (IllegalArgumentException e) {
+            return new File(uri);
+        }
     }
 
     /**

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

@@ -17,15 +17,19 @@
 
 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;
+import java.util.Optional;
+import java.util.function.Predicate;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import org.eclipse.lsp4j.Location;
@@ -35,17 +39,21 @@
 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.ShapeId;
 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 Map<ShapeId, Location> locations = Collections.emptyMap();
+    private final ValidatedResult<Model> model;
     private final File root;
 
     private SmithyProject(List<Path> imports, List<File> smithyFiles, List<File> externalJars, File root,
@@ -55,9 +63,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 +77,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)) {
@@ -102,14 +108,10 @@ public List<SmithyCompletionItem> getCompletions(String token) {
         return this.model.getResult().map(model -> Completions.find(model, token)).orElse(Collections.emptyList());
     }
 
-    public Map<String, List<Location>> getLocations() {
+    public Map<ShapeId, 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 +147,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()));
         }
     }
@@ -159,37 +161,162 @@ public File getRoot() {
         return this.root;
     }
 
-    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;
+    private static Map<ShapeId, Location> collectLocations(Model model) {
+        Map<ShapeId, Location> locations = new HashMap<>();
+        List<String> modelFiles = model.shapes()
+                .map(shape -> shape.getSourceLocation().getFilename())
+                .distinct()
+                .collect(Collectors.toList());
+        for (String modelFile : modelFiles) {
+            List<String> lines = getFileLines(modelFile);
+            int endMarker = getInitialEndMarker(lines);
+            int memberEndMarker = getInitialEndMarker(lines);
+
+             // 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))
+                    // TODO: Once the change in https://github.com/awslabs/smithy/pull/1192 lands, replace with with
+                    // `.sorted(Comparator.comparing(Shape::getSourceLocation).reversed())`.
+                    .sorted(new SourceLocationSorter().reversed())
+                    .collect(Collectors.toList());
 
-            Position pos = new Position(sourceLocation.getLine() - 1, sourceLocation.getColumn() - 1);
-            Location location = new Location(uri, new Range(pos, pos));
 
-            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 = getEndPosition(endMarker, lines);
                 }
-            }
-        });
 
+                // Find the end of a member's location by first trimming trailing commas, empty lines and closing
+                // structure braces.
+                if (shape.getType() == ShapeType.MEMBER) {
+                    int currentMemberEndMarker = memberEndMarker < endMarker ? memberEndMarker : endMarker;
+                    String currentLine = lines.get(currentMemberEndMarker - 1).trim();
+                    while (currentLine.startsWith("//") || currentLine.equals("") || currentLine.equals("}")) {
+                        currentMemberEndMarker = currentMemberEndMarker - 1;
+                        currentLine = lines.get(currentMemberEndMarker - 1).trim();
+                    }
+                    // Set the member's end position.
+                    endPosition = getEndPosition(currentMemberEndMarker, lines);
+                    // Advance the member end marker on any traits on the current member, so that the next member
+                    // location starts in the right place.
+                    List<Trait> traits = new ArrayList<>(shape.getAllTraits().values());
+                    if (!traits.isEmpty()) {
+                        traits.sort(new SourceLocationSorter());
+                        currentMemberEndMarker = traits.get(0).getSourceLocation().getLine();
+                    }
+                    memberEndMarker = currentMemberEndMarker - 1;
+                } else {
+                    // When handling non-member shapes, advance the end marker for traits and comments above the current
+                    // shape.
+                    endMarker = startPosition.getLine();
+                    List<Trait> traits = new ArrayList<>(shape.getAllTraits().values());
+                    // If the shape has traits, advance the end marker again.
+                    if (!traits.isEmpty()) {
+                        // TODO: Replace with Comparator when this class is removed.
+                        traits.sort(new SourceLocationSorter());
+                        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));
+                locations.put(shape.getId(), location);
+            }
+        }
         return locations;
     }
 
+    /**
+     * Returns the shapeId of the shape that corresponds to the file uri and position within the model.
+     *
+     * @param uri String uri of model file
+     * @param position Cursor position within model file
+     * @return ShapeId of corresponding shape defined at location.
+     */
+    public Optional<ShapeId> getShapeIdFromLocation(String uri, Position position) {
+        Comparator<Map.Entry<ShapeId, Location>> rangeSize = Comparator.comparing(entry ->
+                entry.getValue().getRange().getEnd().getLine() - entry.getValue().getRange().getStart().getLine());
+
+        List<Map.Entry<ShapeId, Location>> matchingShapes = locations.entrySet().stream()
+                .filter(entry -> entry.getValue().getUri().endsWith(Paths.get(uri).toString()))
+                .filter(filterByLocation(position))
+                // Since the position is in each of the overlapping shapes, return the smallest range.
+                .sorted(rangeSize)
+                .collect(Collectors.toList());
+        if (matchingShapes.isEmpty()) {
+            return Optional.empty();
+        }
+        return Optional.ofNullable(matchingShapes.get(0).getKey());
+    }
+
+    private Predicate<Map.Entry<ShapeId, Location>> filterByLocation(Position position) {
+        return entry -> {
+            Range range = entry.getValue().getRange();
+            if (range.getStart().getLine() > position.getLine()) {
+                return false;
+            }
+            if (range.getEnd().getLine() < position.getLine()) {
+                return false;
+            }
+            if (range.getStart().getLine() == position.getLine()) {
+                return range.getStart().getCharacter() <= position.getCharacter();
+            } else if (range.getEnd().getLine() == position.getLine()) {
+                return range.getEnd().getCharacter() >= position.getCharacter();
+            }
+            return true;
+        };
+    }
+
+    private static int getInitialEndMarker(List<String> lines) {
+        int endMarker = lines.size();
+        // Remove empty lines from the end of the file.
+        if (lines.size() > 0) {
+            while (lines.get(endMarker - 1).trim().equals("")) {
+                endMarker = endMarker - 1;
+            }
+        }
+        return endMarker;
+    }
+
+    // If the lines of the model were successfully loaded, return the end position of the actual shape line,
+    // otherwise set it to the start of the next line.
+    private static Position getEndPosition(int endMarker, List<String> lines) {
+        if (lines.size() >= endMarker) {
+            return new Position(endMarker - 1, lines.get(endMarker - 1).length());
+        }
+        return new Position(endMarker, 0);
+    }
+
+    private static List<String> getFileLines(String file) {
+        try {
+            if (Utils.isSmithyJarFile(file) || Utils.isJarFile(file)) {
+                return Utils.jarFileContents(Utils.toSmithyJarFile(file));
+            } else {
+                return Files.readAllLines(Paths.get(file));
+            }
+        } catch (IOException e) {
+            LspLog.println("File " + file + " could not be loaded.");
+        }
+        return Collections.emptyList();
+    }
+
+    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 +359,24 @@ private static List<File> downloadExternalDependencies(SmithyBuildExtensions ext
     private static List<File> onlyExistingFiles(Collection<File> files) {
         return files.stream().filter(File::isFile).collect(Collectors.toList());
     }
+
+    // TODO: Remove this Class once the change in https://github.com/awslabs/smithy/pull/1192 is available.
+    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());
+        }
+    }
 }