diff --git a/compile/compile.go b/compile/compile.go
index eab5141972..6484a2b014 100644
--- a/compile/compile.go
+++ b/compile/compile.go
@@ -562,7 +562,7 @@ func (c *Compiler) compilePlan(context.Context) error {
 	}
 
 	// Prepare modules and builtins for the planner.
-	modules := []*ast.Module{}
+	modules := make([]*ast.Module, 0, len(c.compiler.Modules))
 	for _, module := range c.compiler.Modules {
 		modules = append(modules, module)
 	}
diff --git a/internal/compiler/wasm/wasm.go b/internal/compiler/wasm/wasm.go
index e6d7885c61..3e39792520 100644
--- a/internal/compiler/wasm/wasm.go
+++ b/internal/compiler/wasm/wasm.go
@@ -887,12 +887,9 @@ func (c *Compiler) compileFunc(fn *ir.Func) error {
 }
 
 func mapFunc(mapping ast.Object, fn *ir.Func, index int) (ast.Object, bool) {
-	curr := ast.NewObject()
-	curr.Insert(ast.StringTerm(fn.Path[len(fn.Path)-1]), ast.IntNumberTerm(index))
+	curr := ast.NewObject(ast.Item(ast.StringTerm(fn.Path[len(fn.Path)-1]), ast.IntNumberTerm(index)))
 	for i := len(fn.Path) - 2; i >= 0; i-- {
-		o := ast.NewObject()
-		o.Insert(ast.StringTerm(fn.Path[i]), ast.NewTerm(curr))
-		curr = o
+		curr = ast.NewObject(ast.Item(ast.StringTerm(fn.Path[i]), ast.NewTerm(curr)))
 	}
 	return mapping.Merge(curr)
 }
diff --git a/internal/planner/planner.go b/internal/planner/planner.go
index 74fc56282d..93c935b352 100644
--- a/internal/planner/planner.go
+++ b/internal/planner/planner.go
@@ -147,39 +147,48 @@ func (p *Planner) buildFunctrie() error {
 		}
 
 		for _, rule := range module.Rules {
-			val := p.rules.LookupOrInsert(rule.Ref())
+			r := rule.Ref()
+			switch r[len(r)-1].Value.(type) {
+			case ast.String: // pass
+			default: // cut off
+				r = r[:len(r)-1]
+			}
+			val := p.rules.LookupOrInsert(r)
 			val.rules = append(val.rules, rule)
 		}
 	}
-
 	return nil
 }
 
