feat(symbol-table): read in template context with variables (#85)

* feat(symbol-table): read in template context with variables

* fix: handle root variable ($) at different stage

Makefile

@@ -68,7 +68,7 @@ test:
 	@$(GO) test ./... -v -race -tags=integration
 
 coverage:
-	@$(GO) test -coverprofile=.coverage -tags=integration ./internal/... && go tool cover -html=.coverage
+	@$(GO) test -coverprofile=.coverage -tags=integration -coverpkg=./internal/... ./internal/... && go tool cover -html=.coverage
 
 .PHONY: build-release
 build-release:

internal/handler/hover_main_test.go

@@ -2,7 +2,6 @@ package handler
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -24,6 +23,24 @@ func TestHoverMain(t *testing.T) {
 		expected      string
 		expectedError error
 	}{
+		{
+			desc: "Test hover on template context with variables",
+			position: lsp.Position{
+				Line:      74,
+				Character: 50,
+			},
+			expected:      "$root.Values.deployments",
+			expectedError: nil,
+		},
+		{
+			desc: "Test hover on template context with variables in range loop",
+			position: lsp.Position{
+				Line:      80,
+				Character: 35,
+			},
+			expected:      "$config.hpa.minReplicas",
+			expectedError: nil,
+		},
 		{
 			desc: "Test hover on dot",
 			position: lsp.Position{
@@ -114,15 +131,6 @@ func TestHoverMain(t *testing.T) {
 			expected:      fmt.Sprintf("### %s\n%s\n\n", filepath.Join("..", "..", "testdata", "example", "values.yaml"), "1"),
 			expectedError: nil,
 		},
-		{
-			desc: "Test hover on template context with variables",
-			position: lsp.Position{
-				Line:      74,
-				Character: 50,
-			},
-			expected:      "",
-			expectedError: errors.New("no template context found"),
-		},
 	}
 	for _, tt := range testCases {
 		t.Run(tt.desc, func(t *testing.T) {

internal/lsp/symbol_table.go

@@ -18,7 +18,7 @@ func (t TemplateContext) Tail() TemplateContext {
 }
 
 func (t TemplateContext) IsVariable() bool {
-	return len(t) > 0 && t[0] == "$"
+	return len(t) > 0 && strings.HasPrefix(t[0], "$")
 }
 
 func (t TemplateContext) AppendSuffix(suffix string) TemplateContext {
@@ -45,7 +45,12 @@ func NewSymbolTable(ast *sitter.Tree, content []byte) *SymbolTable {
 }
 
 func (s *SymbolTable) AddTemplateContext(templateContext TemplateContext, pointRange sitter.Range) {
-	s.contexts[templateContext.Format()] = append(s.contexts[strings.Join(templateContext, ".")], pointRange)
+	if templateContext.IsVariable() && templateContext[0] == "$" {
+		// $ is a special variable that resolves to the root context
+		// we can just remove it from the template context
+		templateContext = templateContext.Tail()
+	}
+	s.contexts[templateContext.Format()] = append(s.contexts[templateContext.Format()], pointRange)
 	sliceCopy := make(TemplateContext, len(templateContext))
 	copy(sliceCopy, templateContext)
 	s.contextsReversed[pointRange] = sliceCopy

internal/lsp/symbol_table_template_context.go

@@ -60,15 +60,16 @@ func (v *TemplateContextVisitor) Enter(node *sitter.Node) {
 		v.symbolTable.AddTemplateContext(append(v.currentContext, content), GetRangeForNode(node.ChildByFieldName("name")))
 	case gotemplate.NodeTypeUnfinishedSelectorExpression:
 		operandNode := node.ChildByFieldName("operand")
-		if operandNode.Type() == gotemplate.NodeTypeVariable && operandNode.Content(v.content) == "$" {
+		if operandNode.Type() == gotemplate.NodeTypeVariable {
 			v.StashContext()
 		}
 		v.symbolTable.AddTemplateContext(append(getContextForSelectorExpression(operandNode, v.content), ""),
 			GetRangeForNode(node.Child(int(node.ChildCount())-1)))
 	case gotemplate.NodeTypeSelectorExpression:
 		operandNode := node.ChildByFieldName("operand")
-		if operandNode.Type() == gotemplate.NodeTypeVariable && operandNode.Content(v.content) == "$" {
+		if operandNode.Type() == gotemplate.NodeTypeVariable {
 			v.StashContext()
+			v.PushContext(operandNode.Content(v.content))
 		}
 	}
 }
@@ -77,7 +78,8 @@ func (v *TemplateContextVisitor) Exit(node *sitter.Node) {
 	switch node.Type() {
 	case gotemplate.NodeTypeSelectorExpression, gotemplate.NodeTypeUnfinishedSelectorExpression:
 		operandNode := node.ChildByFieldName("operand")
-		if operandNode.Type() == gotemplate.NodeTypeVariable && operandNode.Content(v.content) == "$" {
+		if operandNode.Type() == gotemplate.NodeTypeVariable {
+			v.PopContext()
 			v.RestoreStashedContext()
 		}
 	}
@@ -97,7 +99,6 @@ func (v *TemplateContextVisitor) EnterContextShift(node *sitter.Node, suffix str
 			s = s.AppendSuffix(suffix)
 			if s.IsVariable() {
 				v.StashContext()
-				s = s.Tail()
 			}
 		}
 		v.PushContextMany(s)

internal/lsp/symbol_table_template_context_test.go

@@ -0,0 +1,64 @@
+package lsp
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetContextForSelectorExpression(t *testing.T) {
+	testCases := []struct {
+		desc        string
+		template    string
+		nodeContent string
+		expected    TemplateContext
+	}{
+		{
+			desc:        "Selects simple selector expression correctly",
+			template:    `{{ .Values.test }}`,
+			nodeContent: ".Values.test",
+			expected:    TemplateContext{"Values", "test"},
+		},
+		{
+			desc:        "Selects unfinished selector expression correctly",
+			template:    `{{ .Values.test. }}`,
+			nodeContent: ".Values.test.",
+			expected:    TemplateContext{"Values", "test"},
+		},
+		{
+			desc:        "Selects selector expression with $ correctly",
+			template:    `{{ $.Values.test }}`,
+			nodeContent: "$.Values.test",
+			expected:    TemplateContext{"$", "Values", "test"},
+		},
+		{
+			desc:        "Selects unfinished selector expression with $ correctly",
+			template:    `{{ $.Values.test. }}`,
+			nodeContent: "$.Values.test.",
+			expected:    TemplateContext{"$", "Values", "test"},
+		},
+		{
+			desc:        "Selects selector expression with variable correctly",
+			template:    `{{ $x.test }}`,
+			nodeContent: "$x.test",
+			expected:    TemplateContext{"$x", "test"},
+		},
+		{
+			desc:        "Selects unfinished selector expression with variable correctly",
+			template:    `{{ $x.test. }}`,
+			nodeContent: "$x.test.",
+			expected:    TemplateContext{"$x", "test"},
+		},
+	}
+	for _, tC := range testCases {
+		t.Run(tC.desc, func(t *testing.T) {
+			ast := ParseAst(nil, tC.template)
+			node := ast.RootNode().Child(1)
+
+			assert.Equal(t, tC.nodeContent, node.Content([]byte(tC.template)))
+			result := getContextForSelectorExpression(node, []byte(tC.template))
+
+			assert.Equal(t, tC.expected, result)
+		})
+	}
+}

internal/lsp/symbol_table_test.go

@@ -217,19 +217,63 @@ func TestSymbolTableForValuesTestFile(t *testing.T) {
 
 func TestSymbolTableForValuesSingleTests(t *testing.T) {
 	type testCase struct {
-		template   string
-		path       []string
-		startPoint sitter.Point
+		template         string
+		path             []string
+		startPoint       sitter.Point
+		foundContextsLen int
 	}
 
 	testCases := []testCase{
+		{
+			template: `
+      {{- $root := . -}}
+      {{- range $type, $config := $root.Values.deployments }}
+				{{- .InLoop }}
+			{{- end }}
+			{{ .Values.test }}
+`,
+			path: []string{"$root", "Values"},
+			startPoint: sitter.Point{
+				Row:    2,
+				Column: 40,
+			},
+			foundContextsLen: 6,
+		},
+		{
+			template: `{{ $x := .Values }}{{ $x.test }}{{ .Values.test }}`,
+			path:     []string{"$x", "test"},
+			startPoint: sitter.Point{
+				Row:    0,
+				Column: 25,
+			},
+			foundContextsLen: 3,
+		},
+		{
+			template: `{{ $x.test }}`,
+			path:     []string{"$x", "test"},
+			startPoint: sitter.Point{
+				Row:    0,
+				Column: 6,
+			},
+			foundContextsLen: 1,
+		},
+		{
+			template: `{{ $x.test. }}`,
+			path:     []string{"$x", "test", ""},
+			startPoint: sitter.Point{
+				Row:    0,
+				Column: 10,
+			},
+			foundContextsLen: 2,
+		},
 		{
 			template: `{{ if (and .Values. ) }} {{ end }} `,
 			path:     []string{"Values"},
 			startPoint: sitter.Point{
 				Row:    0,
 				Column: 12,
 			},
+			foundContextsLen: 2,
 		},
 		{
 			template: `{{ if (and .Values. ) }} {{ end }} `,
@@ -238,17 +282,21 @@ func TestSymbolTableForValuesSingleTests(t *testing.T) {
 				Row:    0,
 				Column: 18,
 			},
+			foundContextsLen: 2,
 		},
 	}
 
 	for _, v := range testCases {
-		ast := ParseAst(nil, v.template)
-		symbolTable := NewSymbolTable(ast, []byte(v.template))
-		values := symbolTable.GetTemplateContextRanges(v.path)
-		points := []sitter.Point{}
-		for _, v := range values {
-			points = append(points, v.StartPoint)
-		}
-		assert.Contains(t, points, v.startPoint)
+		t.Run(v.template, func(t *testing.T) {
+			ast := ParseAst(nil, v.template)
+			symbolTable := NewSymbolTable(ast, []byte(v.template))
+			values := symbolTable.GetTemplateContextRanges(v.path)
+			points := []sitter.Point{}
+			for _, v := range values {
+				points = append(points, v.StartPoint)
+			}
+			assert.Contains(t, points, v.startPoint)
+			assert.Len(t, symbolTable.contexts, v.foundContextsLen)
+		})
 	}
 }

internal/lsp/visitor.go

@@ -43,7 +43,13 @@ func (v *Visitors) visitNodesRecursiveWithScopeShift(node *sitter.Node) {
 	case gotemplate.NodeTypeRangeAction:
 		rangeNode := node.ChildByFieldName("range")
 		if rangeNode == nil {
-			break // range is optional (e.g. {{ range $index, $element := pipeline }})
+			// for {{- range $type, $config := $root.Values.deployments }} the range node is in the
+			// range_variable_definition node an not in the range_action node
+			rangeNode = node.NamedChild(0).ChildByFieldName("range")
+			if rangeNode == nil {
+				logger.Error("Could not find range node")
+				break
+			}
 		}
 		v.visitNodesRecursiveWithScopeShift(rangeNode)
 		for _, visitor := range v.visitors {

testdata/example/templates/deployment.yaml

@@ -79,5 +79,6 @@ spec:
         name: my-app-{{ $type }}
       spec:
         replicas: {{ $config.hpa.minReplicas }}
+        test:  {{ $.Values.ingress.hosts }}
       ---
       {{- end }}