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"
@@ -1031,10 +1033,8 @@ func (c *Compiler) buildRequiredCapabilities() {
 	c.Required.FutureKeywords = util.KeysSorted(keywords)
 
 	// extract required features from modules
-
 	for _, name := range c.sorted {
 		mod := c.Modules[name]
-
 		if c.moduleIsRegoV1(mod) {
 			features[FeatureRegoV1] = struct{}{}
 		} else {
@@ -1225,12 +1225,11 @@ func (c *Compiler) checkRuleConflicts() {
 }
 
 func (c *Compiler) checkUndefinedFuncs() {
-	for _, name := range c.sorted {
-		m := c.Modules[name]
-		for _, err := range checkUndefinedFuncs(c.TypeEnv, m, c.GetArity, c.RewrittenVars) {
+	c.forEachModule(func(mod *Module) {
+		for _, err := range checkUndefinedFuncs(c.TypeEnv, mod, c.GetArity, c.RewrittenVars) {
 			c.err(err)
 		}
-	}
+	})
 }
 
 func checkUndefinedFuncs(env *TypeEnv, x any, arity func(Ref) int, rwVars map[Var]Var) Errors {
@@ -1285,17 +1284,16 @@ func arityMismatchError(env *TypeEnv, f Ref, expr *Expr, exp, act int) *Error {
 // positions of built-in expressions will be bound when evaluating the rule from left
 // to right, re-ordering as necessary.
 func (c *Compiler) checkSafetyRuleBodies() {
-	for _, name := range c.sorted {
-		m := c.Modules[name]
-		WalkRules(m, func(r *Rule) bool {
+	c.forEachModule(func(mod *Module) {
+		WalkRules(mod, func(r *Rule) bool {
 			safe := ReservedVars.Copy()
 			if len(r.Head.Args) > 0 {
 				safe.Update(r.Head.Args.Vars())
 			}
 			r.Body = c.checkBodySafety(safe, r.Body)
 			return false
 		})
-	}
+	})
 }
 
 func (c *Compiler) checkBodySafety(safe VarSet, b Body) Body {
@@ -1320,8 +1318,8 @@ var SafetyCheckVisitorParams = VarVisitorParams{
 // checkSafetyRuleHeads ensures that variables appearing in the head of a
 // rule also appear in the body.
 func (c *Compiler) checkSafetyRuleHeads() {
-	for _, name := range c.sorted {
-		WalkRules(c.Modules[name], func(r *Rule) bool {
+	c.forEachModule(func(mod *Module) {
+		WalkRules(mod, func(r *Rule) bool {
 			safe := r.Body.Vars(SafetyCheckVisitorParams)
 			if len(r.Head.Args) > 0 {
 				safe.Update(r.Head.Args.Vars())
@@ -1342,7 +1340,7 @@ func (c *Compiler) checkSafetyRuleHeads() {
 			}
 			return false
 		})
-	}
+	})
 }
 
 func compileSchema(goSchema any, allowNet []string) (*gojsonschema.Schema, error) {
@@ -1609,13 +1607,12 @@ func (c *Compiler) checkUnsafeBuiltins() {
 	if len(c.unsafeBuiltinsMap) == 0 {
 		return
 	}
-
-	for _, name := range c.sorted {
-		errs := checkUnsafeBuiltins(c.unsafeBuiltinsMap, c.Modules[name])
+	c.forEachModule(func(mod *Module) {
+		errs := checkUnsafeBuiltins(c.unsafeBuiltinsMap, mod)
 		for _, err := range errs {
 			c.err(err)
 		}
-	}
+	})
 }
 
 func (c *Compiler) checkDeprecatedBuiltins() {
@@ -1630,15 +1627,14 @@ func (c *Compiler) checkDeprecatedBuiltins() {
 		return
 	}
 
-	for _, name := range c.sorted {
-		mod := c.Modules[name]
+	c.forEachModule(func(mod *Module) {
 		if c.strict || mod.regoV1Compatible() {
 			errs := checkDeprecatedBuiltins(c.deprecatedBuiltinsMap, mod)
 			for _, err := range errs {
 				c.err(err)
 			}
 		}
-	}
+	})
 }
 
 func (c *Compiler) runStage(metricName string, f func()) {
@@ -1756,6 +1752,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 +1962,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 +1974,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 +2045,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 +2065,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 +2083,7 @@ func (c *Compiler) rewritePrintCalls() {
 				WalkBodies(r.Body, vis)
 				return false
 			})
-		}
+		})
 	}
 	if modified {
 		c.Required.addBuiltinSorted(Print)
@@ -2295,8 +2296,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 +2317,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 +2366,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 +2404,7 @@ func (c *Compiler) parseMetadataBlocks() {
 
 				attachRuleAnnotations(mod)
 			}
-		}
+		})
 	}
 }
 
@@ -2419,9 +2414,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 +2492,7 @@ func (c *Compiler) rewriteRegoMetadataCalls() {
 
 			return false
 		})
-	}
+	})
 }
 
 func getPrimaryRuleAnnotations(as *AnnotationSet, rule *Rule) *Annotations {
@@ -2893,8 +2886,7 @@ func (vis *ruleArgLocalRewriter) Visit(x any) Visitor {
 
 func (c *Compiler) rewriteWithModifiers() {
 	f := newEqualityFactory(c.localvargen)
-	for _, name := range c.sorted {
-		mod := c.Modules[name]
+	c.forEachModule(func(mod *Module) {
 		t := NewGenericTransformer(func(x any) (any, error) {
 			body, ok := x.(Body)
 			if !ok {
@@ -2908,7 +2900,7 @@ func (c *Compiler) rewriteWithModifiers() {
 			return body, nil
 		})
 		_, _ = Transform(t, mod) // ignore error
-	}
+	})
 }
 
 func (c *Compiler) setModuleTree() {
@@ -4333,27 +4325,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
 		}