fix: when the cursor is on an object lookup expression, return its type instead of the type of the key

Author: CaiJimmy

.changeset/loose-dragons-sit.md

@@ -0,0 +1,5 @@
+---
+"rumble-lsp-wrapper": minor
+---
+
+fix: when the cursor is on an object lookup expression, return its type instead of the type of the key

packages/rumble-lsp-wrapper/src/main/java/org/jsoniq/lsp/wrapper/handlers/TypeAtPosition.java

@@ -15,6 +15,7 @@
 import org.rumbledb.expressions.Expression;
 import org.rumbledb.expressions.Node;
 import org.rumbledb.expressions.module.MainModule;
+import org.rumbledb.expressions.postfix.ObjectLookupExpression;
 
 public final class TypeAtPosition implements RequestHandler {
     public static final String REQUEST_TYPE = "type-at-position";
@@ -78,6 +79,11 @@ public Result findType(String query, URI documentUri, Position position) {
     }
 
     private static Expression findExpression(Node root, Position position) {
+        Candidate objectLookupKeyCandidate = findBestObjectLookupKeyCandidate(root, position, null);
+        if (objectLookupKeyCandidate != null) {
+            return objectLookupKeyCandidate.expression();
+        }
+
         Candidate bestEnding = findBestCandidate(root, position, true, null);
         if (bestEnding != null) {
             return bestEnding.expression();
@@ -87,6 +93,30 @@ private static Expression findExpression(Node root, Position position) {
         return bestContaining == null ? null : bestContaining.expression();
     }
 
+    private static Candidate findBestObjectLookupKeyCandidate(Node node, Position position, Candidate best) {
+        if (node == null) {
+            return best;
+        }
+
+        if (node instanceof ObjectLookupExpression objectLookupExpression) {
+            Expression lookupExpression = objectLookupExpression.getLookupExpression();
+            Candidate lookupCandidate = lookupExpression == null ? null : Candidate.create(lookupExpression);
+            Candidate objectLookupCandidate = Candidate.create(objectLookupExpression);
+            boolean positionIsInsideLookupKey = lookupCandidate != null
+                    && objectLookupCandidate != null
+                    && lookupCandidate.matches(position, false);
+
+            if (positionIsInsideLookupKey) {
+                best = Candidate.prefer(best, objectLookupCandidate, false);
+            }
+        }
+
+        for (Node child : node.getChildren()) {
+            best = findBestObjectLookupKeyCandidate(child, position, best);
+        }
+        return best;
+    }
+
     private static Candidate findBestCandidate(
             Node node,
             Position position,

packages/rumble-lsp-wrapper/src/test/java/org/jsoniq/lsp/wrapper/TypeAtPositionTest.java

@@ -48,6 +48,65 @@ void returnsNestedObjectLookupType() {
         assertEquals("xs:integer", result.sequenceType().itemType().fields().get("value").toString());
     }
 
+    @Test
+    void returnsObjectLookupExpressionTypeWhenPositionIsInsideLookupKey() {
+        String query = """
+                declare variable $a := {
+                    "name": "ada",
+                    "details": {
+                        "tlf": 1233.2
+                    }
+                };
+                ($a.details.tlf)
+                """;
+        int detailsOffset = query.indexOf("details.tlf") + "det".length();
+
+        TypeAtPosition.Result result = this.typeAtPosition.findType(
+            query,
+            positionAtOffset(query, detailsOffset)
+        );
+
+        assertNotNull(result.sequenceType());
+        assertEquals("object", result.sequenceType().itemType().kind());
+        assertEquals("xs:decimal", result.sequenceType().itemType().fields().get("tlf").toString());
+        assertEquals(
+            new Range(
+                    positionAtOffset(query, query.indexOf("$a.details")),
+                    positionAtOffset(query, query.indexOf("$a.details") + "$a.details".length())
+            ),
+            result.range()
+        );
+    }
+
+    @Test
+    void returnsNestedObjectLookupExpressionTypeWhenPositionIsInsideNestedLookupKey() {
+        String query = """
+                declare variable $a := {
+                    "name": "ada",
+                    "details": {
+                        "tlf": 1233.2
+                    }
+                };
+                ($a.details.tlf)
+                """;
+        int tlfOffset = query.indexOf("tlf)") + "t".length();
+
+        TypeAtPosition.Result result = this.typeAtPosition.findType(
+            query,
+            positionAtOffset(query, tlfOffset)
+        );
+
+        assertNotNull(result.sequenceType());
+        assertEquals("xs:decimal", result.sequenceType().toString());
+        assertEquals(
+            new Range(
+                    positionAtOffset(query, query.indexOf("$a.details.tlf")),
+                    positionAtOffset(query, query.indexOf("$a.details.tlf") + "$a.details.tlf".length())
+            ),
+            result.range()
+        );
+    }
+
     @Test
     void returnsSmallestExpressionContainingPositionWhenNothingEndsThere() {
         String query = "1 + 20";