Use local ref. targets for count.index

hashicorp · Nov 22, 2022 · 5076ad2 · 5076ad2
1 parent f8c8d2e
commit 5076ad2
Show file tree

Hide file tree

Showing 8 changed files with 80 additions and 50 deletions.
diff --git a/decoder/body_extensions_test.go b/decoder/body_extensions_test.go
@@ -60,7 +60,7 @@ func TestCompletionAtPos_BodySchema_Extensions(t *testing.T) {
 				{
 					Label: "count",
 					Description: lang.MarkupContent{
-						Value: "The distinct index number (starting with 0) corresponding to the instance",
+						Value: "Total number of instances of this block.\n\n**Note**: A given block cannot use both `count` and `for_each`.",
 						Kind:  lang.MarkdownKind,
 					},
 					Detail: "optional, number",
@@ -136,7 +136,7 @@ func TestCompletionAtPos_BodySchema_Extensions(t *testing.T) {
 				{
 					Label: "count",
 					Description: lang.MarkupContent{
-						Value: "The distinct index number (starting with 0) corresponding to the instance",
+						Value: "Total number of instances of this block.\n\n**Note**: A given block cannot use both `count` and `for_each`.",
 						Kind:  lang.MarkdownKind,
 					},
 					Detail: "optional, number",
@@ -289,7 +289,16 @@ func TestCompletionAtPos_BodySchema_Extensions(t *testing.T) {
 					},
 				},
 			},
-			reference.Targets{},
+			reference.Targets{
+				{
+					LocalAddr: lang.Address{
+						lang.RootStep{Name: "count"},
+						lang.AttrStep{Name: "index"},
+					},
+					Type:        cty.Number,
+					Description: lang.PlainText("The distinct index number (starting with 0) corresponding to the instance"),
+				},
+			},
 			`resource "aws_instance" "foo" {
 	count = 4
 	cpu_count =
@@ -440,7 +449,42 @@ func TestCompletionAtPos_BodySchema_Extensions(t *testing.T) {
 					},
 				},
 			},
-			reference.Targets{},
+			reference.Targets{
+				{
+					LocalAddr: lang.Address{
+						lang.RootStep{Name: "count"},
+						lang.AttrStep{Name: "index"},
+					},
+					Type:        cty.Number,
+					Description: lang.PlainText("The distinct index number (starting with 0) corresponding to the instance"),
+					RangePtr: &hcl.Range{
+						Filename: "test.tf",
+						Start: hcl.Pos{
+							Line:   2,
+							Column: 3,
+							Byte:   34,
+						},
+						End: hcl.Pos{
+							Line:   2,
+							Column: 12,
+							Byte:   43,
+						},
+					},
+					DefRangePtr: &hcl.Range{
+						Filename: "test.tf",
+						Start: hcl.Pos{
+							Line:   2,
+							Column: 3,
+							Byte:   34,
+						},
+						End: hcl.Pos{
+							Line:   2,
+							Column: 8,
+							Byte:   39,
+						},
+					},
+				},
+			},
 			`resource "aws_instance" "foo" {
   count = 4
   foo {

diff --git a/decoder/decoder.go b/decoder/decoder.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 
 	"github.com/hashicorp/hcl-lang/lang"
+	"github.com/hashicorp/hcl-lang/reference"
 	"github.com/hashicorp/hcl-lang/schema"
 	"github.com/hashicorp/hcl/v2"
 	"github.com/hashicorp/hcl/v2/hclsyntax"
@@ -180,7 +181,8 @@ func countAttributeSchema() *schema.AttributeSchema {
 			schema.TraversalExpr{OfType: cty.Number},
 			schema.LiteralTypeExpr{Type: cty.Number},
 		},
-		Description: lang.Markdown("The distinct index number (starting with 0) corresponding to the instance"),
+		Description: lang.Markdown("Total number of instances of this block.\n\n" +
+			"**Note**: A given block cannot use both `count` and `for_each`."),
 	}
 }
 
@@ -248,24 +250,17 @@ func dynamicBlockSchema() *schema.BlockSchema {
 	}
 }
 
-func countIndexHoverData(rng hcl.Range) *lang.HoverData {
-	return &lang.HoverData{
-		Content: lang.Markdown("`count.index` _number_\n\nThe distinct index number (starting with 0) corresponding to the instance"),
-		Range:   rng,
-	}
-}
-
-func countIndexCandidate(editRng hcl.Range) lang.Candidate {
-	return lang.Candidate{
-		Label:       "count.index",
-		Detail:      "number",
-		Description: lang.PlainText("The distinct index number (starting with 0) corresponding to the instance"),
-		Kind:        lang.TraversalCandidateKind,
-		TextEdit: lang.TextEdit{
-			NewText: "count.index",
-			Snippet: "count.index",
-			Range:   editRng,
+func countIndexReferenceTarget(attr *hcl.Attribute, bodyRange hcl.Range) reference.Target {
+	return reference.Target{
+		LocalAddr: lang.Address{
+			lang.RootStep{Name: "count"},
+			lang.AttrStep{Name: "index"},
 		},
+		TargetableFromRangePtr: bodyRange.Ptr(),
+		Type:                   cty.Number,
+		Description:            lang.Markdown("The distinct index number (starting with 0) corresponding to the instance"),
+		RangePtr:               attr.Range.Ptr(),
+		DefRangePtr:            attr.NameRange.Ptr(),
 	}
 }
 

diff --git a/decoder/expression_candidates.go b/decoder/expression_candidates.go
@@ -328,9 +328,6 @@ func (d *PathDecoder) constraintToCandidates(ctx context.Context, constraint sch
 			},
 		})
 	case schema.TraversalExpr:
-		if schema.ActiveCountFromContext(ctx) && attr.Name != "count" {
-			candidates = append(candidates, countIndexCandidate(editRng))
-		}
 		if schema.ActiveForEachFromContext(ctx) && attr.Name != "for_each" {
 			candidates = append(candidates, foreachEachCandidate(editRng)...)
 		}

diff --git a/decoder/hover.go b/decoder/hover.go
@@ -269,13 +269,10 @@ func (d *PathDecoder) hoverDataForExpr(ctx context.Context, expr hcl.Expression,
 			return nil, err
 		}
 
-		countIndexAddr := lang.Address{lang.RootStep{Name: "count"}, lang.AttrStep{Name: "index"}}
 		eachKeyAddr := lang.Address{lang.RootStep{Name: "each"}, lang.AttrStep{Name: "key"}}
 		eachValueAddr := lang.Address{lang.RootStep{Name: "each"}, lang.AttrStep{Name: "value"}}
 
-		if address.Equals(countIndexAddr) && schema.ActiveCountFromContext(ctx) {
-			return countIndexHoverData(expr.Range()), nil
-		} else if address.Equals(eachKeyAddr) && schema.ActiveForEachFromContext(ctx) {
+		if address.Equals(eachKeyAddr) && schema.ActiveForEachFromContext(ctx) {
 			return eachKeyHoverData(expr.Range()), nil
 		} else if address.Equals(eachValueAddr) && schema.ActiveForEachFromContext(ctx) {
 			return eachValueHoverData(expr.Range()), nil

diff --git a/decoder/hover_test.go b/decoder/hover_test.go
@@ -970,7 +970,7 @@ func TestDecoder_HoverAtPos_extension(t *testing.T) {
 `,
 			hcl.Pos{Line: 2, Column: 5, Byte: 24},
 			&lang.HoverData{
-				Content: lang.Markdown("**count** _optional, number_\n\nThe distinct index number (starting with 0) corresponding to the instance"),
+				Content: lang.Markdown("**count** _optional, number_\n\nTotal number of instances of this block.\n\n**Note**: A given block cannot use both `count` and `for_each`."),
 				Range: hcl.Range{
 					Filename: "test.tf",
 					Start:    hcl.Pos{Line: 2, Column: 3, Byte: 24},
@@ -1034,7 +1034,7 @@ func TestDecoder_HoverAtPos_extension(t *testing.T) {
 `,
 			hcl.Pos{Line: 3, Column: 15, Byte: 48},
 			&lang.HoverData{
-				Content: lang.Markdown("`count.index` _number_\n\nThe distinct index number (starting with 0) corresponding to the instance"),
+				Content: lang.Markdown("`count.index`\n_number_\n\nThe distinct index number (starting with 0) corresponding to the instance"),
 				Range: hcl.Range{
 					Filename: "test.tf",
 					Start:    hcl.Pos{Line: 3, Column: 11, Byte: 44},

diff --git a/decoder/reference_targets.go b/decoder/reference_targets.go
@@ -107,6 +107,12 @@ func (d *PathDecoder) decodeReferenceTargetsForBody(body hcl.Body, parentBlock *
 	content := decodeBody(body, bodySchema)
 
 	for _, attr := range content.Attributes {
+		if bodySchema.Extensions != nil {
+			if bodySchema.Extensions.Count && attr.Name == "count" && content.RangePtr != nil {
+				refs = append(refs, countIndexReferenceTarget(attr, *content.RangePtr))
+				continue
+			}
+		}
 		attrSchema, ok := bodySchema.Attributes[attr.Name]
 		if !ok {
 			if bodySchema.AnyAttribute == nil {

diff --git a/decoder/semantic_tokens.go b/decoder/semantic_tokens.go
@@ -187,17 +187,6 @@ func (d *PathDecoder) tokensForExpression(ctx context.Context, expr hclsyntax.Ex
 			return tokens
 		}
 
-		countAvailable := schema.ActiveCountFromContext(ctx)
-		countIndexAttr := lang.Address{
-			lang.RootStep{Name: "count"}, lang.AttrStep{Name: "index"},
-		}
-
-		if address.Equals(countIndexAttr) && countAvailable {
-			tokens = append(tokens, semanticTokensForTraversalExpression(eType.AsTraversal())...)
-
-			return tokens
-		}
-
 		foreachAvailable := schema.ActiveForEachFromContext(ctx)
 		eachKeyAddress := lang.Address{
 			lang.RootStep{Name: "each"}, lang.AttrStep{Name: "key"},

diff --git a/decoder/semantic_tokens_test.go b/decoder/semantic_tokens_test.go
@@ -853,7 +853,7 @@ resource "vault_auth_backend" "blah" {
 	}
 }
 
-func TestDecoder_SemanticTokensInFile_extensions(t *testing.T) {
+func TestDecoder_SemanticTokensInFile_extensions_basic(t *testing.T) {
 	bodySchema := &schema.BodySchema{
 		Blocks: map[string]*schema.BlockSchema{
 			"resource": {
@@ -1079,8 +1079,8 @@ resource "aws_instance" "app_server" {
 				},
 				End: hcl.Pos{
 					Line:   4,
-					Column: 32,
-					Byte:   92,
+					Column: 31,
+					Byte:   91,
 				},
 			},
 		},
@@ -1329,8 +1329,8 @@ resource "aws_instance" "app_server" {
 				},
 				End: hcl.Pos{
 					Line:   4,
-					Column: 32,
-					Byte:   92,
+					Column: 31,
+					Byte:   91,
 				},
 			},
 		},
@@ -1492,7 +1492,9 @@ func TestDecoder_SemanticTokensInFile_extensions_countIndexInSubBlock(t *testing
 							Body: &schema.BodySchema{
 								Attributes: map[string]*schema.AttributeSchema{
 									"attr": {
-										Expr: schema.LiteralTypeOnly(cty.Number),
+										Expr: schema.ExprConstraints{
+											schema.TraversalExpr{OfType: cty.Number},
+										},
 									},
 								},
 							},
@@ -1726,8 +1728,8 @@ resource "foobar" "name" {
 				},
 				End: hcl.Pos{
 					Line:   5,
-					Column: 22,
-					Byte:   69,
+					Column: 21,
+					Byte:   68,
 				},
 			},
 		},