planner: address ref head issue, don't optimize if impossible (#7439)

When planning rules like these:

```
package authz

p.allow[action][resource] if { action := "list"; resource := "fruit" }

p.unrelated.eat.veggies if true

resp := p[input.rule][input.action][input.resource]
```

we ended up with a broken CallDynamic statement. Since the first ref
rule is planned as `g0.data.authz.p.allow` and builds an object return
value, and the second rule is planned as
`g0.data.authz.p.unrelated.eat.veggies` with a boolean return value, we cannot
dynamically dispatch their calls.

With this change, the previously existing "unbalanced ruletrie" check now
also hits before reaching the end of the ref. It'll catch this situation
and avoid optimizing the dispatch. We'll end up with a longer, less
efficient, but correct plan.

Signed-off-by: Stephan Renatus <[email protected]>

Author: srenatus

internal/planner/planner.go

@@ -2410,6 +2410,10 @@ outer:
 			opt = true
 			// take all children, they might match
 			for _, node := range nodes {
+				if nr := node.Rules(); len(nr) > 0 {
+					p.debugf("no optimization of %s: node with rules (%v)", ref, refsOfRules(nr))
+					return dont()
+				}
 				for _, child := range node.Children() {
 					if node := node.Get(child); node != nil {
 						nextNodes = append(nextNodes, node)
@@ -2419,6 +2423,10 @@ outer:
 		case ast.String:
 			// take all children that either match or have a var key // TODO(sr): Where's the code for the second part, having a var key?
 			for _, node := range nodes {
+				if nr := node.Rules(); len(nr) > 0 {
+					p.debugf("no optimization of %s: node with rules (%v)", ref, refsOfRules(nr))
+					return dont()
+				}
 				if node := node.Get(r); node != nil {
 					nextNodes = append(nextNodes, node)
 				}
@@ -2559,3 +2567,11 @@ func (p *Planner) isFunction(r ast.Ref) bool {
 func op(v ir.Val) ir.Operand {
 	return ir.Operand{Value: v}
 }
+
+func refsOfRules(rs []*ast.Rule) []string {
+	refs := make([]string, len(rs))
+	for i := range rs {
+		refs[i] = rs[i].Head.Ref().String()
+	}
+	return refs
+}

internal/planner/planner_test.go

@@ -1101,6 +1101,28 @@ func TestOptimizeLookup(t *testing.T) {
 			t.Fatalf("expected %d rules in ruleset[0], got %d\n", exp, act)
 		}
 	})
+
+	t.Run("ref heads, mixed-length rules in ruletrie", func(t *testing.T) {
+		r0 := ast.MustParseRule("allow[x].something { x := \"show\" }")
+		r1 := ast.MustParseRule("allow.see.something_else { true }")
+		r := newRuletrie()
+		val := r.LookupOrInsert(ref("primary.allow"))
+		val.rules = append(val.rules, r0)
+		val = r.LookupOrInsert(ref("secondary.allow.see.something_else"))
+		val.rules = append(val.rules, r1)
+
+		if testing.Verbose() {
+			t.Logf("rules: %v", r)
+		}
+
+		p := planner()
+		p.vars.Put(ast.Var("x"), p.newLocal())
+		_, _, _, opt := p.optimizeLookup(r, ast.MustParseRef("data[x].allow.see.something_else"))
+
+		if exp, act := false, opt; exp != act {
+			t.Errorf("expected 'optimize' %v, got %v\n", exp, act)
+		}
+	})
 }
 
 func TestPlannerCallDynamic(t *testing.T) {

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

@@ -132,3 +132,37 @@ cases:
         allow.other.stuff if true
     want_result:
       - x: true
+  - note: ir/call-dynamic not used with ref heads and mixed-length rules (issue 7399)
+    query: data.test.q = x
+    modules:
+      - |
+        package test
+
+        p.allow[action][resource] if { action := "list"; resource := "fruit" }
+
+        p.unrelated.eat.veggies if true
+
+        q := p[input.rule][input.action][input.resource]
+    input:
+      rule: allow
+      action: list
+      resource: fruit
+    want_result:
+      - x: true
+  - note: ir/call-dynamic not used with ref heads and mixed-length rules (issue 7399), second rule
+    query: data.test.q = x
+    modules:
+      - |
+        package test
+
+        p.allow[action][resource] if { action := "list"; resource := "fruit" }
+
+        p.unrelated.eat.veggies if true
+
+        q := p[input.rule][input.action][input.resource]
+    input:
+      rule: unrelated
+      action: eat
+      resource: veggies
+    want_result:
+      - x: true