test+eval: add helper to smuggle compiler through context

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

cmd/eval.go

@@ -473,7 +473,9 @@ func evalOnce(ctx context.Context, ectx *evalContext) pr.Output {
 	if ectx.profiler != nil {
 		ectx.profiler.reset()
 	}
-	r := rego.New(ectx.regoArgs...)
+	r := rego.New(append(ectx.regoArgs, rego.CompilerHook(func(c *ast.Compiler) {
+		ctx = ast.WithCompiler(ctx, c)
+	}))...)
 
 	if !ectx.params.partial {
 		var pq rego.PreparedEvalQuery

v1/ast/rego_compiler.go

@@ -0,0 +1,17 @@
+package ast
+
+import "context"
+
+type regoCompileCtx struct{}
+
+func WithCompiler(ctx context.Context, c *Compiler) context.Context {
+	return context.WithValue(ctx, regoCompileCtx{}, c)
+}
+
+func CompilerFromContext(ctx context.Context) (*Compiler, bool) {
+	if ctx == nil {
+		return nil, false
+	}
+	v, ok := ctx.Value(regoCompileCtx{}).(*Compiler)
+	return v, ok
+}

v1/rego/rego.go

@@ -668,6 +668,8 @@ type Rego struct {
 	plugins                     []TargetPlugin
 	targetPrepState             TargetPluginEval
 	regoVersion                 ast.RegoVersion
+	compilerHook                func(*ast.Compiler)
+	evalMode                    *ast.CompilerEvalMode
 }
 
 func (r *Rego) RegoVersion() ast.RegoVersion {
@@ -1320,6 +1322,21 @@ func SetRegoVersion(version ast.RegoVersion) func(r *Rego) {
 	}
 }
 
+// CompilerHook sets a hook function that will be called after the compiler is initialized.
+// This is only called if the compiler has not been provided already.
+func CompilerHook(hook func(*ast.Compiler)) func(r *Rego) {
+	return func(r *Rego) {
+		r.compilerHook = hook
+	}
+}
+
+// EvalMode lets you override the evaluation mode.
+func EvalMode(mode ast.CompilerEvalMode) func(r *Rego) {
+	return func(r *Rego) {
+		r.evalMode = &mode
+	}
+}
+
 // New returns a new Rego object.
 func New(options ...func(r *Rego)) *Rego {
 
@@ -1336,6 +1353,8 @@ func New(options ...func(r *Rego)) *Rego {
 		option(r)
 	}
 
+	callHook := r.compiler == nil // call hook only if we created the compiler here
+
 	if r.compiler == nil {
 		r.compiler = ast.NewCompiler().
 			WithUnsafeBuiltins(r.unsafeBuiltins).
@@ -1404,6 +1423,14 @@ func New(options ...func(r *Rego)) *Rego {
 		r.compiler = r.compiler.WithEvalMode(ast.EvalModeIR)
 	}
 
+	if r.evalMode != nil {
+		r.compiler = r.compiler.WithEvalMode(*r.evalMode)
+	}
+
+	if r.compilerHook != nil && callHook {
+		r.compilerHook(r.compiler)
+	}
+
 	return r
 }
 

v1/rego/rego_test.go

@@ -3371,7 +3371,7 @@ func TestDescriptionRegisterBuiltinDyn(t *testing.T) {
 		),
 	}
 
-	RegisterBuiltinDyn(decl, func(_ BuiltinContext, _ []*ast.Term) (*ast.Term, error) {
+	RegisterBuiltinDyn(decl, func(BuiltinContext, []*ast.Term) (*ast.Term, error) {
 		return ast.StringTerm("bar"), nil
 	})
 	defer unregisterBuiltin("foo")
@@ -3385,12 +3385,67 @@ func TestDescriptionRegisterBuiltinDyn(t *testing.T) {
 // unregisterBuiltin removes the builtin of the given name from ast.Builtins. This assists in
 // cleaning up custom functions added as part of certain test cases.
 func unregisterBuiltin(name string) {
-	builtins := make([]*ast.Builtin, 0, len(ast.Builtins))
-	for _, builtin := range ast.Builtins {
-		if builtin.Name == name {
-			continue
-		}
-		builtins = append(builtins, builtin)
-	}
-	ast.Builtins = builtins
+	ast.Builtins = slices.DeleteFunc(ast.Builtins, func(b *ast.Builtin) bool { return b.Name == name })
+}
+
+func TestCompilerContextViaRegoModuleBuiltin(t *testing.T) {
+	moduleSource := `package test
+
+result := test.module("policy.rego")
+`
+
+	t.Run("compiler not passed", func(t *testing.T) {
+		ctx := context.Background()
+		r := New(
+			Query("data.test.result"),
+			CompilerHook(func(c *ast.Compiler) { ctx = ast.WithCompiler(ctx, c) }),
+			Module("policy.rego", moduleSource),
+			Function1(&Function{
+				Name: "test.module",
+				Decl: types.NewFunction(types.Args(types.S), types.S),
+			}, func(bctx BuiltinContext, a *ast.Term) (*ast.Term, error) {
+				moduleName, ok := a.Value.(ast.String)
+				if !ok {
+					return nil, fmt.Errorf("bad arg type: %T", a.Value)
+				}
+
+				comp, ok := ast.CompilerFromContext(bctx.Context)
+				if !ok {
+					return nil, errors.New("no compiler on context")
+				}
+
+				return ast.StringTerm(comp.Modules[string(moduleName)].String()), nil
+			}),
+		)
+
+		rs, err := r.Eval(ctx)
+		if err != nil {
+			t.Fatalf("rego Eval error: %v", err)
+		}
+		if len(rs) == 0 || len(rs[0].Expressions) == 0 {
+			t.Fatalf("No results")
+		}
+		got := rs[0].Expressions[0].Value
+		want := "package test\n\nresult := __local0__ if { true; test.module(\"policy.rego\", __local0__) }"
+		if got != want {
+			t.Errorf("Expected %q, got %q", want, got)
+		}
+	})
+
+	t.Run("compiler passed in", func(t *testing.T) { // when the compiler is passed, no hook is run
+		ctx := context.Background()
+		r := New(
+			Compiler(ast.NewCompiler()),
+			Query("data.test.result"),
+			CompilerHook(func(*ast.Compiler) { t.Fatal("unexpected hook call") }),
+			Module("policy.rego", "package test\nresult:=true"),
+		)
+		rs, err := r.Eval(ctx)
+		if err != nil {
+			t.Fatalf("rego Eval error: %v", err)
+		}
+		if act, exp := rs.Allowed(), true; exp != act {
+			t.Errorf("expected %v, got %v", exp, act)
+		}
+	})
 }

v1/tester/runner.go

@@ -937,6 +937,8 @@ func (r *Runner) runTest(ctx context.Context, txn storage.Transaction, mod *ast.
 		v.Func(rg)
 	}
 
+	ctx = ast.WithCompiler(ctx, r.compiler)
+
 	t0 := time.Now()
 	rs, err := rg.Eval(ctx)
 	dt := time.Since(t0)