Better support rule head keyword completions

internal/lsp/completions/providers/ruleheadkeyword.go

@@ -12,6 +12,7 @@ import (
 // RuleHeadKeyword will return completions for the keywords when starting a new rule.
 // The current cases are supported:
 // - [rule-name] if
+// - [rule-name] :=
 // - [rule-name] contains
 // - [rule-name] contains if
 // These completions are mandatory, that means they are the only ones to be shown.
@@ -42,29 +43,49 @@ func (*RuleHeadKeyword) Run(c *cache.Cache, params types.CompletionParams, _ *Op
 
 	const keyWdIf = "if"
 
-	var label string
-
-	switch {
-	// suggest contains after the name of the rule in the rule head
-	//nolint:gocritic
-	case len(words) == 2 && strings.HasPrefix(keyWdContains, lastWord):
-		label = "contains"
-	// suggest if at the end of the rule head
-	case len(words) == 4 && words[1] == keyWdContains:
-		label = keyWdIf
-	// suggest if after the rule name
-	//nolint:gocritic
-	case len(words) == 2 && strings.HasPrefix(keyWdIf, lastWord):
-		label = keyWdIf
+	const keyWdAssign = ":="
+
+	mandatory := false
+	keywords := map[string]bool{
+		keyWdIf:       true,
+		keyWdContains: true,
+		keyWdAssign:   true,
 	}
 
-	if label == "" {
-		return []types.CompletionItem{}, nil
+	if lastWord != "" {
+		mandatory = true
+		keywords = map[string]bool{
+			keyWdIf:       false,
+			keyWdContains: false,
+		}
+
+		switch {
+		// suggest the assignment operator after the name of the rule
+		case len(words) == 2 && strings.HasPrefix(keyWdAssign, lastWord):
+			keywords[keyWdAssign] = true
+		// suggest contains after the name of the rule in the rule head
+		//nolint:gocritic
+		case len(words) == 2 && strings.HasPrefix(keyWdContains, lastWord):
+			keywords[keyWdContains] = true
+		// suggest if at the end of the rule head
+		case len(words) == 4 && words[1] == keyWdContains:
+			keywords[keyWdIf] = true
+		// suggest if after the rule name
+		//nolint:gocritic
+		case len(words) == 2 && strings.HasPrefix(keyWdIf, lastWord):
+			keywords[keyWdIf] = true
+		}
 	}
 
-	return []types.CompletionItem{
-		{
-			Label: label,
+	ret := make([]types.CompletionItem, 0)
+
+	for k, v := range keywords {
+		if !v {
+			continue
+		}
+
+		ret = append(ret, types.CompletionItem{
+			Label: k,
 			Kind:  completion.Keyword,
 			TextEdit: &types.TextEdit{
 				Range: types.Range{
@@ -77,9 +98,11 @@ func (*RuleHeadKeyword) Run(c *cache.Cache, params types.CompletionParams, _ *Op
 						Character: uint(len(currentLine)),
 					},
 				},
-				NewText: label + " ",
+				NewText: k + " ",
 			},
-			Mandatory: true,
-		},
-	}, nil
+			Mandatory: mandatory,
+		})
+	}
+
+	return ret, nil
 }

internal/lsp/completions/providers/ruleheadkeyword_test.go

@@ -8,6 +8,57 @@ import (
 	"github.com/styrainc/regal/internal/lsp/types"
 )
 
+func TestRuleHeadKeyword_ManualTriggerAfterRuleName(t *testing.T) {
+	t.Parallel()
+
+	c := cache.NewCache()
+
+	fileContents := `package policy
+
+deny` + " "
+
+	c.SetFileContents(testCaseFileURI, fileContents)
+
+	p := &RuleHeadKeyword{}
+
+	completionParams := types.CompletionParams{
+		TextDocument: types.TextDocumentIdentifier{
+			URI: testCaseFileURI,
+		},
+		Position: types.Position{
+			Line:      2,
+			Character: 4,
+		},
+	}
+
+	completions, err := p.Run(c, completionParams, nil)
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	expectedLabels := []string{"contains", "if", ":="}
+
+	if len(completions) != len(expectedLabels) {
+		t.Fatalf("Expected %v completions, got: %v", len(expectedLabels), len(completions))
+	}
+
+	for _, comp := range completions {
+		found := false
+
+		for _, label := range expectedLabels {
+			if comp.Label == label {
+				found = true
+
+				break
+			}
+		}
+
+		if !found {
+			t.Fatalf("Expected label to be one of %v, got: %v", expectedLabels, comp.Label)
+		}
+	}
+}
+
 func TestRuleHeadKeyword_TypedIAfterRuleName(t *testing.T) {
 	t.Parallel()