refactor(autoland): merge land runs as a stack through a plan optimization policy

src/stack_pr/autoland/__init__.py

@@ -32,7 +32,9 @@
 - ``native_stack``: merging a run of land steps as one GitHub stack.
 - ``services``: the tools a step uses as it executes (console, GitHub, the
   waits, the repo), and their production bindings.
-- ``engine``: executing a plan.
+- ``optimize``: execution units, the plan optimization policies that group
+  steps into them (e.g. ``--merge-as-stack``), and the registry of policies.
+- ``engine``: optimizing and executing a plan.
 - ``status``: the ``--status`` report.
 - ``commands``: the subparser, and the fresh / resume / replan flows.
 

src/stack_pr/autoland/commands.py

@@ -15,7 +15,6 @@
 from stack_pr.autoland import discovery, engine, runtime
 from stack_pr.autoland.display import _describe_step, _escape_markup, print_status
 from stack_pr.autoland.model import LandingContext, LandStep, PlanStep, StackEntry
-from stack_pr.autoland.native_stack import print_native_stack_runs
 from stack_pr.autoland.options import AutolandOptions
 from stack_pr.autoland.plan import (
     PLAN_SUFFIX,
@@ -417,7 +416,7 @@ def _run_fresh(common: cli.CommonArgs, opts: AutolandOptions) -> None:
         )
 
         print_status(ctx)
-        print_native_stack_runs(ctx, opts)
+        engine.print_execution_preview(ctx, opts)
         if opts.dry_run:
             runtime.console.print("\n[yellow]Dry run — exiting.[/yellow]")
             return
@@ -496,7 +495,7 @@ def _run_resume(common: cli.CommonArgs, opts: AutolandOptions) -> None:
             return
 
         print_status(ctx)
