planner: adjust check introduced in open-policy-agent#5839

So the check introduced before was too broad: it aborted optimizations
at the wrong spot -- the resulting plan didn't add up: path lengths used
in CallDynamicStmt didn't match path lengths that the planned funcs
had.

So, while this change looks like over-fitting (also to me), we're really
trying to make test previous fix more specific.

Generally looking at that section of the planner, it feels like the intro
of general refs would be a good moment to nuke and start over: the way
that refs-with-vars are put into the ruletrie seems like the root cause of
our trouble here.

Fixes open-policy-agent#5964.

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

internal/planner/planner.go

@@ -2286,10 +2286,28 @@ outer:
 			break
 		}
 
-		for _, node := range nodes {
-			if len(node.Rules()) > 0 {
-				p.debugf("ref %s: at least one rule has 1+ rules, break", ref[0:index+1])
-				break outer
+		// NOTE(sr): we only need this check for the penultimate part:
+		// When planning the ref data.pkg.a[input.x][input.y],
+		// We want to capture this situation:
+		//   a.b[c] := "x" if c := "c"
+		//   a.b.d := "y"
+		//
+		// Not this:
+		//   a.b[c] := "x" if c := "c"
+		//   a.d := "y"
+		// since the length doesn't add up. Even if input.x was "d", the second
+		// rule (a.d) wouldn't contribute anything to the result, since we cannot
+		// "dot it".
+		if index == len(ref)-2 {
+			for _, node := range nodes {
+				anyNonGround := false
+				for _, r := range node.Rules() {
+					anyNonGround = anyNonGround || !r.Ref().IsGround()
+				}
+				if anyNonGround {
+					p.debugf("ref %s: at least one node has 1+ non-ground ref rules, break", ref[0:index+1])
+					break outer
+				}
 			}
 		}
 	}

test/cases/testdata/planner-ir/test-call-dynamic.yaml

@@ -40,7 +40,7 @@ cases:
     package test
     import future.keywords.if
     p := a[b].c if b := "b"
-    a.b.c := "C" if b := "b"
+    a.b.c := "C"
     a.x.d := "D"
   query: data.test.p = x
   want_result:
@@ -52,7 +52,7 @@ cases:
     import future.keywords.if
     p := a.b[c] if c := "c"
     a.b[c] := "C" if c := "c"
-    a.b["d"] := "D"
+    a.b.d := "D"
   query: data.test.p = x
   want_result:
   - x: C
@@ -66,4 +66,28 @@ cases:
     a.x.d := "D"
   query: data.test.p = x
   want_result:
-  - x: C
+  - x: C
+- note: ir/call-dynamic with ref heads, issue 5839
+  modules:
+  - |
+    package test
+    import future.keywords.if
+    p := a[b][c].allow if { b := "b"; c := "c" }
+    a.b.c.allow if true
+    a.x.d.allow if true
+    a.y[x] := "E" if x := "x"
+  query: data.test.p = x
+  want_result:
+  - x: true
+- note: ir/call-dynamic with ref heads, issue 5839, penultimate
+  modules:
+  - |
+    package test
+    import future.keywords.if
+    p := a[b][c] if { b := "b"; c := "c" }
+    a.b.c if true
+    a.x.d if true
+    a.y := "E"
+  query: data.test.p = x
+  want_result:
+  - x: true