topdown: fix panic in partial eval with ref head rule

Fixes open-policy-agent#6027.

Signed-off-by: Stephan Renatus <stephan@styra.com>

topdown/eval.go

@@ -2503,6 +2503,10 @@ func (e evalVirtualPartial) evalOneRulePreUnify(iter unifyIterator, rule *ast.Ru
 }
 
 func (e evalVirtualPartial) evalOneRulePostUnify(iter unifyIterator, rule *ast.Rule) error {
+	headKey := rule.Head.Key
+	if headKey == nil {
+		headKey = rule.Head.Reference[len(rule.Head.Reference)-1]
+	}
 
 	key := e.ref[e.pos+1]
 	child := e.e.child(rule.Body)
@@ -2512,7 +2516,7 @@ func (e evalVirtualPartial) evalOneRulePostUnify(iter unifyIterator, rule *ast.R
 
 	err := child.eval(func(child *eval) error {
 		defined = true
-		return e.e.biunify(rule.Head.Key, key, child.bindings, e.bindings, func() error {
+		return e.e.biunify(headKey, key, child.bindings, e.bindings, func() error {
 			return e.evalOneRuleContinue(iter, rule, child)
 		})
 	})

topdown/topdown_partial_test.go

@@ -3320,6 +3320,25 @@ func TestTopDownPartialEval(t *testing.T) {
 			}`},
 			wantQueries: []string{`every __local0__1, __local1__1 in input.ys { __local1__1 = input.foo }; x1 = input.foo`},
 		},
+		{ // https://github.com/open-policy-agent/opa/issues/6027
+			note:  "ref heads: \"double\" unification, single-value rule",
+			query: "data.test.foo[input.a][input.b]",
+			modules: []string{`package test
+			foo.bar[baz] {
+				baz := "baz"
+			}`},
+			wantQueries: []string{`"bar" = input.a; "baz" = input.b`},
+		},
+		{ // https://github.com/open-policy-agent/opa/issues/6027
+			note:  "ref heads: \"double\" unification, multi-value rule",
+			query: "data.test.foo[input.a][input.b]",
+			modules: []string{`package test
+			import future.keywords.contains
+			foo.bar contains baz {
+				baz := "baz"
+			}`},
+			wantQueries: []string{`"bar" = input.a; "baz" = input.b`},
+		},
 	}
 
 	ctx := context.Background()