ast: add DefaultModuleLoader

This provides the means to inject modules _for all invocations of the
ast package in an executable_. It's handy if you want to include your
own set of base modules, with helper functions.

Since the module loader can be provided via other means, we need to
check both possibilities:

1. If there hasn't been a module loader set up before, just use
   the default module loader.

2. If there had been one set up before, run that first, then run
   the default one, and merge the results. This can be iterated,
   of course, if need be -- and should reach a fix point given
   that the individual loaders do.

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

v1/ast/compile.go

@@ -1681,6 +1681,31 @@ func (c *Compiler) init() {
 		return
 	}
 
+	if defaultModuleLoader != nil {
+		if c.moduleLoader == nil {
+			c.moduleLoader = defaultModuleLoader
+		} else {
+			first := c.moduleLoader
+			c.moduleLoader = func(res map[string]*Module) (map[string]*Module, error) {
+				res0, err := first(res)
+				if err != nil {
+					return nil, err
+				}
+				res1, err := defaultModuleLoader(res)
+				if err != nil {
+					return nil, err
+				}
+				// merge res1 into res0, based on module "file" names, to avoid clashes
+				for k, v := range res1 {
+					if _, ok := res0[k]; !ok {
+						res0[k] = v
+					}
+				}
+				return res0, nil
+			}
+		}
+	}
+
 	if c.capabilities == nil {
 		c.capabilities = CapabilitiesForThisVersion()
 	}

v1/ast/compile_test.go

@@ -11500,3 +11500,63 @@ func TestCompile_DefaultRegoVersion(t *testing.T) {
 		})
 	}
 }
+
+func TestCompilerInitWithDefaultModuleLoader(t *testing.T) {
+	// Reset the global variable after the test
+	defer func() { defaultModuleLoader = nil }()
+
+	// a dummy loader that adds "foo"
+	loader1 := func(res map[string]*Module) (map[string]*Module, error) {
+		mod := MustParseModule(`package foo`)
+		resCopy := map[string]*Module{}
+		maps.Copy(resCopy, res)
+		resCopy["foo.rego"] = mod
+		return resCopy, nil
+	}
+
+	// a dummy loader that adds "bar"
+	loader2 := func(res map[string]*Module) (map[string]*Module, error) {
+		mod := MustParseModule(`package bar`)
+		resCopy := map[string]*Module{}
+		maps.Copy(resCopy, res)
+		resCopy["bar.rego"] = mod
+		return resCopy, nil
+	}
+
+	DefaultModuleLoader(loader2)
+
+	c := NewCompiler().WithModuleLoader(loader1)
+	c.init()
+
+	got, err := c.moduleLoader(make(map[string]*Module))
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	expected := map[string]*Module{
+		"foo.rego": MustParseModule(`package foo`),
+		"bar.rego": MustParseModule(`package bar`),
+	}
+	// check both modules are present
+	for k, v := range expected {
+		gotMod, ok := got[k]
+		if !ok {
+			t.Errorf("expected key %q in result", k)
+			continue
+		}
+		if !reflect.DeepEqual(gotMod, v) {
+			t.Errorf("unexpected module for %q: got %v want %v", k, gotMod, v)
+		}
+	}
+
+	// Now, test defaultModuleLoader only
+	c2 := NewCompiler()
+	c2.init()
+	got2, err := c2.moduleLoader(make(map[string]*Module))
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if _, ok := got2["bar.rego"]; !ok {
+		t.Error("expected bar.rego from defaultModuleLoader in result")
+	}
+}

v1/ast/default_module_loader.go

@@ -0,0 +1,14 @@
+// Copyright 2025 The OPA Authors.  All rights reserved.
+// Use of this source code is governed by an Apache2
+// license that can be found in the LICENSE file.
+
+package ast
+
+var defaultModuleLoader ModuleLoader
+
+// DefaultModuleLoader lets you inject an `ast.ModuleLoader` that will
+// always be used. If another one is provided with the ast package,
+// they will both be consulted to enrich the set of modules dynamically.
+func DefaultModuleLoader(ml ModuleLoader) {
+	defaultModuleLoader = ml
+}