experiment: per-module fan-out in compiler stages

Signed-off-by: Stephan Renatus <[email protected]>

Author: srenatus

go.mod

@@ -1,6 +1,6 @@
 module github.com/open-policy-agent/opa
 
-go 1.24.6
+go 1.25
 
 require (
 	github.com/agnivade/levenshtein v1.2.1

v1/ast/compile.go

@@ -13,6 +13,8 @@ import (
 	"sort"
 	"strconv"
 	"strings"
+	"sync"
+	"sync/atomic"
 
 	"github.com/open-policy-agent/opa/internal/debug"
 	"github.com/open-policy-agent/opa/internal/gojsonschema"
@@ -1756,6 +1758,15 @@ func (c *Compiler) init() {
 	c.initialized = true
 }
 
+func (c *Compiler) forEachModule(f func(mod *Module)) {
+	wg := &sync.WaitGroup{}
+	for _, name := range c.sorted {
+		wg.Go(func() { f(c.Modules[name]) })
+	}
+	wg.Wait()
+}
+
+// TODO(sr): Fix this, it's not concurrency-safe. And the panic is bad form, too.
 func (c *Compiler) err(err *Error) {
 	if c.maxErrs > 0 && len(c.Errors) >= c.maxErrs {
 		c.Errors = append(c.Errors, errLimitReached)
@@ -1957,7 +1968,7 @@ func (c *Compiler) resolveAllRefs() {
 
 func (c *Compiler) removeImports() {
 	c.imports = make(map[string][]*Import, len(c.Modules))
-	for name := range c.Modules {
+	for name := range c.Modules { // Trivial. No fan-out for this.
 		c.imports[name] = c.Modules[name].Imports
 		c.Modules[name].Imports = nil
 	}
@@ -1969,27 +1980,25 @@ func (c *Compiler) initLocalVarGen() {
 
 func (c *Compiler) rewriteComprehensionTerms() {
 	f := newEqualityFactory(c.localvargen)
-	for _, name := range c.sorted {
-		mod := c.Modules[name]
+	c.forEachModule(func(mod *Module) {
 		_, _ = rewriteComprehensionTerms(f, mod) // ignore error
-	}
+	})
 }
 
 func (c *Compiler) rewriteExprTerms() {
-	for _, name := range c.sorted {
-		mod := c.Modules[name]
+	c.forEachModule(func(mod *Module) {
 		WalkRules(mod, func(rule *Rule) bool {
 			rewriteExprTermsInHead(c.localvargen, rule)
 			rule.Body = rewriteExprTermsInBody(c.localvargen, rule.Body)
 			return false
 		})
-	}
+	})
 }
 
 func (c *Compiler) rewriteRuleHeadRefs() {
 	f := newEqualityFactory(c.localvargen)
-	for _, name := range c.sorted {
-		WalkRules(c.Modules[name], func(rule *Rule) bool {
+	c.forEachModule(func(mod *Module) {
+		WalkRules(mod, func(rule *Rule) bool {
 
 			ref := rule.Head.Ref()
 			// NOTE(sr): We're backfilling Refs here -- all parser code paths would have them, but
@@ -2042,16 +2051,15 @@ func (c *Compiler) rewriteRuleHeadRefs() {
 
 			return true
 		})
-	}
+	})
 }
 
 func (c *Compiler) checkVoidCalls() {
-	for _, name := range c.sorted {
-		mod := c.Modules[name]
+	c.forEachModule(func(mod *Module) {
 		for _, err := range checkVoidCalls(c.TypeEnv, mod) {
 			c.err(err)
 		}
-	}
+	})
 }
 
 func (c *Compiler) rewritePrintCalls() {
@@ -2063,8 +2071,7 @@ func (c *Compiler) rewritePrintCalls() {
 			}
 		}
 	} else {
-		for _, name := range c.sorted {
-			mod := c.Modules[name]
+		c.forEachModule(func(mod *Module) {
 			WalkRules(mod, func(r *Rule) bool {
 				safe := r.Head.Args.Vars()
 				safe.Update(ReservedVars)
@@ -2082,7 +2089,7 @@ func (c *Compiler) rewritePrintCalls() {
 				WalkBodies(r.Body, vis)
 				return false
 			})
-		}
+		})
 	}
 	if modified {
 		c.Required.addBuiltinSorted(Print)
@@ -2295,8 +2302,7 @@ func isPrintCall(x *Expr) bool {
 // p[__local0__] { i < 100; __local0__ = {"foo": data.foo[i]} }
 func (c *Compiler) rewriteRefsInHead() {
 	f := newEqualityFactory(c.localvargen)
-	for _, name := range c.sorted {
-		mod := c.Modules[name]
+	c.forEachModule(func(mod *Module) {
 		WalkRules(mod, func(rule *Rule) bool {
 			if requiresEval(rule.Head.Key) {
 				expr := f.Generate(rule.Head.Key)
@@ -2317,27 +2323,27 @@ func (c *Compiler) rewriteRefsInHead() {
 			}
 			return false
 		})
-	}
+	})
 }
 
 func (c *Compiler) rewriteEquals() {
 	modified := false
-	for _, name := range c.sorted {
-		modified = rewriteEquals(c.Modules[name]) || modified
-	}
+	c.forEachModule(func(mod *Module) {
+		modified = rewriteEquals(mod) || modified
+	})
 	if modified {
 		c.Required.addBuiltinSorted(Equal)
 	}
 }
 
 func (c *Compiler) rewriteDynamicTerms() {
 	f := newEqualityFactory(c.localvargen)
-	for _, name := range c.sorted {
-		WalkRules(c.Modules[name], func(rule *Rule) bool {
+	c.forEachModule(func(mod *Module) {
+		WalkRules(mod, func(rule *Rule) bool {
 			rule.Body = rewriteDynamics(f, rule.Body)
 			return false
 		})
-	}
+	})
 }
 
 // rewriteTestRuleEqualities rewrites equality expressions in test rule bodies to create local vars for statements that would otherwise
@@ -2366,39 +2372,34 @@ func (c *Compiler) rewriteTestRuleEqualities() {
 	}
 
 	f := newEqualityFactory(c.localvargen)
-	for _, name := range c.sorted {
-		mod := c.Modules[name]
+	c.forEachModule(func(mod *Module) {
 		WalkRules(mod, func(rule *Rule) bool {
 			if strings.HasPrefix(string(rule.Head.Name), "test_") {
 				rule.Body = rewriteTestEqualities(f, rule.Body)
 			}
 			return false
 		})
-	}
+	})
 }
 
 func (c *Compiler) parseMetadataBlocks() {
 	// Only parse annotations if rego.metadata built-ins are called
 	regoMetadataCalled := false
-	for _, name := range c.sorted {
-		mod := c.Modules[name]
+	c.forEachModule(func(mod *Module) {
+		if regoMetadataCalled {
+			return
+		}
 		WalkExprs(mod, func(expr *Expr) bool {
 			if isRegoMetadataChainCall(expr) || isRegoMetadataRuleCall(expr) {
 				regoMetadataCalled = true
 			}
 			return regoMetadataCalled
 		})
-
-		if regoMetadataCalled {
-			break
-		}
-	}
+	})
 
 	if regoMetadataCalled {
 		// NOTE: Possible optimization: only parse annotations for modules on the path of rego.metadata-calling module
-		for _, name := range c.sorted {
-			mod := c.Modules[name]
-
+		c.forEachModule(func(mod *Module) {
 			if len(mod.Annotations) == 0 {
 				var errs Errors
 				mod.Annotations, errs = parseAnnotations(mod.Comments)
@@ -2409,7 +2410,7 @@ func (c *Compiler) parseMetadataBlocks() {
 
 				attachRuleAnnotations(mod)
 			}
-		}
+		})
 	}
 }
 
@@ -2419,9 +2420,7 @@ func (c *Compiler) rewriteRegoMetadataCalls() {
 	_, chainFuncAllowed := c.builtins[RegoMetadataChain.Name]
 	_, ruleFuncAllowed := c.builtins[RegoMetadataRule.Name]
 
-	for _, name := range c.sorted {
-		mod := c.Modules[name]
-
+	c.forEachModule(func(mod *Module) {
 		WalkRules(mod, func(rule *Rule) bool {
 			var firstChainCall *Expr
 			var firstRuleCall *Expr
@@ -2499,7 +2498,7 @@ func (c *Compiler) rewriteRegoMetadataCalls() {
 
 			return false
 		})
-	}
+	})
 }
 
 func getPrimaryRuleAnnotations(as *AnnotationSet, rule *Rule) *Annotations {
@@ -4333,27 +4332,27 @@ const LocalVarPrefix = "__local"
 type localVarGenerator struct {
 	exclude VarSet
 	suffix  string
-	next    int
+	next    *atomic.Int32
 }
 
 func newLocalVarGeneratorForModuleSet(sorted []string, modules map[string]*Module) *localVarGenerator {
 	vis := NewVarVisitor()
 	for _, key := range sorted {
 		vis.Walk(modules[key])
 	}
-	return &localVarGenerator{exclude: vis.vars, next: 0}
+	return &localVarGenerator{exclude: vis.vars, next: &atomic.Int32{}}
 }
 
 func newLocalVarGenerator(suffix string, node any) *localVarGenerator {
 	vis := NewVarVisitor()
 	vis.Walk(node)
-	return &localVarGenerator{exclude: vis.vars, suffix: suffix, next: 0}
+	return &localVarGenerator{exclude: vis.vars, suffix: suffix, next: &atomic.Int32{}}
 }
 
 func (l *localVarGenerator) Generate() Var {
 	for {
-		result := Var(LocalVarPrefix + l.suffix + strconv.Itoa(l.next) + "__")
-		l.next++
+		next := l.next.Add(1) - 1 // we want the old number
+		result := Var(LocalVarPrefix + l.suffix + strconv.Itoa(int(next)) + "__")
 		if !l.exclude.Contains(result) {
 			return result
 		}