Completion: Fix completion through an import pointing to a local bind (#118)

* Completion: Fix completion through an import pointing to a local bind
Closes #113

This also fixes the go-to-definition functionality for this case, since it's the same code

* Change order. Var ref is probably more exact than a stack search

* Make it work nested

Author: julienduchesne

pkg/ast/processing/top_level_objects.go

@@ -37,15 +37,37 @@ func FindTopLevelObjects(stack *nodestack.NodeStack, vm *jsonnet.VM) []*ast.Desu
 			rootNode, _, _ := vm.ImportAST(string(curr.Loc().File.DiagnosticFileName), filename)
 			stack.Push(rootNode)
 		case *ast.Index:
-			container := stack.Peek()
-			if containerObj, containerIsObj := container.(*ast.DesugaredObject); containerIsObj {
-				indexValue, indexIsString := curr.Index.(*ast.LiteralString)
-				if !indexIsString {
+			indexValue, indexIsString := curr.Index.(*ast.LiteralString)
+			if !indexIsString {
+				continue
+			}
+
+			var container ast.Node
+			// If our target is a var, the container for the index is the var ref
+			if varTarget, targetIsVar := curr.Target.(*ast.Var); targetIsVar {
+				ref, err := FindVarReference(varTarget, vm)
+				if err != nil {
+					log.WithError(err).Errorf("Error finding var reference, ignoring this node")
 					continue
 				}
-				objs := findObjectFieldsInObject(containerObj, indexValue.Value, false)
-				if len(objs) > 0 {
-					stack.Push(objs[0].Body)
+				container = ref
+			}
+
+			// If we have not found a viable container, peek at the next object on the stack
+			if container == nil {
+				container = stack.Peek()
+			}
+
+			var possibleObjects []*ast.DesugaredObject
+			if containerObj, containerIsObj := container.(*ast.DesugaredObject); containerIsObj {
+				possibleObjects = []*ast.DesugaredObject{containerObj}
+			} else if containerImport, containerIsImport := container.(*ast.Import); containerIsImport {
+				possibleObjects = FindTopLevelObjectsInFile(vm, containerImport.File.Value, string(containerImport.Loc().File.DiagnosticFileName))
+			}
+
+			for _, obj := range possibleObjects {
+				for _, field := range findObjectFieldsInObject(obj, indexValue.Value, false) {
+					stack.Push(field.Body)
 				}
 			}
 		case *ast.Var:

pkg/server/completion_test.go

@@ -454,28 +454,6 @@ func TestCompletion(t *testing.T) {
 				},
 			},
 		},
-		// TODO: This one doesn't work yet
-		// Issue: https://github.com/grafana/jsonnet-language-server/issues/113
-		// {
-		// 	name:            "autocomplete local at root 2",
-		// 	filename:        "testdata/local-at-root-2.jsonnet",
-		// 	replaceString:   "hello.to",
-		// 	replaceByString: "hello.",
-		// 	expected: protocol.CompletionList{
-		// 		IsIncomplete: false,
-		// 		Items: []protocol.CompletionItem{
-		// 			{
-		// 				Label:      "to",
-		// 				Kind:       protocol.FieldCompletion,
-		// 				Detail:     "hello.to",
-		// 				InsertText: "to",
-		// 				LabelDetails: protocol.CompletionItemLabelDetails{
-		// 					Description: "object",
-		// 				},
-		// 			},
-		// 		},
-		// 	},
-		// },
 		{
 			// This checks that we don't match on `hello.hello.*` if we autocomplete on `hello.hel.`
 			name:            "autocomplete local at root, no partial match if full match exists",
@@ -508,6 +486,86 @@ func TestCompletion(t *testing.T) {
 				Items:        nil,
 			},
 		},
+		{
+			name:            "autocomplete local at root 2",
+			filename:        "testdata/local-at-root-2.jsonnet",
+			replaceString:   "hello.to",
+			replaceByString: "hello.",
+			expected: protocol.CompletionList{
+				IsIncomplete: false,
+				Items: []protocol.CompletionItem{
+					{
+						Label:      "to",
+						Kind:       protocol.FieldCompletion,
+						Detail:     "hello.to",
+						InsertText: "to",
+						LabelDetails: protocol.CompletionItemLabelDetails{
+							Description: "object",
+						},
+					},
+				},
+			},
+		},
+		{
+			name:            "autocomplete local at root 2, nested",
+			filename:        "testdata/local-at-root-2.jsonnet",
+			replaceString:   "hello.to",
+			replaceByString: "hello.to.",
+			expected: protocol.CompletionList{
+				IsIncomplete: false,
+				Items: []protocol.CompletionItem{
+					{
+						Label:      "the",
+						Kind:       protocol.FieldCompletion,
+						Detail:     "hello.to.the",
+						InsertText: "the",
+						LabelDetails: protocol.CompletionItemLabelDetails{
+							Description: "object",
+						},
+					},
+				},
+			},
+		},
+		{
+			name:            "autocomplete local at root 3, import chain",
+			filename:        "testdata/local-at-root-3.jsonnet",
+			replaceString:   "hello2.the",
+			replaceByString: "hello2.",
+			expected: protocol.CompletionList{
+				IsIncomplete: false,
+				Items: []protocol.CompletionItem{
+					{
+						Label:      "the",
+						Kind:       protocol.FieldCompletion,
+						Detail:     "hello2.the",
+						InsertText: "the",
+						LabelDetails: protocol.CompletionItemLabelDetails{
+							Description: "object",
+						},
+					},
+				},
+			},
+		},
+		{
+			name:            "autocomplete local at root 4, import chain",
+			filename:        "testdata/local-at-root-4.jsonnet",
+			replaceString:   "hello3.world",
+			replaceByString: "hello3.",
+			expected: protocol.CompletionList{
+				IsIncomplete: false,
+				Items: []protocol.CompletionItem{
+					{
+						Label:      "world",
+						Kind:       protocol.FieldCompletion,
+						Detail:     "hello3.world",
+						InsertText: "world",
+						LabelDetails: protocol.CompletionItemLabelDetails{
+							Description: "string",
+						},
+					},
+				},
+			},
+		},
 	}
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {

pkg/server/testdata/local-at-root-2.jsonnet

@@ -1,5 +1,3 @@
 local hello = import 'local-at-root.jsonnet';
 
-{
-  a: hello.to,
-}
+hello.to

pkg/server/testdata/local-at-root-3.jsonnet

@@ -0,0 +1,3 @@
+local hello2 = import 'local-at-root-2.jsonnet';
+
+hello2.the

pkg/server/testdata/local-at-root-4.jsonnet

@@ -0,0 +1,3 @@
+local hello3 = import 'local-at-root-3.jsonnet';
+
+hello3.world

pkg/server/testdata/local-at-root.jsonnet

@@ -5,7 +5,9 @@ local hello = {
   },
   hello: {
     to: {
-      the: 'world',
+      the: {
+        world: 'hello',
+      },
     },
   },
 };