From a2c769fe36dd5cce92b8dc8f5ad09709c2e06d2c Mon Sep 17 00:00:00 2001
From: Stephan Renatus <stephan@styra.com>
Date: Thu, 17 Aug 2023 11:25:32 +0200
Subject: [PATCH] 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               | 25 ++++++++++++++
 v1/ast/compile_test.go          | 60 +++++++++++++++++++++++++++++++++
 v1/ast/default_module_loader.go | 14 ++++++++
 3 files changed, 99 insertions(+)
 create mode 100644 v1/ast/default_module_loader.go

diff --git a/v1/ast/compile.go b/v1/ast/compile.go
index f3ca101735..2ffc0f59a2 100644
--- a/v1/ast/compile.go
+++ b/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()
 	}
diff --git a/v1/ast/compile_test.go b/v1/ast/compile_test.go
index 5c7e01b348..04514957eb 100644
--- a/v1/ast/compile_test.go
+++ b/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")
+	}
+}
diff --git a/v1/ast/default_module_loader.go b/v1/ast/default_module_loader.go
new file mode 100644
index 0000000000..528c253e16
--- /dev/null
+++ b/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
+}