Fix shape search resolution

Fixes an issue with ShapeSearch::findShape where the local namespace was
checked before imports. It was inconsistent with the way the model
loader resolves shapes, so you could end up jumping to the wrong
definition/references, or renaming the wrong shape.

Author: milesziemer

src/main/java/software/amazon/smithy/lsp/language/ShapeSearch.java

@@ -9,6 +9,7 @@
 import java.util.List;
 import java.util.Optional;
 import software.amazon.smithy.lsp.document.DocumentId;
+import software.amazon.smithy.lsp.document.DocumentImports;
 import software.amazon.smithy.lsp.project.SmithyFile;
 import software.amazon.smithy.lsp.syntax.NodeCursor;
 import software.amazon.smithy.lsp.syntax.StatementView;
@@ -34,9 +35,9 @@ private ShapeSearch() {
      * Attempts to find a shape using a token, {@code nameOrId}.
      *
      * <p>When {@code nameOrId} does not contain a '#', this searches for shapes
-     * either in {@code idlParse}'s namespace, in {@code idlParse}'s
-     * imports, or the prelude, in that order. When {@code nameOrId} does contain
-     * a '#', it is assumed to be a full shape id and is searched for directly.
+     * either in {@code idlParse}'s imports, in {@code idlParse}'s namespace, or
+     * the prelude, in that order. When {@code nameOrId} does contain a '#', it
+     * is assumed to be a full shape id and is searched for directly.
      *
      * @param parseResult The parse result of the file {@code nameOrId} is within.
      * @param nameOrId    The name or shape id of the shape to find.
@@ -45,58 +46,69 @@ private ShapeSearch() {
      */
     static Optional<Shape> findShape(Syntax.IdlParseResult parseResult, String nameOrId, Model model) {
         return switch (nameOrId) {
+            case null -> Optional.empty();
+
             case String s when s.isEmpty() -> Optional.empty();
-            case String s when s.contains("#") -> tryFrom(s).flatMap(model::getShape);
-            case String s -> {
-                Optional<Shape> fromCurrent = tryFromParts(parseResult.namespace().namespace(), s)
-                        .flatMap(model::getShape);
-                if (fromCurrent.isPresent()) {
-                    yield fromCurrent;
-                }
 
-                if (nameOrId.contains("$")) {
-                    // Relative member id, so it could be a member of an imported shape
-                    String[] split = nameOrId.split("\\$");
-                    String containerName = split[0];
-                    String memberName = split[1];
-                    for (String fileImport : parseResult.imports().imports()) {
-                        var imported = tryFrom(fileImport)
-                                .filter(importId -> importId.getName().equals(containerName))
-                                .map(importId -> importId.withMember(memberName))
-                                .flatMap(model::getShape);
-                        if (imported.isPresent()) {
-                            yield imported;
-                        }
-                    }
-                } else {
-                    for (String fileImport : parseResult.imports().imports()) {
-                        Optional<Shape> imported = tryFrom(fileImport)
-                                .filter(importId -> importId.getName().equals(s))
-                                .flatMap(model::getShape);
-                        if (imported.isPresent()) {
-                            yield imported;
-                        }
-                    }
-                }
+            case String s when s.contains("#") -> tryFrom(s, model);
 
-                yield tryFromParts(Prelude.NAMESPACE, s).flatMap(model::getShape);
-            }
-            case null -> Optional.empty();
+            default -> fromImports(parseResult.imports(), nameOrId, model)
+                    .or(() -> tryFromRelative(parseResult.namespace().namespace(), nameOrId, model))
+                    .or(() -> tryFromRelative(Prelude.NAMESPACE, nameOrId, model));
         };
     }
 
-    private static Optional<ShapeId> tryFrom(String id) {
+    private static Optional<Shape> fromImports(DocumentImports imports, String nameOrId, Model model) {
+        if (imports.imports().isEmpty()) {
+            return Optional.empty();
+        }
+
+        if (nameOrId.contains("$")) {
+            // Relative member id, so it could be a member of an imported shape
+            String[] split = nameOrId.split("\\$");
+            String containerName = split[0];
+            String memberName = split[1];
+            String matchString = "#" + containerName;
+            for (String fileImport : imports.imports()) {
+                if (fileImport.endsWith(matchString)) {
+                    return tryWithMember(fileImport, memberName, model);
+                }
+            }
+        } else {
+            String matchString = "#" + nameOrId;
+            for (String fileImport : imports.imports()) {
+                if (fileImport.endsWith(matchString)) {
+                    return tryFrom(fileImport, model);
+                }
+            }
+        }
+
+        return Optional.empty();
+    }
+
+    private static Optional<Shape> tryFrom(String id, Model model) {
+        try {
+            ShapeId shapeId = ShapeId.from(id);
+            return model.getShape(shapeId);
+        } catch (ShapeIdSyntaxException e) {
+            return Optional.empty();
+        }
+    }
+
+    private static Optional<Shape> tryWithMember(String rootId, String memberName, Model model) {
         try {
-            return Optional.of(ShapeId.from(id));
-        } catch (ShapeIdSyntaxException ignored) {
+            ShapeId shapeId = ShapeId.from(rootId).withMember(memberName);
+            return model.getShape(shapeId);
+        } catch (ShapeIdSyntaxException e) {
             return Optional.empty();
         }
     }
 
-    private static Optional<ShapeId> tryFromParts(String namespace, String name) {
+    private static Optional<Shape> tryFromRelative(String namespace, String name, Model model) {
         try {
-            return Optional.of(ShapeId.fromRelative(namespace, name));
-        } catch (ShapeIdSyntaxException ignored) {
+            ShapeId shapeId = ShapeId.fromRelative(namespace, name);
+            return model.getShape(shapeId);
+        } catch (ShapeIdSyntaxException e) {
             return Optional.empty();
         }
     }