self complex scope improvements

pkg/nodestack/nodestack.go

@@ -40,6 +40,13 @@ func (s *NodeStack) Pop() (*NodeStack, ast.Node) {
 	return s, n
 }
 
+func (s *NodeStack) Peek() ast.Node {
+	if len(s.Stack) == 0 {
+		return nil
+	}
+	return s.Stack[len(s.Stack)-1]
+}
+
 func (s *NodeStack) IsEmpty() bool {
 	return len(s.Stack) == 0
 }

pkg/processing/find_field.go

@@ -48,9 +48,14 @@ func FindRangesFromIndexList(stack *nodestack.NodeStack, indexList []string, vm
 		}
 		foundDesugaredObjects = append(foundDesugaredObjects, lhsObject)
 	} else if start == "self" {
-		tmpStack := nodestack.NewNodeStack(stack.From)
-		tmpStack.Stack = make([]ast.Node, len(stack.Stack))
-		copy(tmpStack.Stack, stack.Stack)
+		tmpStack := stack.Clone()
+
+		// Special case. If the index was part of a binary node (ex: self.foo + {...}),
+		//   then the second element's content should not be considered to find the index's reference
+		if _, ok := tmpStack.Peek().(*ast.Binary); ok {
+			tmpStack.Pop()
+		}
+
 		foundDesugaredObjects = filterSelfScope(findTopLevelObjects(tmpStack, vm))
 	} else if start == "std" {
 		return nil, fmt.Errorf("cannot get definition of std lib")
@@ -189,6 +194,18 @@ func findTopLevelObjects(stack *nodestack.NodeStack, vm *jsonnet.VM) []*ast.Desu
 			filename := curr.File.Value
 			rootNode, _, _ := vm.ImportAST(string(curr.Loc().File.DiagnosticFileName), filename)
 			stack = stack.Push(rootNode)
+		case *ast.Index:
+			container := stack.Peek()
+			if containerObj, containerIsObj := container.(*ast.DesugaredObject); containerIsObj {
+				indexValue, indexIsString := curr.Index.(*ast.LiteralString)
+				if !indexIsString {
+					continue
+				}
+				obj := findObjectFieldInObject(containerObj, indexValue.Value)
+				if obj != nil {
+					stack.Push(obj.Body)
+				}
+			}
 		}
 	}
 	return objects

pkg/server/definition_test.go

@@ -343,6 +343,32 @@ func TestDefinition(t *testing.T) {
 				End:   protocol.Position{Line: 0, Character: 12},
 			},
 		},
+		{
+			name:     "goto self complex scope 1",
+			filename: "testdata/goto-self-complex-scoping.jsonnet",
+			position: protocol.Position{Line: 10, Character: 15},
+			targetRange: protocol.Range{
+				Start: protocol.Position{Line: 6, Character: 2},
+				End:   protocol.Position{Line: 8, Character: 3},
+			},
+			targetSelectionRange: protocol.Range{
+				Start: protocol.Position{Line: 6, Character: 2},
+				End:   protocol.Position{Line: 6, Character: 6},
+			},
+		},
+		{
+			name:     "goto self complex scope 2",
+			filename: "testdata/goto-self-complex-scoping.jsonnet",
+			position: protocol.Position{Line: 11, Character: 19},
+			targetRange: protocol.Range{
+				Start: protocol.Position{Line: 7, Character: 4},
+				End:   protocol.Position{Line: 7, Character: 18},
+			},
+			targetSelectionRange: protocol.Range{
+				Start: protocol.Position{Line: 7, Character: 4},
+				End:   protocol.Position{Line: 7, Character: 9},
+			},
+		},
 	}
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
@@ -471,7 +497,7 @@ func TestDefinitionFail(t *testing.T) {
 			params: protocol.DefinitionParams{
 				TextDocumentPositionParams: protocol.TextDocumentPositionParams{
 					TextDocument: protocol.TextDocumentIdentifier{
-						URI: "testdata/goto-self-out-of-scope.jsonnet",
+						URI: "testdata/goto-self-complex-scoping.jsonnet",
 					},
 					Position: protocol.Position{
 						Line:      3,

pkg/server/testdata/goto-self-complex-scoping.jsonnet

@@ -0,0 +1,14 @@
+{
+  test: 'test',
+  sub: {
+    test2: self.test,  // Should not be found
+  },
+
+  sub2: {
+    test3: 'test3',
+  },
+
+  sub3: self.sub2 {  // sub2 should be found
+    test4: self.test3,  // test3 should be found
+  },
+}