General refs in rule heads (#5913)

* Adding support for multiple variables at arbitrary locations in rule refs
* Updating type-checker to handle general ref heads

Fixes: #5993
Fixes: #5994

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

format/format.go

@@ -461,7 +461,7 @@ func (w *writer) writeElse(rule *ast.Rule, o fmtOpts, comments []*ast.Comment) [
 
 func (w *writer) writeHead(head *ast.Head, isDefault, isExpandedConst bool, o fmtOpts, comments []*ast.Comment) []*ast.Comment {
 	ref := head.Ref()
-	if head.Key != nil && head.Value == nil {
+	if head.Key != nil && head.Value == nil && !head.HasDynamicRef() {
 		ref = ref.GroundPrefix()
 	}
 	if o.refHeads || len(ref) == 1 {

format/format_test.go

@@ -78,6 +78,8 @@ func TestFormatSourceError(t *testing.T) {
 }
 
 func TestFormatSource(t *testing.T) {
+	t.Setenv("EXPERIMENTAL_GENERAL_RULE_REFS", "true")
+
 	regoFiles, err := filepath.Glob("testfiles/*.rego")
 	if err != nil {
 		panic(err)

format/testfiles/test_ref_heads.rego

@@ -11,3 +11,10 @@ q[1] = y if true
 r[x] if x := 10
 p.q.r[x] if x := 10
 p.q.r[2] if true
+
+g[h].i[j].k { true }
+g[h].i[j].k { h := 1; j = 2 }
+g[3].i[j].k = x { j := 3; x = 4 }
+g[h].i[j].k[l] if { true }
+g[h].i[j].k[l] contains x { x = "foo" }
+g[h].i[j].k[l] contains x { h := 5; j := 6; l = 7; x = "foo" }

format/testfiles/test_ref_heads.rego.formatted

@@ -17,3 +17,26 @@ r[x] if x := 10
 p.q.r[x] if x := 10
 
 p.q.r[2] = true
+
+g[h].i[j].k = true
+
+g[h].i[j].k if {
+	h := 1
+	j = 2
+}
+
+g[3].i[j].k = x if {
+	j := 3
+	x = 4
+}
+
+g[h].i[j].k[l] = true
+
+g[h].i[j].k[l] contains x if x = "foo"
+
+g[h].i[j].k[l] contains x if {
+	h := 5
+	j := 6
+	l = 7
+	x = "foo"
+}