Wrap dispatch — how Stage-2 sugars reach both Bases

Stage2Pipeline<Base, L> is one struct. Its sugar surface (Stage2SugarsShared and Stage2SugarsLocal) is one trait per domain. Yet its chain L operates over different node types depending on the Base:

• Stage2Pipeline<TreeishPipeline<…>, L> — chain runs over the user’s N.
• Stage2Pipeline<SeedPipeline<…>, L> — chain runs over SeedNode<N>, because SeedLift prepends the synthetic EntryRoot at run time.

A user-facing closure types at &N. The chain expects &<wrapped N>. Bridging the two is the job of Wrap.

The trait

#![allow(unused)]
fn main() {
/// Type-level dispatch for the chain's input N. Each
/// [`Stage2Base`](super::Stage2Base) declares which `Wrap` it uses;
/// `WrapShared` / `WrapLocal` impls carry the per-domain lift
/// construction.
pub trait Wrap {
    /// The wrapped node type for a given user-facing N.
    type Of<UN: Clone + 'static>: Clone + 'static;
}
}

Two impls:

• Identity::Of<UN> = UN — used by TreeishPipeline-rooted chains.
• SeedWrap::Of<UN> = SeedNode<UN> — used by SeedPipeline-rooted chains.

Stage2Base declares which Wrap a Base uses:

#![allow(unused)]
fn main() {
/// A Stage-1 pipeline that can drive a Stage-2 chain. Carries the
/// `Wrap` selection plus the run-time machinery (pre-lift, root
/// reference, run-input shape).
///
/// Inherits `TreeishSource` so the `(treeish<N>, fold<N, H, R>)` pair is
/// yielded through one canonical path; `with_treeish` is the single
/// place per-base storage shapes are read.
///
/// `PreLift` is intentionally unbounded at the trait level. The
/// `Stage2Pipeline::run` impl adds the `Lift<…, N2 = <Wrap>::Of<N>>`
/// bound at use time; that keeps the supertrait surface free of the
/// `Domain<<Wrap>::Of<N>>` obligation that would otherwise propagate
/// through every site naming `Stage2Base`.
pub trait Stage2Base: TreeishSource + Sized {
    /// Type-level dispatcher for the chain's input N.
    /// `Identity` → `Of<UN> = UN` (treeish-rooted).
    /// `SeedWrap` → `Of<UN> = SeedNode<UN>` (seed-rooted).
    type Wrap: Wrap;

    /// The user-facing N (the type user lambdas type at). Equal to
    /// `Self::N` for every shipped base; kept distinct for
    /// documentation symmetry with the sugar surface, which threads
    /// `UN` as a method-level parameter.
    type UserN: Clone + 'static;

    /// What `.run(...)` accepts as its second argument. Parameterised
    /// by `CurN`, the user-facing N at the chain tip (i.e. after any
    /// `map_n_bi` lifts; `CurN = Self::N` if the chain doesn't change
    /// the user N).
    ///
    /// `Identity`-Wrap bases: `&'i CurN` (a borrowed post-chain root).
    /// `SeedWrap` bases: an owned `(seeds, entry_heap)` pair (the
    /// `CurN` parameter is unused at the value level — `EntryRoot` is
    /// constructible at any inner type).
    type RunInputs<'i, CurN: Clone + 'static>;

    /// The lift composed at the head of the run-time chain.
    /// `IdentityLift` for treeish-rooted, `SeedLift` for seed-rooted.
    /// Pre-lift transforms `(treeish<N>, fold<N,H,R>)` into
    /// `(treeish<Wrap::Of<N>>, fold<Wrap::Of<N>, H, R>)` without
    /// touching H or R.
    ///
    /// Unbounded at the trait level — see the trait-level note.
    /// The `Stage2Pipeline::run` impl adds
    /// `Self::PreLift: Lift<…, N2 = <Wrap>::Of<N>, MapH = H, MapR = R>`
    /// at use time.
    type PreLift;

    /// Build the pre-lift from inputs (consuming the parts of inputs
    /// the lift captures), then yield it together with the executor's
    /// post-chain root reference to the continuation.
    ///
    /// The continuation receives the pre-lift by value (consumed when
    /// applied to the (treeish, fold) pair) and the root by reference,
    /// at the post-chain type `<Self::Wrap as Wrap>::Of<CurN>`. The
    /// reference is valid for the entire duration of `cont`.
    ///
    /// `Identity` case: pre-lift is `IdentityLift`; the root is the
    /// `&CurN` extracted from `inputs`.
    /// `SeedWrap` case: pre-lift is `SeedLift::from_*_grow(...)`,
    /// consuming `inputs.0` (entry seeds) and `inputs.1` (entry heap);
    /// the root is `&SeedNode::entry_root::<CurN>()`, constructed
    /// locally in this frame and alive for `cont`'s lifetime.
    fn provide_run_essentials<CurN: Clone + 'static, T>(
        &self,
        inputs: Self::RunInputs<'_, CurN>,
        cont: impl FnOnce(Self::PreLift,
                          &<Self::Wrap as Wrap>::Of<CurN>) -> T,
    ) -> T;
}
}

So in the type system: <<Self::Base as Stage2Base>::Wrap as Wrap>::Of<UN> is the chain’s input N — equal to UN for treeish-rooted, equal to SeedNode<UN> for seed-rooted. This two-hop projection appears verbatim in every Stage-2 sugar’s signature.