-        print_native_stack_runs(ctx, opts)
+        engine.print_execution_preview(ctx, opts)
         runtime.console.print(f"[dim]State file: {sf_path}[/dim]\n")
         _install_signal_handler(ctx, checkpointer, worktree, opts)
         _finish(
@@ -681,7 +680,7 @@ def _replan(common: cli.CommonArgs, opts: AutolandOptions, state_path: Path) ->
         ctx = LandingContext(stack=stack, plan=plan)
 
         print_status(ctx)
-        print_native_stack_runs(ctx, opts)
+        engine.print_execution_preview(ctx, opts)
         if lost:
             runtime.console.print(
                 "\n[yellow]Done in the previous run, but not carried over (changed, "

src/stack_pr/autoland/engine.py

@@ -1,24 +1,28 @@
-"""Executing a landing plan step by step.
-
-The engine owns the run: it walks the plan from ``ctx.current_step``,
-checkpoints before each step, hands each step a ``StepContext``, records the
-run-level facts a step reports in its ``StepResult``, and stops on an abort.
-What a step does is up to the step (``AutolandStep.execute``).
+"""Executing a landing plan.
+
+The engine owns the run. It optimizes the rest of the plan into execution
+units (``optimize.optimize_plan``), then executes them in order. Through the
+unit's ``UnitContext``, each step the unit runs is made the current step and
+checkpointed before it starts, is handed a ``StepContext``, and has the
+run-level facts it reports in its ``StepResult`` recorded. The run stops on
+the first abort. What a step does is up to the step (``AutolandStep.execute``),
+and what a unit does across its steps is up to the policy that built it.
 """
 
 from __future__ import annotations
 
 import functools
 
 from stack_pr import cli
+from stack_pr.autoland import runtime
 from stack_pr.autoland.display import _describe_step
 from stack_pr.autoland.model import LandingContext
-from stack_pr.autoland.native_stack import land_as_native_stack, native_stack_run
+from stack_pr.autoland.optimize import PolicyRegistry, UnitContext, optimize_plan
 from stack_pr.autoland.options import AutolandOptions
 from stack_pr.autoland.plan import landed_prefix_end
 from stack_pr.autoland.services import Services
 from stack_pr.autoland.state import AutolandCheckpointer
-from stack_pr.autoland.steps import StepContext
+from stack_pr.autoland.steps import StepContext, StepResult
 
 
 def execute_plan(
@@ -28,12 +32,15 @@ def execute_plan(
     checkpointer: AutolandCheckpointer,
     *,
     services: Services | None = None,
+    policies: PolicyRegistry | None = None,
 ) -> bool:
-    """Execute the landing plan from ctx.current_step. Returns True on success."""
+    """Execute the landing plan from ctx.current_step. Returns True on success.
+
+    *policies* defaults to every registered optimization policy.
+    """
     if services is None:
         services = Services.production()
     prefix_end = landed_prefix_end(ctx.plan)
-    native_stack = _NativeStackMerge(ctx, common, opts, checkpointer)
 
     def upcoming(index: int) -> list[str]:
         return [_describe_step(s, ctx) for s in ctx.plan[index + 1 :]]
@@ -43,38 +50,63 @@ def landed_sha() -> str:
             _record_landed_sha(ctx, services.repo.landed_sha(common))
         return ctx.last_landed_sha
 
-    for step_idx in range(ctx.current_step, len(ctx.plan)):
-        step = ctx.plan[step_idx]
-        ctx.current_step = step_idx
+    def checkpoint() -> None:
         checkpointer.save(ctx)
 
-        abort_reason = native_stack.before_step(step_idx)
-        if abort_reason:
-            return _abort(ctx, checkpointer, abort_reason)
+    def enter(index: int) -> None:
+        ctx.current_step = index
+        checkpoint()
 
-        result = step.execute(
+    def execute_step(index: int) -> StepResult:
+        result = ctx.plan[index].execute(
             StepContext(
-                index=step_idx,
+                index=index,
                 total=len(ctx.plan),
                 run=ctx,
                 opts=opts,
                 common=common,
                 services=services,
-                checkpoint=lambda: checkpointer.save(ctx),
-                upcoming=functools.partial(upcoming, step_idx),
+                checkpoint=checkpoint,
+                upcoming=functools.partial(upcoming, index),
                 landed_sha=landed_sha,
-                ends_landed_prefix=step_idx == prefix_end,
+                ends_landed_prefix=index == prefix_end,
             )
         )
         _record_landed_sha(ctx, result.landed_sha)
-        if result.aborted:
-            return _abort(ctx, checkpointer, result.abort_reason)
+        return result
+
+    uc = UnitContext(
+        run=ctx,
+        opts=opts,
+        common=common,
+        services=services,
+        enter=enter,
+        execute_step=execute_step,
+        checkpoint=checkpoint,
+    )
+    for unit in optimize_plan(ctx, opts, policies):
+        abort_reason = unit.execute(uc)
+        if abort_reason:
+            return _abort(ctx, checkpointer, abort_reason)
 
     ctx.current_step = len(ctx.plan)
     checkpointer.save(ctx)
     return True
 
 
+def print_execution_preview(
+    ctx: LandingContext,
+    opts: AutolandOptions,
+    *,
+    policies: PolicyRegistry | None = None,
+) -> None:
+    """Announce what the optimized plan does beyond executing its steps one by
+    one. Derived from the same units ``execute_plan`` executes."""
+    for unit in optimize_plan(ctx, opts, policies):
+        for line in unit.preview():
+            runtime.console.print(line)
+
+
 def _record_landed_sha(ctx: LandingContext, sha: str) -> None:
     # An unknown SHA ("") keeps the last one known.
     if sha:
@@ -88,50 +120,3 @@ def _abort(
     ctx.aborted = True
     checkpointer.save(ctx)
     return False
-
-
-class _NativeStackMerge:
-    """TEMPORARY: merge a run of consecutive land steps as one GitHub stack.
-
-    With ``--merge-as-stack``, before each step the engine asks whether a run
-    of two or more open land steps starts there, and if so lands the run with
-    one stack merge (``native_stack.land_as_native_stack``). The land steps
-    then find their PRs merged and only skip. A run whose stack merge fell
-    through is landed one step at a time instead, without retrying the merge.
-
-    This lives in the engine, not in ``LandStep``, because steps are
-    independent of each other. It is a stopgap: a plan-optimization pass that
-    groups land steps into one execution unit replaces it.
-    """
-
-    def __init__(
-        self,
-        ctx: LandingContext,
-        common: cli.CommonArgs,
-        opts: AutolandOptions,
-        checkpointer: AutolandCheckpointer,
-    ) -> None:
-        self._ctx = ctx
-        self._common = common
-        self._opts = opts
-        self._checkpointer = checkpointer
-        # The last step of a run whose stack merge fell through.
-        self._unstacked_through = -1
-
-    def before_step(self, step_idx: int) -> str:
-        """Stack-merge the run starting at *step_idx*, if any. Returns the
-        reason the plan must stop, or ""."""
-        if not self._opts.merge_as_stack or step_idx <= self._unstacked_through:
-            return ""
-        run_entries = native_stack_run(self._ctx, step_idx)
-        if not run_entries:
-            return ""
-        result = land_as_native_stack(
-            run_entries, ctx=self._ctx, common=self._common, opts=self._opts
-        )
-        if result.abort_reason:
-            return result.abort_reason
-        if not result.landed:
-            self._unstacked_through = step_idx + len(run_entries) - 1
-        self._checkpointer.save(self._ctx)
-        return ""

src/stack_pr/autoland/native_stack.py

@@ -1,4 +1,8 @@
-"""Merging a run of consecutive land steps as one GitHub stack."""
+"""Merging a run of consecutive land steps as one GitHub stack.
+
+This is GitHub's side of a stack merge. Which steps one covers is decided by
+the ``optimize.stack_merge`` policy.
+"""
 
 from __future__ import annotations
 
@@ -7,7 +11,7 @@
 
 from stack_pr import cli
 from stack_pr.autoland import gh, runtime, waits
-from stack_pr.autoland.model import LandingContext, LandStep, PRState, StackEntry
+from stack_pr.autoland.model import LandingContext, PRState, StackEntry
 from stack_pr.autoland.options import AutolandOptions
 from stack_pr.autoland.waits import _poll_read
 
@@ -16,48 +20,6 @@
 # ---------------------------------------------------------------------------
 
 
-def native_stack_run(ctx: LandingContext, start: int) -> list[StackEntry]:
-    """The PRs a stack merge starting at plan step *start* would land.
-
-    That is the run of consecutive ``l`` steps from *start* whose PRs are still
-    open. A ``w`` or ``c`` step ends the run: it has to see the PRs above it
-    land separately from the ones below it. A run of fewer than two PRs gains
-    nothing from a stack merge, so it comes back empty.
-    """
-    entries: list[StackEntry] = []
-    for step in ctx.plan[start:]:
-        if not isinstance(step, LandStep) or step.already_landed:
-            break
-        entry = ctx.stack[step.entry_index]
-        if entry.state == PRState.MERGED:
-            break
-        entries.append(entry)
-    return entries if len(entries) > 1 else []
-
-
-def native_stack_runs(ctx: LandingContext) -> list[tuple[int, list[StackEntry]]]:
-    """Every ``(first step, PRs)`` run the rest of the plan would stack-merge."""
-    runs = []
-    index = ctx.current_step
-    while index < len(ctx.plan):
-        entries = native_stack_run(ctx, index)
-        if entries:
-            runs.append((index, entries))
-        index += max(len(entries), 1)
-    return runs
-
-
-def print_native_stack_runs(ctx: LandingContext, opts: AutolandOptions) -> None:
-    if not opts.merge_as_stack:
-        return
-    for first, entries in native_stack_runs(ctx):
-        prs = ", ".join(f"#{e.pr_number}" for e in entries)
-        runtime.console.print(
-            f"[dim]Steps {first + 1}-{first + len(entries)} ({prs}) will merge "
-            "together as a GitHub stack.[/dim]"
-        )
-
-
 def _open_native_stack_prs(native_stack: dict) -> list[int]:
     """The numbers of a GitHub stack's unmerged PRs, bottom first."""
     return [

src/stack_pr/autoland/optimize/__init__.py

@@ -0,0 +1,35 @@
+"""Plan optimization: execution units and the policies that build them.
+
+Importing this package imports every policy module, so each policy is in the
+registry before a plan is executed.
+"""
+
+from __future__ import annotations
+
+from stack_pr.autoland.optimize.base import (
+    POLICIES,
+    ExecutionUnit,
+    OptimizationPolicy,
+    PolicyRegistry,
+    StepUnit,
+    UnitContext,
+    optimize_plan,
+    register_policy,
+)
+from stack_pr.autoland.optimize.stack_merge import (
+    NativeStackMergePolicy,
+    StackMergeUnit,
+)
+
+__all__ = [
+    "POLICIES",
+    "ExecutionUnit",
+    "NativeStackMergePolicy",
+    "OptimizationPolicy",
+    "PolicyRegistry",
+    "StackMergeUnit",
+    "StepUnit",
+    "UnitContext",
+    "optimize_plan",
+    "register_policy",
+]