fix: also visit range variable definition

mrjosh · May 29, 2024 · 552c020 · 552c020
1 parent 9de28c4
commit 552c020
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 7 deletions.
diff --git a/internal/lsp/symbol_table_variables_test.go b/internal/lsp/symbol_table_variables_test.go
@@ -1,6 +1,7 @@
 package lsp
 
 import (
+	"fmt"
 	"testing"
 
 	sitter "github.com/smacker/go-tree-sitter"
@@ -37,16 +38,16 @@ func TestSymbolTableForVariableDefinitions(t *testing.T) {
 			expectedVariableDefinitions: map[string][]VariableDefinition{
 				"$type": {
 					{
-						Value:        ".$root.Values.deployments",
-						VariableType: VariableTypeRangeKey,
+						Value:        "$root.Values.deployments",
+						VariableType: VariableTypeRangeKeyOrIndex,
 						Range:        sitter.Range{
 							// TODO: fill this with the correct result
 						},
 					},
 				},
 				"$config": {
 					{
-						Value:        ".$root.Values.deployments",
+						Value:        "$root.Values.deployments",
 						VariableType: VariableTypeRangeValue,
 						Range:        sitter.Range{
 							// TODO: fill this with the correct result
@@ -57,14 +58,26 @@ func TestSymbolTableForVariableDefinitions(t *testing.T) {
 		},
 		{
 			template: `{{ $x := .Values }}{{ $x.test }}{{ .Values.test }}`,
+			expectedVariableDefinitions: map[string][]VariableDefinition{
+				"$x": {
+					{
+						Value:        ".Values",
+						VariableType: VariableTypeAssigment,
+						Range:        sitter.Range{
+							// TODO: fill this with the correct result
+						},
+					},
+				},
+			},
 		},
 	}
 
 	for _, v := range testCases {
 		t.Run(v.template, func(t *testing.T) {
 			ast := ParseAst(nil, v.template)
 			symbolTable := NewSymbolTable(ast, []byte(v.template))
-			assert.Equal(t, v.expectedVariableDefinitions, symbolTable.variableDefinitions)
+			assert.Equal(t, v.expectedVariableDefinitions, symbolTable.variableDefinitions,
+				fmt.Sprintf("Ast was %s", ast.RootNode()))
 		})
 	}
 }
diff --git a/internal/lsp/symobl_table_variables.go → internal/lsp/symobol_table_variables.go b/internal/lsp/symobl_table_variables.go → internal/lsp/symobol_table_variables.go
@@ -9,7 +9,7 @@ type VariableType int64
 
 const (
 	VariableTypeAssigment VariableType = iota
-	VariableTypeRangeKey
+	VariableTypeRangeKeyOrIndex
 	VariableTypeRangeValue
 )
 
@@ -37,7 +37,29 @@ func NewVariablesVisitor(symbolTable *SymbolTable, content []byte) *VariablesVis
 func (v *VariablesVisitor) Enter(node *sitter.Node) {
 	switch node.Type() {
 	case gotemplate.NodeTypeRangeVariableDefinition:
-	// TODO:
+		keyOrIndexVariableName := node.ChildByFieldName("index")
+		valueVariableName := node.ChildByFieldName("element")
+		variableValueNode := node.ChildByFieldName("range")
+		if variableValueNode == nil {
+			return
+		}
+		variableValue := variableValueNode.Content(v.content)
+		if keyOrIndexVariableName != nil {
+			v.symbolTable.AddVariableDefinition(keyOrIndexVariableName.Content(v.content), VariableDefinition{
+				Value:        variableValue,
+				VariableType: VariableTypeRangeKeyOrIndex,
+				Range:        sitter.Range{},
+			})
+		}
+
+		if valueVariableName != nil {
+			v.symbolTable.AddVariableDefinition(valueVariableName.Content(v.content), VariableDefinition{
+				Value:        variableValue,
+				VariableType: VariableTypeRangeValue,
+				Range:        sitter.Range{},
+			})
+		}
+
 	case gotemplate.NodeTypeVariableDefinition:
 		if node.NamedChildCount() < 2 {
 			return

diff --git a/internal/lsp/visitor.go b/internal/lsp/visitor.go
@@ -51,7 +51,7 @@ func (v *Visitors) visitNodesRecursiveWithScopeShift(node *sitter.Node) {
 				break
 			}
 		}
-		v.visitNodesRecursiveWithScopeShift(rangeNode)
+		v.visitNodesRecursiveWithScopeShift(node.NamedChild(0))
 		for _, visitor := range v.visitors {
 			visitor.EnterContextShift(rangeNode, "[]")
 		}