Updating all vars in rule ref

Fixing issue where strict-mode would only see the last var in a rule head's ref as used,
if that var is assigned in the rule body but only used in the ref.

Signed-off-by: Johan Fylling <johan.dev@fylling.se>

ast/compile.go

@@ -2385,8 +2385,9 @@ func (c *Compiler) rewriteLocalVarsInRule(rule *Rule, unusedArgs VarSet, argsSta
 	// Rewrite assignments in body.
 	used := NewVarSet()
 
-	last := rule.Head.Ref()[len(rule.Head.Ref())-1]
-	used.Update(last.Vars())
+	for _, t := range rule.Head.Ref()[1:] {
+		used.Update(t.Vars())
+	}
 
 	if rule.Head.Key != nil {
 		used.Update(rule.Head.Key.Vars())

ast/compile_test.go

@@ -6861,6 +6861,8 @@ func TestCompilerMockVirtualDocumentPartially(t *testing.T) {
 }
 
 func TestCompilerCheckUnusedAssignedVar(t *testing.T) {
+	t.Setenv("EXPERIMENTAL_GENERAL_RULE_REFS", "true")
+
 	type testCase struct {
 		note           string
 		module         string
@@ -7278,6 +7280,29 @@ func TestCompilerCheckUnusedAssignedVar(t *testing.T) {
 				&Error{Message: "assigned var y unused"},
 			},
 		},
+		{
+			note: "general ref in rule head",
+			module: `package test
+						p[q].r[s] := 1 {
+							q := "foo"
+							s := "bar"
+							t := "baz"
+						}
+		`,
+			expectedErrors: Errors{
+				&Error{Message: "assigned var t unused"},
+			},
+		},
+		{
+			note: "general ref in rule head (no errors)",
+			module: `package test
+						p[q].r[s] := 1 {
+							q := "foo"
+							s := "bar"
+						}
+		`,
+			expectedErrors: Errors{},
+		},
 	}
 
 	makeTestRunner := func(tc testCase, strict bool) func(t *testing.T) {