Per-domain build subtraits

Wrap is type-only: it fixes a type family, not how to construct lifts. The constructors live on per-domain subtraits:

#![allow(unused)]
fn main() {
    fn build_wrap_init<UN, H, R, W>(w: W)
        -> ShapeLift<Shared, Self::Of<UN>, H, R, Self::Of<UN>, H, R>
    where
        UN: Clone + Send + Sync + 'static,
        H:  Clone + Send + Sync + 'static,
        R:  Clone + Send + Sync + 'static,
        Self::Of<UN>: Clone + Send + Sync + 'static,
        W: Fn(&UN, &dyn Fn(&UN) -> H) -> H + Send + Sync + 'static;
}

(Plus one method per Stage-2 sugar; see stage2/wrap/shared.rs for the full set, and stage2/wrap/local.rs for the Local mirror.)

The split is forced by the Send + Sync axis: Shared user closures must be Send + Sync (Arc storage; parallel executors); Local must not require it (Rc storage; supports non-Send captured state). WrapShared/WrapLocal are how that single asymmetry is expressed without macros.

Identity: pass-through

#![allow(unused)]
fn main() {
impl WrapShared for Identity {
    fn build_wrap_init<UN, H, R, W>(w: W)
        -> ShapeLift<Shared, UN, H, R, UN, H, R>
    where
        UN: Clone + Send + Sync + 'static,
        H:  Clone + Send + Sync + 'static,
        R:  Clone + Send + Sync + 'static,
        W: Fn(&UN, &dyn Fn(&UN) -> H) -> H + Send + Sync + 'static,
    {
        Shared::wrap_init_lift::<UN, H, R, _>(w)   // pass-through
    }
}

User closure goes straight to Shared::wrap_init_lift. Of<UN> = UN, so no adaptation is needed.

SeedWrap: peel Node, pass EntryRoot

#![allow(unused)]
fn main() {
impl WrapShared for SeedWrap {
    fn build_wrap_init<UN, H, R, W>(w: W)
        -> ShapeLift<Shared, SeedNode<UN>, H, R, SeedNode<UN>, H, R>
    where
        UN: Clone + Send + Sync + 'static,
        H:  Clone + Send + Sync + 'static,
        R:  Clone + Send + Sync + 'static,
        W: Fn(&UN, &dyn Fn(&UN) -> H) -> H + Send + Sync + 'static,
    {
        let user = Arc::new(w);
        // Adapter for the SeedNode<UN>-typed chain: peel Node(_), pass EntryRoot.
        let lifted = move |ln: &SeedNode<UN>,
                           orig: &dyn Fn(&SeedNode<UN>) -> H| -> H
        {
            match sn_int::inner(ln) {
                SeedNodeInner::Node(n) => {
                    let user = user.clone();
                    user(n, &|inner: &UN| orig(&sn_int::node(inner.clone())))
                }
                SeedNodeInner::EntryRoot => orig(ln),
            }
        };
        Shared::wrap_init_lift::<SeedNode<UN>, H, R, _>(lifted)
    }
}

The user types Fn(&UN, …) -> H. The chain expects Fn(&SeedNode<UN>, …) -> H. The body adapts: when the row is Node(n), call the user’s closure with &n; when it’s EntryRoot, call through to the chain’s orig continuation directly (the user closure has nothing to do with the synthetic root).

The same pattern recurs for every N-aware sugar: build_filter_edges, build_memoize_by, build_wrap_visit, build_map_n_bi. Sugars without &N in their signature (wrap_accumulate, wrap_finalize, zipmap, map_r_bi, explain) need no peeling — both impls forward unchanged.

How the sugar trait forwards

A representative Stage2SugarsShared body — the unified surface that covers both Base shapes:

#![allow(unused)]
fn main() {
    fn wrap_init<W>(self, w: W) -> Self::With<ShapeLift<Shared,
        <<Self::Base as Stage2Base>::Wrap as Wrap>::Of<UN>, H, R,
        <<Self::Base as Stage2Base>::Wrap as Wrap>::Of<UN>, H, R>>
    where
        <Self::Base as Stage2Base>::Wrap: WrapShared,
        <<Self::Base as Stage2Base>::Wrap as Wrap>::Of<UN>: Clone + Send + Sync + 'static,
        W: Fn(&UN, &dyn Fn(&UN) -> H) -> H + Send + Sync + 'static,
    {
        self.then_lift(<<Self::Base as Stage2Base>::Wrap as WrapShared>::build_wrap_init::<UN, H, R, _>(w))
    }
}

The body is one line. The surrounding where clauses repeat the projection chain so Rust’s solver can verify each junction; that’s where the verbosity sits. See the type-level deep dive for why the projection has to be spelled out symmetrically here.

What the user sees

Nothing of the above. From the call site:

seed_pipeline
    .lift()
    .wrap_init(|n: &N, orig| orig(n) + 1)   // typed at &N, not &SeedNode<N>
    .filter_edges(|n: &N| !is_excluded(n))
    .run_from_slice(&exec, &seeds, h0);

Wrap dispatch is invisible. The user picks a Base; the trait routes through the right impl; closures stay typed at the user’s N. Switching Base shape — say, building the same chain over a TreeishPipeline — costs no code at the sugar layer.