Fix rebuilding when SourceLocation is NONE (#196)

The rebuild index checks which shapes have traits that are applied in
other files (i.e. via an `apply`), so we can preserve those traits in a
reload. However, if a trait's source location is `SourceLocation.NONE`,
which can happen if it wasn't set on the trait's builder, the rebuild
index would consider that "a trait applied from another file". If the
trait isn't actually being applied from another file, this would cause
a duplicate trait conflict, since the trait would be added both from the
rebuild index, and the ModelAssembler reparsing.

This commit changes two things:
1. The trait's _node_ source location is now used for the comparison,
   since it should always have a source location as it is constructed
   within smithy-model, not by trait provider implementations
2. This source location is checked to see if it is NONE.

src/main/java/software/amazon/smithy/lsp/project/Project.java

@@ -345,13 +345,15 @@ private void removeFileForReload(
             // those traits.
 
             // This shape's dependencies files will be removed and re-loaded
-            this.rebuildIndex.getDependenciesFiles(toShapeId).forEach((depPath) ->
-                    removeFileForReload(assembler, builder, depPath, visited));
+            for (String depPath : this.rebuildIndex.getDependenciesFiles(toShapeId)) {
+                removeFileForReload(assembler, builder, depPath, visited);
+            }
 
             // Traits applied in other files are re-added to the assembler so if/when the shape
             // is reloaded, it will have those traits
-            this.rebuildIndex.getTraitsAppliedInOtherFiles(toShapeId).forEach((trait) ->
-                    assembler.addTrait(toShapeId.toShapeId(), trait));
+            for (Trait trait : this.rebuildIndex.getTraitsAppliedInOtherFiles(toShapeId)) {
+                assembler.addTrait(toShapeId.toShapeId(), trait);
+            }
         }
     }
 
@@ -507,8 +509,9 @@ RebuildIndex recompute(ValidatedResult<Model> modelResult) {
                     Node traitNode = traitApplication.toNode();
                     if (traitNode.isArrayNode()) {
                         for (Node element : traitNode.expectArrayNode()) {
-                            String elementSourceFilename = element.getSourceLocation().getFilename();
-                            if (!elementSourceFilename.equals(shapeSourceFilename)) {
+                            SourceLocation elementSourceLocation = element.getSourceLocation();
+                            String elementSourceFilename = elementSourceLocation.getFilename();
+                            if (!isNone(elementSourceLocation) && !elementSourceFilename.equals(shapeSourceFilename)) {
                                 newIndex.filesToDependentFiles
                                         .computeIfAbsent(elementSourceFilename, (f) -> new HashSet<>())
                                         .add(shapeSourceFilename);
@@ -518,8 +521,9 @@ RebuildIndex recompute(ValidatedResult<Model> modelResult) {
                             }
                         }
                     } else {
-                        String traitSourceFilename = traitApplication.getSourceLocation().getFilename();
-                        if (!traitSourceFilename.equals(shapeSourceFilename)) {
+                        SourceLocation traitSourceLocation = traitNode.getSourceLocation();
+                        String traitSourceFilename = traitSourceLocation.getFilename();
+                        if (!isNone(traitSourceLocation) && !traitSourceFilename.equals(shapeSourceFilename)) {
                             newIndex.shapesToAppliedTraitsInOtherFiles
                                     .computeIfAbsent(shape.getId(), (i) -> new ArrayList<>())
                                     .add(traitApplication);
@@ -534,5 +538,9 @@ RebuildIndex recompute(ValidatedResult<Model> modelResult) {
 
             return newIndex;
         }
+
+        private static boolean isNone(SourceLocation sourceLocation) {
+            return sourceLocation.equals(SourceLocation.NONE);
+        }
     }
 }

src/test/java/software/amazon/smithy/lsp/project/ProjectTest.java

@@ -9,23 +9,32 @@
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.containsInAnyOrder;
+import static org.hamcrest.Matchers.empty;
 import static software.amazon.smithy.lsp.UtilMatchers.anOptionalOf;
 
 import java.net.URISyntaxException;
 import java.net.URL;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
 import org.junit.jupiter.api.Test;
 import software.amazon.smithy.lsp.ServerState;
+import software.amazon.smithy.lsp.SmithyMatchers;
 import software.amazon.smithy.lsp.TestWorkspace;
 import software.amazon.smithy.lsp.document.Document;
 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;
 import software.amazon.smithy.model.shapes.ShapeId;
+import software.amazon.smithy.model.shapes.StringShape;
 import software.amazon.smithy.model.traits.LengthTrait;
 import software.amazon.smithy.model.traits.PatternTrait;
 import software.amazon.smithy.model.traits.TagsTrait;
+import software.amazon.smithy.model.validation.ValidatedResult;
 
 public class ProjectTest {
     @Test
@@ -426,6 +435,56 @@ public void loadsEmptyProjectWhenThereAreNoConfigFiles() throws Exception {
         assertThat(project.type(), equalTo(Project.Type.EMPTY));
     }
 
+    @Test
+    public void changingTraitWithSourceLocationNone() {
+        // Manually construct a Project with a model containing a trait with SourceLocation.NONE,
+        // since this test can't rely on any specific trait always having SourceLocation.NONE, as
+        // it may be fixed upstream.
+        Path root = Path.of("foo").toAbsolutePath();
+        String fooPath = root.resolve("foo.smithy").toString();
+        SmithyFile fooSmithyFile = SmithyFile.create(fooPath, Document.of("""
+                $version: "2"
+                namespace com.foo
+                @length(max: 10)
+                string Foo
+                """));
+        Map<String, SmithyFile> smithyFiles = new HashMap<>();
+        smithyFiles.put(fooPath, fooSmithyFile);
+        Model model = Model.builder()
+                .addShape(StringShape.builder()
+                        .id("com.foo#Foo")
+                        .source(fooPath, 4, 1)
+                        .addTrait(LengthTrait.builder()
+                                .sourceLocation(SourceLocation.NONE)
+                                .min(1L)
+                                .build())
+                        .build())
+                .build();
+        ValidatedResult<Model> modelResult = ValidatedResult.fromValue(model);
+        Project.RebuildIndex rebuildIndex = Project.RebuildIndex.create(modelResult);
+
+        Project project = new Project(
+                root,
+                ProjectConfig.empty(),
+                BuildFiles.of(List.of()),
+                smithyFiles,
+                Model::assembler,
+                Project.Type.DETACHED,
+                modelResult,
+                rebuildIndex,
+                List.of()
+        );
+
+        assertThat(project.modelResult(), SmithyMatchers.hasValue(SmithyMatchers.hasShapeWithId("com.foo#Foo")));
+        assertThat(project.modelResult().getValidationEvents(), empty());
+
+        String fooUri = LspAdapter.toUri(fooPath);
+        project.updateModelWithoutValidating(fooUri);
+
+        assertThat(project.modelResult(), SmithyMatchers.hasValue(SmithyMatchers.hasShapeWithId("com.foo#Foo")));
+        assertThat(project.modelResult().getValidationEvents(), empty());
+    }
+
     public static Project load(Path root) {
         try {
             return ProjectLoader.load(root, new ServerState());