-func (p *Planner) planRules(rules []*ast.Rule, cut bool) (string, error) {
-	pathRef := rules[0].Ref() // NOTE(sr): no longer the same for all those rules, respect `cut`?
-	path := pathRef.String()
+func (p *Planner) planRules(rules []*ast.Rule) (string, error) {
+	pathRef := rules[0].Ref()
 
-	var pathPieces []string
+	// figure out what our rules' collective name/path is:
+	// if we're planning both p.q.r and p.q[s], we'll name
+	// the function p.q (for the mapping table)
 	// TODO(sr): this has to change when allowing `p[v].q.r[w]` ref rules
 	// including the mapping lookup structure and lookup functions
-
-	// if we're planning both p.q.r and p.q[s], we'll name the function p.q (for the mapping table)
 	pieces := len(pathRef)
-	if cut {
-		pieces--
+	for i := range rules {
+		r := rules[i].Ref()
+		if _, ok := r[len(r)-1].Value.(ast.String); !ok {
+			pieces = len(r) - 1
+		}
 	}
+	// control if p.a = 1 is to return 1 directly; or insert 1 under key "a" into an object
+	buildObject := pieces != len(pathRef)
+
+	var pathPieces []string
 	for i := 1; /* skip `data` */ i < pieces; i++ {
 		switch q := pathRef[i].Value.(type) {
 		case ast.String:
 			pathPieces = append(pathPieces, string(q))
-		case ast.Var:
-			pathPieces = append(pathPieces, fmt.Sprintf("[%s]", q))
 		default:
-			// Needs to be fixed if we allow non-string ref pieces, like `p.q[3][4].r = x`
-			pathPieces = append(pathPieces, q.String())
+			panic("impossible")
 		}
 	}
 
+	path := pathRef[:pieces].String()
 	if funcName, ok := p.funcs.Get(path); ok {
 		return funcName, nil
 	}
@@ -219,24 +228,11 @@ func (p *Planner) planRules(rules []*ast.Rule, cut bool) (string, error) {
 
 	params := fn.Params[2:]
 
-	// control if p.a = 1 is to return 1 directly; or insert 1 under key "a" into an object
-	buildObject := false
-
 	// Initialize return value for partial set/object rules. Complete document
 	// rules assign directly to `fn.Return`.
 	switch rules[0].Head.RuleKind() {
 	case ast.SingleValue:
-		// if any rule has a non-ground last key, create an object, insert into it
-		any := false
-		for _, rule := range rules {
-			ref := rule.Head.Ref()
-			if last := ref[len(ref)-1]; len(ref) > 1 && !last.IsGround() {
-				any = true
-				break
-			}
-		}
-		if any {
-			buildObject = true
+		if buildObject {
 			fn.Blocks = append(fn.Blocks, p.blockWithStmt(&ir.MakeObjectStmt{Target: fn.Return}))
 		}
 	case ast.MultiValue:
@@ -870,7 +866,7 @@ func (p *Planner) planExprCall(e *ast.Expr, iter planiter) error {
 				if node := p.rules.Lookup(r); node != nil {
 					if node.Arity() > 0 {
 						p.mocks.Push() // new scope
-						name, err = p.planRules(node.Rules(), false)
+						name, err = p.planRules(node.Rules())
 						if err != nil {
 							return err
 						}
@@ -892,7 +888,7 @@ func (p *Planner) planExprCall(e *ast.Expr, iter planiter) error {
 		}
 
 		if node := p.rules.Lookup(op); node != nil {
-			name, err = p.planRules(node.Rules(), false)
+			name, err = p.planRules(node.Rules())
 			if err != nil {
 				return err
 			}
@@ -1658,7 +1654,7 @@ func (p *Planner) planRefData(virtual *ruletrie, base *baseptr, ref ast.Ref, ind
 	// NOTE(sr): we do it on the first index because later on, the recursion
 	// on subtrees of virtual already lost parts of the path we've taken.
 	if index == 1 && virtual != nil {
-		rulesets, path, index, optimize := p.optimizeLookup(virtual, ref.GroundPrefix())
+		rulesets, path, index, optimize := p.optimizeLookup(virtual, ref)
 		if optimize {
 			// If there are no rulesets in a situation that otherwise would
 			// allow for a call_indirect optimization, then there's nothing
@@ -1668,7 +1664,7 @@ func (p *Planner) planRefData(virtual *ruletrie, base *baseptr, ref ast.Ref, ind
 			}
 			// plan rules
 			for _, rules := range rulesets {
-				if _, err := p.planRules(rules, false); err != nil {
+				if _, err := p.planRules(rules); err != nil {
 					return err
 				}
 			}
@@ -1757,34 +1753,16 @@ func (p *Planner) planRefData(virtual *ruletrie, base *baseptr, ref ast.Ref, ind
 		var vchild *ruletrie
 		var rules []*ast.Rule
 
-		// If there's any non-ground key among the vchild.Children, like
-		// p[x] and p.a (x being non-ground), we'll collect all 'p' rules,
-		// plan them.
-		anyKeyNonGround := false
-
 		if virtual != nil {
 
 			vchild = virtual.Get(ref[index].Value)
-
-			for _, key := range vchild.Children() {
-				if !key.IsGround() {
-					anyKeyNonGround = true
-					break
-				}
-			}
-			if anyKeyNonGround {
-				for _, key := range vchild.Children() {
-					rules = append(rules, vchild.Get(key).Rules()...)
-				}
-			} else {
-				rules = vchild.Rules() // hit or miss
-			}
+			rules = vchild.Rules() // hit or miss
 		}
 
 		if len(rules) > 0 {
 			p.ltarget = p.newOperand()
 
-			funcName, err := p.planRules(rules, anyKeyNonGround)
+			funcName, err := p.planRules(rules)
 			if err != nil {
 				return err
 			}
@@ -1922,7 +1900,7 @@ func (p *Planner) planRefDataExtent(virtual *ruletrie, base *baseptr, iter plani
 				rules = append(rules, virtual.Get(key).Rules()...)
 			}
 
-			funcName, err := p.planRules(rules, true)
+			funcName, err := p.planRules(rules)
 			if err != nil {
 				return err
 			}
@@ -1963,7 +1941,7 @@ func (p *Planner) planRefDataExtent(virtual *ruletrie, base *baseptr, iter plani
 					// Generate virtual document for leaf.
 					lvalue := p.newLocal()
 
-					funcName, err := p.planRules(rules, false)
+					funcName, err := p.planRules(rules)
 					if err != nil {
 						return err
 					}
@@ -2317,7 +2295,7 @@ func (p *Planner) optimizeLookup(t *ruletrie, ref ast.Ref) ([][]*ast.Rule, []ir.
 	for _, node := range nodes {
 		// we're done with ref, check if there's only ruleset leaves; collect rules
 		if index == len(ref)-1 {
-			if len(node.Children()) > 0 {
+			if len(node.Rules()) == 0 && len(node.Children()) > 0 {
 				p.debugf("no optimization of %s: unbalanced ruletrie", ref)
 				return dont()
 			}
diff --git a/internal/planner/planner_test.go b/internal/planner/planner_test.go
index 2f589fa697..cb22d67448 100644
--- a/internal/planner/planner_test.go
+++ b/internal/planner/planner_test.go
@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"os"
 	"reflect"
+	"strings"
 	"testing"
 
 	"github.com/open-policy-agent/opa/ast"
@@ -433,7 +434,7 @@ func findInPolicy(needle interface{}, loc string, p interface{}) error {
 }
 
 // Assert some selected statements' location mappings. Note that for debugging,
-// it's worthwhile to no use tabs in the multi-line strings, as they may be
+// it's worthwhile to not use tabs in the multi-line strings, as they may be
 // counted differently in the editor vs. in code.
 func TestPlannerLocations(t *testing.T) {
 
@@ -730,14 +731,20 @@ func ref(r string) ast.Ref {
 }
 
 func TestOptimizeLookup(t *testing.T) {
-	r0, r1, r2 := ast.Rule{}, ast.Rule{}, ast.Rule{}
+	r0, r1, r2 := ast.MustParseRule("p = 0 { true }"), ast.MustParseRule("p = 1 { true }"), ast.MustParseRule("p = 2 { true }")
+	planner := func() *Planner {
+		if testing.Verbose() {
+			return New().WithDebug(os.Stderr)
+		}
+		return New()
+	}
 
 	t.Run("seen variable (last), one ruleset", func(t *testing.T) {
 		r := newRuletrie()
 		val := r.LookupOrInsert(ref("foo.bar"))
-		val.rules = append(val.rules, &r0, &r1, &r2)
+		val.rules = append(val.rules, r0, r1, r2)
 
-		p := New()
+		p := planner()
 		l := p.newLocal()
 		p.vars.Put(ast.Var("x"), l)
 
@@ -770,9 +777,9 @@ func TestOptimizeLookup(t *testing.T) {
 	t.Run("ref shorter than ruletrie depth", func(t *testing.T) {
 		r := newRuletrie()
 		val := r.LookupOrInsert(ref("foo.bar.baz"))
-		val.rules = append(val.rules, &r0, &r1, &r2)
+		val.rules = append(val.rules, r0, r1, r2)
 
-		p := New()
+		p := planner().WithDebug(os.Stderr)
 		l := p.newLocal()
 		p.vars.Put(ast.Var("x"), l)
 
@@ -785,11 +792,11 @@ func TestOptimizeLookup(t *testing.T) {
 	t.Run("seen variable (last), multiple rulesets", func(t *testing.T) {
 		r := newRuletrie()
 		val := r.LookupOrInsert(ref("foo.bar"))
-		val.rules = append(val.rules, &r0, &r1)
+		val.rules = append(val.rules, r0, r1)
 		val = r.LookupOrInsert(ref("foo.baz"))
-		val.rules = append(val.rules, &r2)
+		val.rules = append(val.rules, r2)
 
-		p := New()
+		p := planner()
 		p.vars.Put(ast.Var("x"), p.newLocal())
 
 		rulesets, _, _, opt := p.optimizeLookup(r, ast.MustParseRef("data.foo[x]"))
@@ -810,9 +817,9 @@ func TestOptimizeLookup(t *testing.T) {
 	t.Run("unseen variable (last)", func(t *testing.T) {
 		r := newRuletrie()
 		val := r.LookupOrInsert(ref("foo.bar"))
-		val.rules = append(val.rules, &r0, &r1, &r2)
+		val.rules = append(val.rules, r0, r1, r2)
 
-		p := New()
+		p := planner()
 
 		_, _, _, opt := p.optimizeLookup(r, ast.MustParseRef("data.foo[x]"))
 		if exp, act := false, opt; exp != act {
@@ -823,9 +830,9 @@ func TestOptimizeLookup(t *testing.T) {
 	t.Run("all ground refs", func(t *testing.T) {
 		r := newRuletrie()
 		val := r.LookupOrInsert(ref("foo.bar.baz"))
-		val.rules = append(val.rules, &r0, &r1, &r2)
+		val.rules = append(val.rules, r0, r1, r2)
 
-		p := New()
+		p := planner()
 
 		_, _, _, opt := p.optimizeLookup(r, ast.MustParseRef("data.foo.bar.baz"))
 		if exp, act := false, opt; exp != act {
@@ -836,9 +843,9 @@ func TestOptimizeLookup(t *testing.T) {
 	t.Run("multiple seen vars, one rule set", func(t *testing.T) {
 		r := newRuletrie()
 		val := r.LookupOrInsert(ref("foo.aaa.bar.bbb.q"))
-		val.rules = append(val.rules, &r0, &r1, &r2)
+		val.rules = append(val.rules, r0, r1, r2)
 
-		p := New()
+		p := planner()
 		lx, ly := p.newLocal(), p.newLocal()
 		p.vars.Put(ast.Var("x"), lx)
 		p.vars.Put(ast.Var("y"), ly)
@@ -865,9 +872,9 @@ func TestOptimizeLookup(t *testing.T) {
 	t.Run("one seen var, one unseen, one rule set", func(t *testing.T) {
 		r := newRuletrie()
 		val := r.LookupOrInsert(ref("foo.aaa.bar.bbb.q"))
-		val.rules = append(val.rules, &r0, &r1, &r2)
+		val.rules = append(val.rules, r0, r1, r2)
 
-		p := New()
+		p := planner()
 		p.vars.Put(ast.Var("x"), p.newLocal())
 
 		_, _, _, opt := p.optimizeLookup(r, ast.MustParseRef("data.foo[x].bar[y].q"))
@@ -879,11 +886,11 @@ func TestOptimizeLookup(t *testing.T) {
 	t.Run("one seen var, one rule set and children left", func(t *testing.T) {
 		r := newRuletrie()
 		val := r.LookupOrInsert(ref("foo.aaa.bar.bbb.q"))
-		val.rules = append(val.rules, &r0)
+		val.rules = append(val.rules, r0)
 		val = r.LookupOrInsert(ref("foo.ccc.bar"))
-		val.rules = append(val.rules, &r1, &r2)
+		val.rules = append(val.rules, r1, r2)
 
-		p := New()
+		p := planner()
 		p.vars.Put(ast.Var("x"), p.newLocal())
 
 		_, _, _, opt := p.optimizeLookup(r, ast.MustParseRef("data.foo[x].bar"))
@@ -895,9 +902,9 @@ func TestOptimizeLookup(t *testing.T) {
 	t.Run("ref goes into the rules' result", func(t *testing.T) {
 		r := newRuletrie()
 		val := r.LookupOrInsert(ref("foo.aaa.bar.q"))
-		val.rules = append(val.rules, &r0, &r1, &r2)
+		val.rules = append(val.rules, r0, r1, r2)
 
-		p := New()
+		p := planner()
 		p.vars.Put(ast.Var("x"), p.newLocal())
 
 		rulesets, path, index, opt := p.optimizeLookup(r, ast.MustParseRef("data.foo[x].bar.q.p.r"))
@@ -922,10 +929,10 @@ func TestOptimizeLookup(t *testing.T) {
 	t.Run("one leaf without rules", func(t *testing.T) {
 		r := newRuletrie()
 		val := r.LookupOrInsert(ref("foo.aaa.bar.q"))
-		val.rules = append(val.rules, &r0, &r1)
+		val.rules = append(val.rules, r0, r1)
 		r.LookupOrInsert(ref("foo.bbb.bar.q"))
 
-		p := New()
+		p := planner()
 		p.vars.Put(ast.Var("x"), p.newLocal())
 
 		rulesets, _, index, opt := p.optimizeLookup(r, ast.MustParseRef("data.foo[x].bar.q"))
@@ -948,7 +955,7 @@ func TestOptimizeLookup(t *testing.T) {
 		r.LookupOrInsert(ref("foo.aaa.bar.q"))
 		r.LookupOrInsert(ref("foo.bbb.bar.q"))
 
-		p := New()
+		p := planner()
 		p.vars.Put(ast.Var("x"), p.newLocal())
 
 		rulesets, _, _, opt := p.optimizeLookup(r, ast.MustParseRef("data.foo[x].bar.q"))
@@ -959,4 +966,252 @@ func TestOptimizeLookup(t *testing.T) {
 			t.Fatalf("expected %d rulesets, got %d\n", exp, act)
 		}
 	})
+
+	t.Run("ref heads, mixed case: string and var last term", func(t *testing.T) {
+		r0, r1 := ast.MustParseRule("b.q.s = 1 { true }"), ast.MustParseRule(`b.q[x] = 2 { x = "t" }`)
+		r := newRuletrie()
+		val := r.LookupOrInsert(ref("a.b.q.s")) // b.q.s = 1 (package a)
+		val.rules = append(val.rules, r0)
+		val = r.LookupOrInsert(ref("a.b.q")) // b.q[x] = 2
+		val.rules = append(val.rules, r1)
+
+		rules := r.Lookup(ref("a.b.q")).Rules()
+		if exp, act := 2, len(rules); exp != act {
+			t.Fatalf("ruletrie: expected %d rules, got %d", exp, act)
+		}
+		if testing.Verbose() {
+			t.Logf("rules: %v", r)
+		}
+
+		p := planner()
+		p.vars.Put(ast.Var("x"), p.newLocal())
+		rulesets, _, _, opt := p.optimizeLookup(r, ast.MustParseRef("data.a[x].q"))
+
+		if exp, act := true, opt; exp != act {
+			t.Errorf("expected 'optimize' %v, got %v\n", exp, act)
+		}
+		if exp, act := 1, len(rulesets); exp != act {
+			t.Fatalf("expected %d rulesets, got %d\n", exp, act)
+		}
+
+		if exp, act := 2, len(rulesets[0]); exp != act {
+			t.Fatalf("expected %d rules in ruleset[0], got %d\n", exp, act)
+		}
+	})
+
+	t.Run("ref heads, mixed case: string and number last term", func(t *testing.T) {
+		r0, r1 := ast.MustParseRule("b.q[1] = 1 { true }"), ast.MustParseRule(`b.q[x] = 2 { x = "t" }`)
+		r := newRuletrie()
+		val := r.LookupOrInsert(ref("a.b.q")) // b.q[1] = 1 (package a)
+		val.rules = append(val.rules, r0)
+		val = r.LookupOrInsert(ref("a.b.q")) // b.q[x] = 2
+		val.rules = append(val.rules, r1)
+
+		rules := r.Lookup(ref("a.b.q")).Rules()
+		if exp, act := 2, len(rules); exp != act {
+			t.Fatalf("ruletrie: expected %d rules, got %d", exp, act)
+		}
+		if testing.Verbose() {
+			t.Logf("rules: %v", r)
+		}
+
+		p := planner()
+		p.vars.Put(ast.Var("x"), p.newLocal())
+		rulesets, _, _, opt := p.optimizeLookup(r, ast.MustParseRef("data.a[x].q"))
+
+		if exp, act := true, opt; exp != act {
+			t.Errorf("expected 'optimize' %v, got %v\n", exp, act)
+		}
+		if exp, act := 1, len(rulesets); exp != act {
+			t.Fatalf("expected %d rulesets, got %d\n", exp, act)
+		}
+
+		if exp, act := 2, len(rulesets[0]); exp != act {
+			t.Fatalf("expected %d rules in ruleset[0], got %d\n", exp, act)
+		}
+	})
+}
+
+func TestPlannerCallDynamic(t *testing.T) {
+	tests := []struct {
+		note    string
+		queries []string
+		modules []string
+		path    []interface{}                // path expected on irCallDynamicStmt, string => string const, int => local
+		where   func(*ir.Policy) interface{} // where to start walking search for `exps`
+		extras  []func(interface{}) error
+	}{
+		{
+			note:    "CallDynamicStmt optimization",
+			queries: []string{`x := "a"; data.test[x] = y`},
+			modules: []string{`package test
+a { true }`},
+			path: []interface{}{"g0", "test", 2},
+			extras: []func(interface{}) error{
+				findFunc("g0.data.test.a", "g0.test.a"),
+			},
+		},
+		{
+			note:    "simple single-val ref head",
+			queries: []string{`x := "a"; data.test.a[x].c = y`},
+			modules: []string{`package test
+a.b.c = 1 { true }`},
+			path: []interface{}{"g0", "test", "a", 2, "c"},
+			extras: []func(interface{}) error{
+				findFunc("g0.data.test.a.b.c", "g0.test.a.b.c"),
+			},
+		},
+		{
+			note:    "two single-val ref heads, string+var",
+			queries: []string{`x := "a"; data.test.a[x] = y`},
+			modules: []string{`package test
+a.b.c = 1 { true }
+a.b[t] = 2 { t := input }`},
+			path: []interface{}{"g0", "test", "a", 2},
+			extras: []func(interface{}) error{
+				findFunc("g0.data.test.a.b", "g0.test.a.b"),
+			},
+		},
+		{
+			note:    "two single-val ref heads, number+var",
+			queries: []string{`x := "a"; data.test.a[x] = y`},
+			modules: []string{`package test
+a.b[1] = 1 { true }
+a.b[t] = 2 { t := input }`},
+			path: []interface{}{"g0", "test", "a", 2},
+			extras: []func(interface{}) error{
+				findFunc("g0.data.test.a.b", "g0.test.a.b"),
+			},
+		},
+		{
+			note:    "one single-val ref head, number",
+			queries: []string{`x := "a"; data.test.a[x] = y`},
+			modules: []string{`package test
+a.b[1] = 1 { true }`},
+			path: []interface{}{"g0", "test", "a", 2},
+			extras: []func(interface{}) error{
+				findFunc("g0.data.test.a.b", "g0.test.a.b"),
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.note, func(t *testing.T) {
+			queries := make([]ast.Body, len(tc.queries))
+			for i := range queries {
+				queries[i] = ast.MustParseBody(tc.queries[i])
+			}
+			modules := make([]*ast.Module, len(tc.modules))
+			for i := range modules {
+				file := fmt.Sprintf("module-%d.rego", i)
+				m, err := ast.ParseModule(file, tc.modules[i])
+				if err != nil {
+					t.Fatal(err)
+				}
+				modules[i] = m
+			}
+			planner := New().WithQueries([]QuerySet{
+				{
+					Name:    "test",
+					Queries: queries,
+				},
+			}).WithModules(modules).WithBuiltinDecls(ast.BuiltinMap)
+			if testing.Verbose() {
+				planner = planner.WithDebug(os.Stderr)
+			}
+			policy, err := planner.Plan()
+			if err != nil {
+				t.Fatal(err)
+			}
+			if testing.Verbose() {
+				err = ir.Pretty(os.Stderr, policy)
+				if err != nil {
+					t.Fatal(err)
+				}
+			}
+			start := interface{}(policy)
+			if tc.where != nil {
+				start = tc.where(policy)
+			}
+
+			if tc.path != nil {
+				exp := make([]ir.Operand, len(tc.path))
+				for i := range tc.path {
+					switch x := tc.path[i].(type) {
+					case string:
+						exp[i] = op(ir.StringIndex(planner.getStringConst(x)))
+					case int:
+						exp[i] = op(ir.Local(x))
+					}
+				}
+				if err := findCallDynamic(exp, start); err != nil {
+					t.Error(err)
+				}
+			}
+
+			if tc.extras == nil {
+				return
+			}
+			for _, e := range tc.extras {
+				if err := e(start); err != nil {
+					t.Error(err)
+				}
+			}
+		})
+	}
+}
+
+type stmtCmpWalker struct {
+	stmt  interface{}
+	found bool // stop comparing after first found needle
+}
+
+func (*stmtCmpWalker) Before(interface{}) {}
+func (*stmtCmpWalker) After(interface{})  {}
+func (w *stmtCmpWalker) Visit(x interface{}) (ir.Visitor, error) {
+	if !w.found {
+		switch s := w.stmt.(type) {
+		case *ir.CallDynamicStmt:
+			c, ok := x.(*ir.CallDynamicStmt)
+			if ok {
+				w.found = true
+				if !reflect.DeepEqual(s.Path, c.Path) {
+					return nil, fmt.Errorf("call dynamic %v: expected path %v, got %v", c, s.Path, c.Path)
+				}
+			}
+		case *ir.Func:
+			f, ok := x.(*ir.Func)
+			if ok && s.Name == f.Name {
+				w.found = true
+				if !reflect.DeepEqual(s.Path, f.Path) {
+					return nil, fmt.Errorf("func %v: expected path %v, got %v", f, s.Path, f.Path)
+				}
+			}
+		}
+	}
+	return w, nil
+}
+
+func findCallDynamic(path []ir.Operand, p interface{}) error {
+	w := &stmtCmpWalker{stmt: &ir.CallDynamicStmt{Path: path}}
+	if err := ir.Walk(w, p); err != nil {
+		return err
+	}
+	if !w.found {
+		return fmt.Errorf("not found")
+	}
+	return nil
+}
+
+func findFunc(name, path string) func(interface{}) error {
+	return func(p interface{}) error {
+		w := &stmtCmpWalker{stmt: &ir.Func{Name: name, Path: strings.Split(path, ".")}}
+		if err := ir.Walk(w, p); err != nil {
+			return err
+		}
+		if !w.found {
+			return fmt.Errorf("not found")
+		}
+		return nil
+	}
 }
diff --git a/internal/planner/rules.go b/internal/planner/rules.go
index 6ff78f4c69..114a25ae2b 100644
--- a/internal/planner/rules.go
+++ b/internal/planner/rules.go
@@ -1,6 +1,7 @@
 package planner
 
 import (
+	"fmt"
 	"sort"
 
 	"github.com/open-policy-agent/opa/ast"
@@ -82,7 +83,27 @@ func (t *ruletrie) Arity() int {
 
 func (t *ruletrie) Rules() []*ast.Rule {
 	if t != nil {
-		return t.rules
+		if t.rules == nil {
+			return nil
+		}
+		rules := make([]*ast.Rule, len(t.rules), len(t.rules)+len(t.children)) // could be too little
+		copy(rules, t.rules)
+
+		// NOTE(sr): We pull in one layer of children: the compiler ensures
+		// that these are the only possible, relevant rule sources for a given
+		// ref: If the trie is what we get for
+		//
+		//     a.b.c  = 1 { ... }
+		//     a.b[x] = 2 { ... }
+		//
+		// and we're retrieving a.b, we want Rules() to include the rule body
+		// of a.b.c.
+		for _, rs := range t.children {
+			if r := rs[len(rs)-1].rules; r != nil {
+				rules = append(rules, r...)
+			}
+		}
+		return rules
 	}
 	return nil
 }
@@ -170,6 +191,38 @@ func (t *ruletrie) Get(k ast.Value) *ruletrie {
 	return nodes[len(nodes)-1]
 }
 
+func (t *ruletrie) DepthFirst(f func(*ruletrie) bool) {
+	if f(t) {
+		return
+	}
+	for _, rules := range t.children {
+		for i := range rules {
+			rules[i].DepthFirst(f)
+		}
+	}
+}
+
+func (t *ruletrie) Depth() int {
+	if len(t.Children()) == 0 {
+		return 0
+	}
+	c := make([]int, 0, len(t.Children()))
+	for _, nodes := range t.children {
+		c = append(c, nodes[len(nodes)-1].Depth())
+	}
+	max := 0
+	for i := range c {
+		if max < c[i] {
+			max = c[i]
+		}
+	}
+	return max + 1
+}
+
+func (t *ruletrie) String() string {
+	return fmt.Sprintf("<ruletrie rules:%v children:%v>", t.rules, t.children)
+}
+
 type functionMocksStack struct {
 	stack []*functionMocksElem
 }
diff --git a/rego/rego.go b/rego/rego.go
index 974261af58..c07e5018e5 100644
--- a/rego/rego.go
+++ b/rego/rego.go
@@ -1328,7 +1328,7 @@ func (r *Rego) Compile(ctx context.Context, opts ...CompileOption) (*CompileResu
 	}
 
 	var queries []ast.Body
-	var modules []*ast.Module
+	modules := make([]*ast.Module, 0, len(r.compiler.Modules))
 
 	if cfg.partial {