ast: add scaffolding to introspect and skip compiler stages (#8304)

* ast: allow skipping stages in compiler
* ast: add package function to return all stages
* ast: restructure loops (metrics vs no-metrics)
* ast: add withOnlyStagesUpTo() helper, use in tests

   I thought about exporting this, but I wasn't certain. We can do this
   later, if we urgently want to call this method from the outside.

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

Author: srenatus

v1/ast/compile.go

@@ -158,6 +158,8 @@ type Compiler struct {
 	evalMode                   CompilerEvalMode              //
 	rewriteTestRulesForTracing bool                          // rewrite test rules to capture dynamic values for tracing.
 	defaultRegoVersion         RegoVersion
+	skipStages                 map[StageID]struct{} // stages to skip during compilation
+	plan                       *executionPlan       // computed execution plan (cached)
 }
 
 func (c *Compiler) DefaultRegoVersion() RegoVersion {
@@ -167,6 +169,87 @@ func (c *Compiler) DefaultRegoVersion() RegoVersion {
 // CompilerStage defines the interface for stages in the compiler.
 type CompilerStage func(*Compiler) *Error
 
+// StageID uniquely identifies a compiler stage.
+type StageID string
+
+// Compiler stage identifiers.
+const (
+	StageResolveRefs                StageID = "ResolveRefs"
+	StageInitLocalVarGen            StageID = "InitLocalVarGen"
+	StageRewriteRuleHeadRefs        StageID = "RewriteRuleHeadRefs"
+	StageCheckKeywordOverrides      StageID = "CheckKeywordOverrides"
+	StageCheckDuplicateImports      StageID = "CheckDuplicateImports"
+	StageRemoveImports              StageID = "RemoveImports"
+	StageSetModuleTree              StageID = "SetModuleTree"
+	StageSetRuleTree                StageID = "SetRuleTree"
+	StageRewriteLocalVars           StageID = "RewriteLocalVars"
+	StageRewriteTemplateStrings     StageID = "RewriteTemplateStrings"
+	StageCheckVoidCalls             StageID = "CheckVoidCalls"
+	StageRewritePrintCalls          StageID = "RewritePrintCalls"
+	StageRewriteExprTerms           StageID = "RewriteExprTerms"
+	StageParseMetadataBlocks        StageID = "ParseMetadataBlocks"
+	StageSetAnnotationSet           StageID = "SetAnnotationSet"
+	StageRewriteRegoMetadataCalls   StageID = "RewriteRegoMetadataCalls"
+	StageSetGraph                   StageID = "SetGraph"
+	StageRewriteComprehensionTerms  StageID = "RewriteComprehensionTerms"
+	StageRewriteRefsInHead          StageID = "RewriteRefsInHead"
+	StageRewriteWithValues          StageID = "RewriteWithValues"
+	StageCheckRuleConflicts         StageID = "CheckRuleConflicts"
+	StageCheckUndefinedFuncs        StageID = "CheckUndefinedFuncs"
+	StageCheckSafetyRuleHeads       StageID = "CheckSafetyRuleHeads"
+	StageCheckSafetyRuleBodies      StageID = "CheckSafetyRuleBodies"
+	StageRewriteEquals              StageID = "RewriteEquals"
+	StageRewriteDynamicTerms        StageID = "RewriteDynamicTerms"
+	StageRewriteTestRulesForTracing StageID = "RewriteTestRulesForTracing"
+	StageCheckRecursion             StageID = "CheckRecursion"
+	StageCheckTypes                 StageID = "CheckTypes"
+	StageCheckUnsafeBuiltins        StageID = "CheckUnsafeBuiltins"
+	StageCheckDeprecatedBuiltins    StageID = "CheckDeprecatedBuiltins"
+	StageBuildRuleIndices           StageID = "BuildRuleIndices"
+	StageBuildComprehensionIndices  StageID = "BuildComprehensionIndices"
+	StageBuildRequiredCapabilities  StageID = "BuildRequiredCapabilities"
+)
+
+// AllStages returns the complete list of compiler stages in execution order.
+func AllStages() []StageID {
+	return []StageID{
+		StageResolveRefs,
+		StageInitLocalVarGen,
+		StageRewriteRuleHeadRefs,
+		StageCheckKeywordOverrides,
+		StageCheckDuplicateImports,
+		StageRemoveImports,
+		StageSetModuleTree,
+		StageSetRuleTree,
+		StageRewriteLocalVars,
+		StageRewriteTemplateStrings,
+		StageCheckVoidCalls,
+		StageRewritePrintCalls,
+		StageRewriteExprTerms,
+		StageParseMetadataBlocks,
+		StageSetAnnotationSet,
+		StageRewriteRegoMetadataCalls,
+		StageSetGraph,
+		StageRewriteComprehensionTerms,
+		StageRewriteRefsInHead,
+		StageRewriteWithValues,
+		StageCheckRuleConflicts,
+		StageCheckUndefinedFuncs,
+		StageCheckSafetyRuleHeads,
+		StageCheckSafetyRuleBodies,
+		StageRewriteEquals,
+		StageRewriteDynamicTerms,
+		StageRewriteTestRulesForTracing,
+		StageCheckRecursion,
+		StageCheckTypes,
+		StageCheckUnsafeBuiltins,
+		StageCheckDeprecatedBuiltins,
+		StageBuildRuleIndices,
+		StageBuildComprehensionIndices,
+		StageBuildRequiredCapabilities,
+	}
+}
+
 // CompilerEvalMode allows toggling certain stages that are only
 // needed for certain modes, Concretely, only "topdown" mode will
 // have the compiler build comprehension and rule indices.
@@ -189,6 +272,18 @@ type CompilerStageDefinition struct {
 	Stage      CompilerStage
 }
 
+// executionPlan represents the complete ordered list of stages to execute.
+type executionPlan struct {
+	stages []plannedStage
+}
+
+// plannedStage represents a single stage in the execution plan.
+type plannedStage struct {
+	name       string
+	metricName string
+	f          func()
+}
+
 // RulesOptions defines the options for retrieving rules by Ref from the
 // compiler.
 type RulesOptions struct {
@@ -406,9 +501,34 @@ func (c *Compiler) WithPathConflictsCheckRoots(rootPaths []string) *Compiler {
 // the named stage.
 func (c *Compiler) WithStageAfter(after string, stage CompilerStageDefinition) *Compiler {
 	c.after[after] = append(c.after[after], stage)
+	c.plan = nil // invalidate cached plan
+	return c
+}
+
+// WithSkipStages configures the compiler to skip the specified stages during
+// compilation. This invalidates any cached execution plan.
+func (c *Compiler) WithSkipStages(stages ...StageID) *Compiler {
+	if c.skipStages == nil {
+		c.skipStages = make(map[StageID]struct{}, len(stages))
+	}
+	for _, s := range stages {
+		c.skipStages[s] = struct{}{}
+	}
+	c.plan = nil // invalidate cached plan
 	return c
 }
 
+// withOnlyStagesUpTo configures the compiler to run only stages up to and
+// including the specified target stage. All stages after the target will be skipped.
+func (c *Compiler) withOnlyStagesUpTo(target StageID) *Compiler {
+	allStages := AllStages()
+	i := slices.Index(allStages, target)
+	if i == -1 {
+		return c
+	}
+	return c.WithSkipStages(allStages[i+1:]...)
+}
+
 // WithMetrics will set a metrics.Metrics and be used for profiling
 // the Compiler instance.
 func (c *Compiler) WithMetrics(metrics metrics.Metrics) *Compiler {
@@ -906,6 +1026,61 @@ func (c *Compiler) WithDefaultRegoVersion(regoVersion RegoVersion) *Compiler {
 	return c
 }
 
+// buildExecutionPlan creates the unified list of stages to execute, including
+// both main stages and "after" stages, with filtering applied.
+func (c *Compiler) buildExecutionPlan() *executionPlan {
+	plan := &executionPlan{
+		stages: make([]plannedStage, 0, len(c.stages)*2),
+	}
+
+	for _, s := range c.stages {
+		if _, skip := c.skipStages[StageID(s.name)]; skip {
+			continue
+		}
+
+		plan.stages = append(plan.stages, plannedStage(s))
+
+		for _, a := range c.after[s.name] {
+			if _, skip := c.skipStages[StageID(a.Name)]; skip {
+				continue
+			}
+
+			afterStage := a // Capture variables in closure properly
+			plan.stages = append(plan.stages, plannedStage{
+				name:       afterStage.Name,
+				metricName: afterStage.MetricName,
+				f: func() {
+					if err := afterStage.Stage(c); err != nil {
+						c.err(err)
+					}
+				},
+			})
+		}
+	}
+
+	return plan
+}
+
+// getOrBuildPlan ensures we have a valid execution plan.
+func (c *Compiler) getOrBuildPlan() *executionPlan {
+	if c.plan == nil {
+		c.plan = c.buildExecutionPlan()
+	}
+	return c.plan
+}
+
+// StagesToRun returns the list of stage IDs that will be executed during
+// compilation, in execution order. This includes both main stages and any
+// registered "after" stages.
+func (c *Compiler) StagesToRun() []StageID {
+	plan := c.getOrBuildPlan()
+	result := make([]StageID, len(plan.stages))
+	for i, s := range plan.stages {
+		result[i] = StageID(s.name)
+	}
+	return result
+}
+
 func (c *Compiler) counterAdd(name string, n uint64) {
 	if c.metrics == nil {
 		return
@@ -1659,42 +1834,22 @@ func (c *Compiler) checkDeprecatedBuiltins() {
 	}
 }
 
-func (c *Compiler) runStage(metricName string, f func()) {
-	if c.metrics != nil {
-		c.metrics.Timer(metricName).Start()
-		defer c.metrics.Timer(metricName).Stop()
-	}
-	f()
-}
+func (c *Compiler) compile() {
+	plan := c.getOrBuildPlan()
 
-func (c *Compiler) runStageAfter(metricName string, s CompilerStage) *Error {
 	if c.metrics != nil {
-		c.metrics.Timer(metricName).Start()
-		defer c.metrics.Timer(metricName).Stop()
-	}
-	return s(c)
-}
-
-func (c *Compiler) compile() {
-	for _, s := range c.stages {
-		if c.evalMode == EvalModeIR {
-			switch s.name {
-			case "BuildRuleIndices", "BuildComprehensionIndices":
-				continue // skip these stages
+		for _, s := range plan.stages {
+			c.metrics.Timer(s.metricName).Start()
+			s.f()
+			c.metrics.Timer(s.metricName).Stop()
+			if c.Failed() {
+				return
 			}
 		}
-
-		if c.allowUndefinedFuncCalls && (s.name == "CheckUndefinedFuncs" || s.name == "CheckSafetyRuleBodies") {
-			continue
-		}
-
-		c.runStage(s.metricName, s.f)
-		if c.Failed() {
-			return
-		}
-		for _, a := range c.after[s.name] {
-			if err := c.runStageAfter(a.MetricName, a.Stage); err != nil {
-				c.err(err)
+	} else {
+		for _, s := range plan.stages {
+			s.f()
+			if c.Failed() {
 				return
 			}
 		}
@@ -1766,6 +1921,14 @@ func (c *Compiler) init() {
 		WithInputType(c.inputType).
 		Env(c.builtins)
 
+	// Configure default stage skips based on existing configuration
+	if c.evalMode == EvalModeIR {
+		c.WithSkipStages(StageBuildRuleIndices, StageBuildComprehensionIndices)
+	}
+	if c.allowUndefinedFuncCalls {
+		c.WithSkipStages(StageCheckUndefinedFuncs, StageCheckSafetyRuleBodies)
+	}
+
 	c.initialized = true
 }