ast+topdown: parametrized (prefix) external rule sources (#8881)

Let a single registered external source serve an unbounded family of
sub-references by adding an optional ParametrizedExternalRuleIndex
interface whose ParamArity(tail Ref) int the evaluator queries at lookup
time. Fixed arity is the degenerate implementation (return N); the
exact-ref path is unchanged (a source that doesn't implement the
interface has arity 0, with no forced no-op method — same optional-
interface idiom as ExternalRuleIndexCloser).

- A single registered prefix can now back an uneven-depth tree — e.g.
  data.reg.user[k] consuming one key while data.reg.pair[a][b] consumes
  two — by keying off a leading discriminator segment.

- Caching is preserved. The contract constrains the count to depend on
  the tail's *shape*, not on parameter *values*, so it's known before
  Lookup and the cache key (prefix + ground params) is still computable
  up front — no need to call Lookup before knowing the boundary.

For example, registering prefix data.foo with ParamArity(data.foo) 1 makes
data.foo[<key>].<rule> resolve <key> as a parameter, handed to Lookup
via LookupOptions.Params. One source can then serve a distinct set of
rules per key without registering each key concretely, so references
whose key only comes into existence at runtime resolve without a
recompile.

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

v1/ast/compile.go

@@ -4363,10 +4363,16 @@ type ExternalIndex struct {
 // ExternalSourceOptions.DistinguishAbsentFromUnknown distinguish absent input
 // from values that are unknown under partial evaluation.
 //
+// params carries the ground key values that followed the registered prefix for
+// a parametrized source (see ParametrizedExternalRuleIndex); it is nil for
+// conventional sources. The returned subtree is always rooted at prefix (the
+// registered ref), regardless of params — the evaluator layers the parameter
+// levels back on top.
+//
 // Like ExternalIndex, Tree is internal plumbing exported only for the topdown
 // evaluator. It is not part of OPA's supported public API and may change
 // without notice.
-func (ei *ExternalIndex) Tree(ctx context.Context, rt *TreeNode, prefix Ref, resolver ValueResolver, m metrics.Metrics, reqMD map[string]any, respMD map[string]any) (*TreeNode, ExternalRuleIndex, error) {
+func (ei *ExternalIndex) Tree(ctx context.Context, rt *TreeNode, prefix Ref, params []Value, resolver ValueResolver, m metrics.Metrics, reqMD map[string]any, respMD map[string]any) (*TreeNode, ExternalRuleIndex, error) {
 	o := ei.Index.Opts()
 
 	// Select the resolver handed to the source. By default we wrap the caller's
@@ -4387,6 +4393,7 @@ func (ei *ExternalIndex) Tree(ctx context.Context, rt *TreeNode, prefix Ref, res
 		LookupMetrics(m),
 		LookupRequestMetadata(reqMD),
 		LookupResponseMetadata(respMD),
+		LookupParams(params),
 	)
 	if err != nil {
 		return nil, nil, err

v1/ast/external_source.go

@@ -46,6 +46,39 @@ type ExternalRuleIndexCloser interface {
 	Close() error
 }
 
+// ParametrizedExternalRuleIndex is an optional interface implemented by external
+// rule indexes that serve a family of sub-references under their registered
+// prefix rather than a single exact ref. The Ref such a source is registered
+// under is treated as a PREFIX: the leading elements of a query reference that
+// follow the prefix are consumed as ground lookup parameters (handed to Lookup
+// via LookupOptions.Params) rather than as descents into a static rule tree.
+// This lets one registered source serve an unbounded family of sub-references —
+// one distinct set of rules per parameter tuple — without registering each
+// concretely, so references whose key only comes into existence at runtime
+// resolve without a recompile.
+//
+// An index that does not implement this interface behaves as a conventional
+// exact-ref source (equivalent to an arity of 0).
+type ParametrizedExternalRuleIndex interface {
+	ExternalRuleIndex
+
+	// ParamArity reports how many elements following the registered prefix this
+	// index consumes as lookup parameters, given the reference tail (the query
+	// reference elements after the prefix, or an empty Ref when none follow).
+	//
+	// The count may vary with the tail's *shape* — e.g. keying off a leading
+	// discriminator segment — which lets a single prefix back an uneven-depth
+	// tree. It must NOT depend on parameter *values*: ParamArity is consulted
+	// before the parameters are plugged, so the tail may contain non-ground
+	// elements, and the count decides the caching boundary. Returning 0 makes
+	// the reference resolve as a conventional exact ref.
+	//
+	// The parameter elements the count selects must be ground at evaluation
+	// time. A non-ground parameter yields an undefined result, except under
+	// partial evaluation where the reference is saved for residualization.
+	ParamArity(tail Ref) int
+}
+
 // ExternalSourceOptions contains options for registering an external rule source.
 type ExternalSourceOptions struct {
 	// VisibleRefs controls which parts of the surrounding rule tree the external
@@ -98,6 +131,7 @@ type LookupOptions struct {
 	resolver         ValueResolver
 	requestMetadata  map[string]any
 	responseMetadata map[string]any
+	params           []Value
 }
 
 // Metrics returns the metrics instance from the options, or nil if not set.
@@ -120,6 +154,14 @@ func (o *LookupOptions) ResponseMetadata() map[string]any {
 	return o.responseMetadata
 }
 
+// Params returns the parameter values for a parametrized external source (see
+// ParametrizedExternalRuleIndex). The slice holds the ground key values that
+// followed the registered prefix in the query reference, in order. It is empty
+// for conventional (non-parametrized) sources.
+func (o *LookupOptions) Params() []Value {
+	return o.params
+}
+
 // LookupMetrics returns a LookupOption that sets the metrics instance
 // for the Lookup call.
 func LookupMetrics(m metrics.Metrics) LookupOption {
@@ -145,3 +187,11 @@ func LookupResponseMetadata(m map[string]any) LookupOption {
 		opts.responseMetadata = m
 	}
 }
+
+// LookupParams returns a LookupOption that sets the parameter values handed to a
+// parametrized external source (see ParametrizedExternalRuleIndex).
+func LookupParams(params []Value) LookupOption {
+	return func(opts *LookupOptions) {
+		opts.params = params
+	}
+}

v1/ast/external_source_test.go

@@ -144,7 +144,7 @@ func TestExternalSourceResolverDistinguishesAbsentFromUnknown(t *testing.T) {
 	t.Run("opted-in source tells absent from unknown", func(t *testing.T) {
 		idx := &resolverCapturingIndex{distinguish: true, got: map[string]resolveResult{}}
 		ei := &ExternalIndex{Index: idx, Ref: prefix}
-		if _, _, err := ei.Tree(context.Background(), rt, prefix, resolver, nil, nil, nil); err != nil {
+		if _, _, err := ei.Tree(context.Background(), rt, prefix, nil, resolver, nil, nil, nil); err != nil {
 			t.Fatal(err)
 		}
 		if got := idx.got["input.foo"]; !got.unknown {
@@ -161,7 +161,7 @@ func TestExternalSourceResolverDistinguishesAbsentFromUnknown(t *testing.T) {
 	t.Run("default collapses absent into unknown", func(t *testing.T) {
 		idx := &resolverCapturingIndex{distinguish: false, got: map[string]resolveResult{}}
 		ei := &ExternalIndex{Index: idx, Ref: prefix}
-		if _, _, err := ei.Tree(context.Background(), rt, prefix, resolver, nil, nil, nil); err != nil {
+		if _, _, err := ei.Tree(context.Background(), rt, prefix, nil, resolver, nil, nil, nil); err != nil {
 			t.Fatal(err)
 		}
 		if got := idx.got["input.foo"]; !got.unknown {
@@ -178,7 +178,7 @@ func TestExternalSourceResolverDistinguishesAbsentFromUnknown(t *testing.T) {
 	t.Run("nil resolver is treated as all-unknown", func(t *testing.T) {
 		idx := &resolverCapturingIndex{distinguish: true, got: map[string]resolveResult{}}
 		ei := &ExternalIndex{Index: idx, Ref: prefix}
-		if _, _, err := ei.Tree(context.Background(), rt, prefix, nil, nil, nil, nil); err != nil {
+		if _, _, err := ei.Tree(context.Background(), rt, prefix, nil, nil, nil, nil, nil); err != nil {
 			t.Fatal(err)
 		}
 		for _, ref := range []string{"input.foo", "input.bar", "input.baz"} {

v1/topdown/eval.go

@@ -2629,36 +2629,104 @@ func (e evalTree) next(iter unifyIterator, plugged *ast.Term) error {
 				externalRef := node.External.Ref
 				externalIndex := node.External.Index
 
-				// Initialize externalTreeStack if needed
-				if e.e.externalTreeStack == nil {
-					e.e.externalTreeStack = newExternalTreeStack(e.e)
+				// For a parametrized (prefix) external source, the leading
+				// elements of the reference after the prefix are lookup
+				// parameters rather than tree descents. The source reports how
+				// many via ParamArity, keyed off the reference tail's shape (so
+				// one prefix can back an uneven-depth tree). They must be ground;
+				// the resolved sub-tree is cached under the full reference
+				// (prefix + params) so distinct parameters do not collide within
+				// a single evaluation.
+				arity := 0
+				if p, ok := externalIndex.(ast.ParametrizedExternalRuleIndex); ok {
+					arity = p.ParamArity(e.ref[e.pos+1:])
+				}
+				var params []ast.Value
+				cacheRef := externalRef
+				expand := true
+
+				if arity > 0 {
+					params = make([]ast.Value, 0, arity)
+					cacheRef = make(ast.Ref, len(externalRef), len(externalRef)+arity)
+					copy(cacheRef, externalRef)
+					for i := 1; i <= arity; i++ {
+						idx := e.pos + i
+						if idx >= len(e.ref) {
+							expand = false
+							break
+						}
+						k := e.bindings.Plug(e.ref[idx])
+						if !k.IsGround() {
+							expand = false
+							break
+						}
+						params = append(params, k.Value)
+						cacheRef = append(cacheRef, k)
+					}
+
+					if !expand {
+						// The parameter key(s) are not ground, so we cannot
+						// select a concrete sub-source. Under partial evaluation
+						// the reference is unknown and must be residualized;
+						// otherwise it is simply undefined and we fall through
+						// with the bare (rule-less) external node.
+						if e.e.partial() {
+							saved := make(ast.Ref, len(e.ref))
+							for i := range e.ref {
+								saved[i] = e.bindings.Plug(e.ref[i])
+							}
+							return e.e.saveUnify(ast.NewTerm(saved), e.rterm, e.bindings, e.rbindings, iter)
+						}
+					}
 				}
 
-				// Check cache first
-				cachedNode, _, found := e.e.externalTreeStack.findCached(externalRef)
-				if found {
-					node = cachedNode
-				} else {
-					// Call Tree() and cache the result
-					e.e.instr.startTimer(evalOpExternalRuleSource)
-					// Pass the eval itself as the resolver: it is save-set
-					// aware, so external sources that opt into
-					// ExternalSourceOptions.DistinguishAbsentFromUnknown can
-					// tell references that are unknown under partial evaluation
-					// apart from references that are simply absent from the
-					// concrete input.
-					tree, updatedIndex, err := node.External.Tree(e.e.ctx, e.e.compiler.RuleTree, externalRef, e.e, e.e.metrics, e.e.requestMetadata, e.e.responseMetadata)
-					e.e.instr.stopTimer(evalOpExternalRuleSource)
-					if err != nil {
-						return err
+				if expand {
+					// Initialize externalTreeStack if needed
+					if e.e.externalTreeStack == nil {
+						e.e.externalTreeStack = newExternalTreeStack(e.e)
 					}
-					if tree != nil {
-						if updatedIndex != nil {
-							externalIndex = updatedIndex
+
+					// Check cache first
+					cachedNode, _, found := e.e.externalTreeStack.findCached(cacheRef)
+					var tree *ast.TreeNode
+					if found {
+						tree = cachedNode
+					} else {
+						// Call Tree() and cache the result.
+						e.e.instr.startTimer(evalOpExternalRuleSource)
+						// Pass the eval itself as the resolver: it is save-set
+						// aware, so external sources that opt into
+						// ExternalSourceOptions.DistinguishAbsentFromUnknown can
+						// tell references that are unknown under partial evaluation
+						// apart from references that are simply absent from the
+						// concrete input. The parameter terms (params) select the
+						// concrete sub-source for a parametrized prefix.
+						t, updatedIndex, err := node.External.Tree(e.e.ctx, e.e.compiler.RuleTree, externalRef, params, e.e, e.e.metrics, e.e.requestMetadata, e.e.responseMetadata)
+						e.e.instr.stopTimer(evalOpExternalRuleSource)
+						if err != nil {
+							return err
+						}
+						if t != nil {
+							if updatedIndex != nil {
+								externalIndex = updatedIndex
+							}
+							e.e.externalTreeStack.Push(cacheRef, t, externalIndex, e.e.input)
+							pushedExternalTree = true
+						}
+						tree = t
+					}
+
+					if tree != nil {
+						if arity > 0 {
+							// The resolved sub-tree is rooted at the prefix, but
+							// the walk still has to consume the parameter
+							// element(s). Re-insert them as ordinary tree levels
+							// so the descent below (and any further descent into
+							// rules) proceeds unchanged.
+							node = wrapExternalParams(cacheRef[len(externalRef):], tree)
+						} else {
+							node = tree
 						}
-						e.e.externalTreeStack.Push(externalRef, tree, externalIndex, e.e.input)
-						node = tree
-						pushedExternalTree = true
 					}
 				}
 			}
@@ -4917,6 +4985,21 @@ func (e *eval) updateSavedMocks(withs []*ast.With) []*ast.With {
 	return ret
 }
 
+// wrapExternalParams re-inserts the parameter levels consumed by a
+// parametrized external source (see ast.ParametrizedExternalRuleIndex) on
+// top of the resolved sub-tree, which is rooted at the registered prefix. The
+// evaluator's descent can then consume the parameter element(s) as ordinary
+// tree levels. keys are the ground parameter terms in reference order.
+func wrapExternalParams(keys []*ast.Term, tree *ast.TreeNode) *ast.TreeNode {
+	node := tree
+	for i := len(keys) - 1; i >= 0; i-- {
+		node = &ast.TreeNode{
+			Children: map[ast.Value]*ast.TreeNode{keys[i].Value: node},
+		}
+	}
+	return node
+}
+
 // simpleTreeNode provides minimal tree structure for navigation
 type simpleTreeNode struct {
 	tree     *ast.TreeNode

v1/topdown/external_source_test.go

@@ -3,6 +3,7 @@ package topdown
 import (
 	"context"
 	"errors"
+	"fmt"
 	"slices"
 	"sync/atomic"
 	"testing"
@@ -404,6 +405,304 @@ check if data.authz.allow`)
 	}
 }
 
+// paramExternalSource is a parametrized (prefix) external source: it is
+// registered under a prefix ref (e.g. data.directory.user) with ParamArity 1, and
+// ParametrizedExternalRuleIndex with arity 1, and
+// synthesizes a distinct module per key on each Lookup from the parameter term.
+// It stands in for a real source that serves a different set of rules per key.
+type paramExternalSource struct {
+	refs      []ast.Ref
+	arity     int
+	callCount int32
+	keys      []string
+}
+
+func (s *paramExternalSource) Refs() []ast.Ref { return s.refs }
+
+func (s *paramExternalSource) getCallCount() int { return int(atomic.LoadInt32(&s.callCount)) }
+
+func (s *paramExternalSource) Init(_ context.Context, ref ast.Ref) (ast.ExternalRuleIndex, error) {
+	return &paramExternalIndex{prefix: ref, arity: s.arity, src: s}, nil
+}
+
+type paramExternalIndex struct {
+	prefix ast.Ref
+	arity  int
+	src    *paramExternalSource
+}
+
+func (*paramExternalIndex) Opts() *ast.ExternalSourceOptions {
+	return &ast.ExternalSourceOptions{}
+}
+
+// ParamArity declares this source parametrized: the registered prefix is
+// followed by idx.arity key segment(s) consumed as lookup parameters.
+func (idx *paramExternalIndex) ParamArity(ast.Ref) int {
+	return idx.arity
+}
+
+func (idx *paramExternalIndex) Lookup(_ context.Context, opts ...ast.LookupOption) ([]*ast.Rule, ast.ExternalRuleIndex, error) {
+	atomic.AddInt32(&idx.src.callCount, 1)
+
+	o := ast.LookupOptions{}
+	for _, opt := range opts {
+		opt(&o)
+	}
+	params := o.Params()
+	if len(params) != idx.arity {
+		return nil, nil, fmt.Errorf("expected %d param(s), got %d", idx.arity, len(params))
+	}
+	key, ok := params[0].(ast.String)
+	if !ok {
+		return nil, nil, fmt.Errorf("expected string key, got %v", params[0])
+	}
+	idx.src.keys = append(idx.src.keys, string(key))
+
+	// The prefix (e.g. data.directory.user) is the Rego package of the synthesized
+	// module; its rules echo the key and gate on input.
+	pkgPath := idx.prefix.String()[len("data."):]
+	mod := ast.MustParseModule(fmt.Sprintf(
+		"package %s\nid := %q\nallow if input.account == %q",
+		pkgPath, string(key), string(key)))
+	return mod.Rules, nil, nil
+}
+
+func TestExternalSourceParametrizedDistinctKeys(t *testing.T) {
+	t.Parallel()
+
+	prefix := ast.MustParseRef("data.directory.user")
+	source := &paramExternalSource{refs: []ast.Ref{prefix}, arity: 1}
+
+	// Two distinct keys in a single evaluation must not collide.
+	staticModule := ast.MustParseModule(`package main
+check if {
+	data.directory.user["123"].id == "123"
+	data.directory.user["456"].id == "456"
+}`)
+
+	compiler := setupCompiler(t, prefix, source, staticModule)
+
+	qrs := runQuery(t, compiler, "data.main.check", nil)
+	if len(qrs) != 1 {
+		t.Fatalf("Expected 1 result, got %d", len(qrs))
+	}
+	if got := source.getCallCount(); got != 2 {
+		t.Errorf("Expected external source to be called twice (once per distinct key), got %d (keys=%v)", got, source.keys)
+	}
+}
+
+func TestExternalSourceParametrizedWrongKeyUndefined(t *testing.T) {
+	t.Parallel()
+
+	prefix := ast.MustParseRef("data.directory.user")
+	source := &paramExternalSource{refs: []ast.Ref{prefix}, arity: 1}
+
+	// The key's module echoes its own id; asking for a mismatched id is undefined.
+	staticModule := ast.MustParseModule(`package main
+check if data.directory.user["123"].id == "999"`)
+
+	compiler := setupCompiler(t, prefix, source, staticModule)
+
+	qrs := runQuery(t, compiler, "data.main.check", nil)
+	if len(qrs) != 0 {
+		t.Errorf("Expected 0 results for mismatched id, got %d", len(qrs))
+	}
+}
+
+func TestExternalSourceParametrizedInputGate(t *testing.T) {
+	t.Parallel()
+
+	prefix := ast.MustParseRef("data.directory.user")
+	source := &paramExternalSource{refs: []ast.Ref{prefix}, arity: 1}
+
+	staticModule := ast.MustParseModule(`package main
+check if data.directory.user["777"].allow`)
+
+	compiler := setupCompiler(t, prefix, source, staticModule)
+
+	t.Run("input matches key", func(t *testing.T) {
+		input := ast.MustParseTerm(`{"account": "777"}`)
+		qrs := runQuery(t, compiler, "data.main.check", input)
+		if len(qrs) != 1 {
+			t.Errorf("Expected 1 result, got %d", len(qrs))
+		}
+	})
+
+	t.Run("input does not match key", func(t *testing.T) {
+		input := ast.MustParseTerm(`{"account": "000"}`)
+		qrs := runQuery(t, compiler, "data.main.check", input)
+		if len(qrs) != 0 {
+			t.Errorf("Expected 0 results, got %d", len(qrs))
+		}
+	})
+}
+
+func TestExternalSourceParametrizedDynamicKeyFromInput(t *testing.T) {
+	t.Parallel()
+
+	prefix := ast.MustParseRef("data.directory.user")
+	source := &paramExternalSource{refs: []ast.Ref{prefix}, arity: 1}
+
+	// The key itself comes from input and is ground at eval time.
+	staticModule := ast.MustParseModule(`package main
+check if data.directory.user[input.account].id == input.account`)
+
+	compiler := setupCompiler(t, prefix, source, staticModule)
+
+	input := ast.MustParseTerm(`{"account": "555"}`)
+	qrs := runQuery(t, compiler, "data.main.check", input)
+	if len(qrs) != 1 {
+		t.Fatalf("Expected 1 result, got %d", len(qrs))
+	}
+	if len(source.keys) == 0 || source.keys[len(source.keys)-1] != "555" {
+		t.Errorf("Expected source to be queried with key 555, got keys=%v", source.keys)
+	}
+}
+
+// unevenExternalSource is a parametrized (prefix) external source whose arity
+// varies with the reference tail, so a single registered prefix (data.reg)
+// serves sub-references of different nesting depth: data.reg.user[<k>] nests one
+// level while data.reg.pair[<a>][<b>] nests two. This is only expressible
+// because ParamArity is consulted at eval time with the reference tail — a
+// compile-time constant, uniform across the prefix subtree, could not describe
+// an uneven-depth tree under one prefix.
+type unevenExternalSource struct {
+	refs  []ast.Ref
+	calls int32
+}
+
+func (s *unevenExternalSource) Refs() []ast.Ref { return s.refs }
+
+func (s *unevenExternalSource) getCallCount() int { return int(atomic.LoadInt32(&s.calls)) }
+
+func (s *unevenExternalSource) Init(_ context.Context, ref ast.Ref) (ast.ExternalRuleIndex, error) {
+	return &unevenExternalIndex{prefix: ref, src: s}, nil
+}
+
+type unevenExternalIndex struct {
+	prefix ast.Ref
+	src    *unevenExternalSource
+}
+
+func (*unevenExternalIndex) Opts() *ast.ExternalSourceOptions { return &ast.ExternalSourceOptions{} }
+
+// ParamArity keys off the leading tail segment to decide how many elements are
+// consumed as parameters: "user" nests one level deep, "pair" nests two.
+func (*unevenExternalIndex) ParamArity(tail ast.Ref) int {
+	if len(tail) == 0 {
+		return 0
+	}
+	switch tail[0].Value {
+	case ast.String("user"):
+		return 2 // "user", <k>
+	case ast.String("pair"):
+		return 3 // "pair", <a>, <b>
+	default:
+		return 0
+	}
+}
+
+func (idx *unevenExternalIndex) Lookup(_ context.Context, opts ...ast.LookupOption) ([]*ast.Rule, ast.ExternalRuleIndex, error) {
+	atomic.AddInt32(&idx.src.calls, 1)
+
+	o := ast.LookupOptions{}
+	for _, opt := range opts {
+		opt(&o)
+	}
+	params := o.Params()
+
+	// Derive a value from the consumed parameters, distinct per depth.
+	var value string
+	switch len(params) {
+	case 2: // data.reg.user[<k>] -> params are "user", <k>
+		value = string(params[1].(ast.String))
+	case 3: // data.reg.pair[<a>][<b>] -> params are "pair", <a>, <b>
+		value = fmt.Sprintf("%s-%s", string(params[1].(ast.String)), string(params[2].(ast.String)))
+	default:
+		return nil, nil, fmt.Errorf("unexpected param count %d", len(params))
+	}
+
+	// The synthesized module is rooted at the registered prefix (data.reg); the
+	// evaluator layers the consumed parameter levels back on top, so id resolves
+	// at data.reg.user[<k>].id or data.reg.pair[<a>][<b>].id respectively.
+	pkgPath := idx.prefix.String()[len("data."):]
+	mod := ast.MustParseModule(fmt.Sprintf("package %s\nid := %q", pkgPath, value))
+	return mod.Rules, nil, nil
+}
+
+var _ ast.ParametrizedExternalRuleIndex = (*unevenExternalIndex)(nil)
+
+func TestExternalSourceParametrizedUnevenDepth(t *testing.T) {
+	t.Parallel()
+
+	prefix := ast.MustParseRef("data.reg")
+
+	for _, tc := range []struct {
+		note   string
+		module string
+		calls  int
+	}{
+		{
+			note: "one level deep (data.reg.user[k])",
+			module: `package main
+check if data.reg.user["u1"].id == "u1"`,
+			calls: 1,
+		},
+		{
+			note: "two levels deep (data.reg.pair[a][b])",
+			module: `package main
+check if data.reg.pair["a"]["b"].id == "a-b"`,
+			calls: 1,
+		},
+		{
+			note: "both depths under one prefix in a single evaluation",
+			module: `package main
+check if {
+	data.reg.user["u1"].id == "u1"
+	data.reg.pair["a"]["b"].id == "a-b"
+}`,
+			calls: 2,
+		},
+	} {
+		t.Run(tc.note, func(t *testing.T) {
+			source := &unevenExternalSource{refs: []ast.Ref{prefix}}
+			compiler := setupCompiler(t, prefix, source, ast.MustParseModule(tc.module))
+
+			qrs := runQuery(t, compiler, "data.main.check", nil)
+			if len(qrs) != 1 {
+				t.Fatalf("Expected 1 result, got %d", len(qrs))
+			}
+			if got := source.getCallCount(); got != tc.calls {
+				t.Errorf("Expected %d lookup(s), got %d", tc.calls, got)
+			}
+		})
+	}
+}
+
+func TestExternalSourceParametrizedUnevenDepthInsufficientDepthUndefined(t *testing.T) {
+	t.Parallel()
+
+	prefix := ast.MustParseRef("data.reg")
+	source := &unevenExternalSource{refs: []ast.Ref{prefix}}
+
+	// "pair" declares arity 3 (the prefix followed by two keys), but this
+	// reference supplies only one key. There aren't enough elements to
+	// parametrize the source, so the reference is undefined and the source is
+	// never consulted.
+	staticModule := ast.MustParseModule(`package main
+check if data.reg.pair["a"]`)
+
+	compiler := setupCompiler(t, prefix, source, staticModule)
+
+	qrs := runQuery(t, compiler, "data.main.check", nil)
+	if len(qrs) != 0 {
+		t.Errorf("Expected 0 results for insufficient depth, got %d", len(qrs))
+	}
+	if got := source.getCallCount(); got != 0 {
+		t.Errorf("Expected source not to be consulted, got %d lookup(s)", got)
+	}
+}
+
 func TestExternalSourceE2EWithInputOverrideNilInput(t *testing.T) {
 	t.Parallel()