fuel_graph/
scan.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//! `unroll_scan`: materialize a bounded [`crate::Op::Scan`] into real
3//! primitive nodes on demand. Used as (a) the FKC/Spec-B numeric oracle and
4//! (b) the fallback lowering for a backend without a scan kernel. NOT
5//! registered as anyone's `.decompose` — `Op::Scan` is a bare primitive that
6//! stays terminal in the base map.
7
8use std::collections::HashMap;
9
10use crate::{Graph, Node, NodeId, Op, ScanEmit, ScanRole};
11
12/// Validate every `Op::ScanPlaceholder` reachable from the scan body's exit
13/// `roots` (`[body_new_carry, body_y]`, optionally the predicate) has an
14/// in-range index: v1 is **single-carry** (`Carry` index must be 0), and every
15/// `Elem` index must address one of the `n_xs` per-step slices. This keeps the
16/// later `clone_body_node`/`build_scan_step` `elem[index]` access infallible by
17/// construction — so a malformed body is a typed **build-time** `Err` (from the
18/// `NodeHandle::scan` / `NodeHandle::scan_until` builders) rather than a mid-realize
19/// `elem[index]` panic on the forward-driver path. Short immutable borrow only.
20pub(crate) fn validate_scan_body_placeholders(
21    graph: &Graph,
22    roots: &[NodeId],
23    n_carries: usize,
24    n_xs: usize,
25) -> std::result::Result<(), fuel_ir::Error> {
26    let reachable = crate::topo_order_multi(graph, roots);
27    for &id in &reachable {
28        if let Op::ScanPlaceholder { role, index } = &graph.node(id).op {
29            match *role {
30                // GAP-303: was `*index != 0` ("v1 is single-carry"). The TYPE was
31                // always multi-carry-shaped -- `index` is a `usize` -- so this is
32                // relaxing a CHECK, not widening a type. Range-checked exactly like
33                // the `Elem` arm below; the two are now symmetric.
34                ScanRole::Carry if *index >= n_carries => {
35                    return Err(fuel_ir::Error::Msg(format!(
36                        "scan: body node {} is ScanPlaceholder{{Carry, {index}}} out of range (n_carries = {n_carries})",
37                        id.0,
38                    )).bt());
39                }
40                ScanRole::Elem if *index >= n_xs => {
41                    return Err(fuel_ir::Error::Msg(format!(
42                        "scan: body node {} is ScanPlaceholder{{Elem, {index}}} out of range (n_xs = {n_xs})",
43                        id.0,
44                    )).bt());
45                }
46                _ => {}
47            }
48        }
49    }
50    Ok(())
51}
52
53/// Unroll `steps` iterations of the `Op::Scan` at `scan_id` into primitives.
54///
55/// Returns `(selected, complementary)`: `emit=All` -> `(stacked_ys,
56/// final_carry)`, `emit=Final` -> `(final_carry, stacked_ys)`. `early_exit =
57/// Some` peels the trailing `pred_exit` and IGNORES it — the build-time
58/// backward/oracle unroll differentiates the full static `bound` (spec C3
59/// static-horizon note); the runtime early-exit is a forward-only optimization
60/// driven by the step driver, not this unroll.
61pub fn unroll_scan(
62    graph: &mut Graph,
63    scan_id: NodeId,
64    steps: usize,
65) -> std::result::Result<(Vec<NodeId>, Vec<NodeId>), fuel_ir::Error> {
66    if scan_id.0 >= graph.len() {
67        return Err(fuel_ir::Error::Msg(format!(
68            "unroll_scan: scan_id {} is out of range (graph has {} nodes)",
69            scan_id.0,
70            graph.len(),
71        ))
72        .bt());
73    }
74    // 1. Read the Scan node's params + input layout in a short borrow.
75    let (n_carries, n_xs, bound, emit, has_exit, inputs) = {
76        let n = graph.node(scan_id);
77        match &n.op {
78            Op::Scan {
79                n_carries,
80                n_xs,
81                bound,
82                emit,
83                early_exit,
84            } => (
85                *n_carries,
86                *n_xs,
87                *bound,
88                *emit,
89                early_exit.is_some(),
90                n.inputs.clone(),
91            ),
92            other => {
93                return Err(fuel_ir::Error::Msg(format!(
94                    "unroll_scan: node {} is not an Op::Scan ({})",
95                    scan_id.0,
96                    other.short_name(),
97                ))
98                .bt());
99            }
100        }
101    };
102    if steps == 0 || steps > bound {
103        return Err(fuel_ir::Error::Msg(format!(
104            "unroll_scan: steps {steps} must be in 1..={bound}",
105        ))
106        .bt());
107    }
108    // inputs = [init_carry, xs_0..xs_{n_xs-1}, consts.., body_new_carry, body_y, [pred_exit]]
109    // Trailing slots: body_new_carry + body_y (+ pred_exit when early_exit = Some).
110    // Minimum well-formed layout: init_carry(1) + n_xs + consts(>=0) + n_trailing.
111    // (One short of the trailing slots — reject it here, before the `consts`
112    // slice below can panic with start > end.)
113    // GAP-303: the leading carry BLOCK is `n_carries` wide (was a literal 1) and
114    // the trailing block carries one `body_new_carry` per carry.
115    //
116    // ⚠️ `early_exit` is dead in VALUE (never `Some` on a live Phase-1 path) and
117    // LIVE IN LAYOUT -- it is the `+ has_exit` term here. That is the one place
118    // the Phase-2 field participates in Phase-1 arithmetic.
119    let n_trailing = n_carries + 1 + usize::from(has_exit); // new_carries.., body_y, [pred]
120    if n_carries == 0 {
121        return Err(fuel_ir::Error::Msg(
122            "unroll_scan: Op::Scan with n_carries = 0 has no recurrent state".to_string(),
123        )
124        .bt());
125    }
126    if inputs.len() < n_carries + n_xs + n_trailing {
127        return Err(fuel_ir::Error::Msg(format!(
128            "unroll_scan: malformed Op::Scan inputs — need >= {} (n_carries={n_carries} + n_xs={n_xs} + {n_trailing} trailing), got {}",
129            n_carries + n_xs + n_trailing, inputs.len(),
130        )).bt());
131    }
132    let init_carries: Vec<NodeId> = inputs[0..n_carries].to_vec();
133    let xs: Vec<NodeId> = inputs[n_carries..n_carries + n_xs].to_vec();
134    let consts: Vec<NodeId> = inputs[n_carries + n_xs..inputs.len() - n_trailing].to_vec();
135    let body_new_carries: Vec<NodeId> =
136        inputs[inputs.len() - n_trailing..inputs.len() - n_trailing + n_carries].to_vec();
137    let body_y = inputs[inputs.len() - n_trailing + n_carries];
138    // pred_exit = inputs[inputs.len() - 1] when has_exit — intentionally NOT read; the
139    // build-time backward unroll differentiates the full static `bound` and ignores the
140    // runtime early-exit predicate (spec C3 "static-horizon note").
141    let consts_set: std::collections::HashSet<NodeId> = consts.iter().copied().collect();
142
143    // 2. Validate every ScanPlaceholder reachable from the body's two exit
144    // nodes has an in-range index, BEFORE any cloning/mutation (shared with the
145    // build-time check in NodeHandle::scan / NodeHandle::scan_until).
146    let mut roots: Vec<NodeId> = body_new_carries.clone();
147    roots.push(body_y);
148    validate_scan_body_placeholders(graph, &roots, n_carries, n_xs)?;
149
150    // 3. Validate every xs[i] has a leading (scan-axis) dim >= steps: the
151    // per-step `Slice { dim: 0, start: t, len: 1 }` below needs `t` in range
152    // for every `t in 0..steps`, and needs a dim 0 to slice at all.
153    for (i, &x) in xs.iter().enumerate() {
154        let dims = graph.node(x).shape.dims().to_vec();
155        if dims.is_empty() {
156            return Err(fuel_ir::Error::Msg(format!(
157                "unroll_scan: xs[{i}] (node {}) is rank-0, needs a leading scan-axis of len >= steps ({steps})",
158                x.0,
159            )).bt());
160        }
161        if dims[0] < steps {
162            return Err(fuel_ir::Error::Msg(format!(
163                "unroll_scan: xs[{i}] (node {}) leading dim {} < steps ({steps})",
164                x.0, dims[0],
165            ))
166            .bt());
167        }
168    }
169
170    let mut carries: Vec<NodeId> = init_carries;
171    let mut ys_steps: Vec<NodeId> = Vec::with_capacity(steps);
172
173    for t in 0..steps {
174        // Per-step xs slices: xs[i] sliced at [t, t+1) on scan-axis 0, then
175        // squeezed to drop the step axis -> the ScanPlaceholder{Elem,i} shape.
176        let mut elem: Vec<NodeId> = Vec::with_capacity(n_xs);
177        for &x in &xs {
178            let (x_shape, x_dtype) = {
179                let n = graph.node(x);
180                (n.shape.clone(), n.dtype)
181            };
182            let sliced_dims: Vec<usize> = std::iter::once(1usize)
183                .chain(x_shape.dims().iter().skip(1).copied())
184                .collect();
185            let sl = graph.push(Node {
186                op: Op::Slice {
187                    dim: 0,
188                    start: t,
189                    len: 1,
190                },
191                inputs: vec![x],
192                shape: fuel_ir::Shape::from_dims(&sliced_dims),
193                dtype: x_dtype,
194            });
195            let sq_dims: Vec<usize> = x_shape.dims().iter().skip(1).copied().collect();
196            let sq = graph.push(Node {
197                op: Op::Squeeze { dim: 0 },
198                inputs: vec![sl],
199                shape: fuel_ir::Shape::from_dims(&sq_dims),
200                dtype: x_dtype,
201            });
202            elem.push(sq);
203        }
204
205        // Clone the body subgraph (rooted at {body_new_carry, body_y}),
206        // substituting placeholders + keeping consts shared.
207        let mut subst: HashMap<NodeId, NodeId> = HashMap::new();
208        // ⚠️ Every new carry is computed from THIS step's `carries`, so the whole
209        // block is built before any assignment -- updating in place mid-loop would
210        // let carry k+1 read carry k's NEXT value instead of its current one.
211        let next_carries: Vec<NodeId> = body_new_carries
212            .iter()
213            .map(|&nc| clone_body_node(graph, nc, &carries, &elem, &consts_set, &mut subst))
214            .collect();
215        let y_t = clone_body_node(graph, body_y, &carries, &elem, &consts_set, &mut subst);
216        carries = next_carries;
217        ys_steps.push(y_t);
218    }
219
220    // stacked_ys = Concat(dim 0) of each y_t unsqueezed at dim 0.
221    let mut unsqueezed: Vec<NodeId> = Vec::with_capacity(ys_steps.len());
222    for &y in &ys_steps {
223        let (y_shape, y_dtype) = {
224            let n = graph.node(y);
225            (n.shape.clone(), n.dtype)
226        };
227        let un_dims: Vec<usize> = std::iter::once(1usize)
228            .chain(y_shape.dims().iter().copied())
229            .collect();
230        let un = graph.push(Node {
231            op: Op::Unsqueeze { dim: 0 },
232            inputs: vec![y],
233            shape: fuel_ir::Shape::from_dims(&un_dims),
234            dtype: y_dtype,
235        });
236        unsqueezed.push(un);
237    }
238    let (y0_shape, y0_dtype) = {
239        let n = graph.node(ys_steps[0]);
240        (n.shape.clone(), n.dtype)
241    };
242    let stacked_dims: Vec<usize> = std::iter::once(ys_steps.len())
243        .chain(y0_shape.dims().iter().copied())
244        .collect();
245    let stacked_ys = graph.push(Node {
246        op: Op::Concat { dim: 0 },
247        inputs: unsqueezed,
248        shape: fuel_ir::Shape::from_dims(&stacked_dims),
249        dtype: y0_dtype,
250    });
251
252    // GAP-303: both sides widened to `Vec`. The (selected, complementary)
253    // contract is UNCHANGED -- only the arity is.
254    //
255    // ⚠️ The old `(NodeId, NodeId)` return was a FOURTH site of the single-carry
256    // assumption, encoded in a `pub fn`'s SIGNATURE rather than in a field: with
257    // `emit = Final` and `n_carries > 1` there is no single "selected" output,
258    // and no type could have expressed that while the pair was scalar.
259    Ok(match emit {
260        ScanEmit::All => (vec![stacked_ys], carries),
261        ScanEmit::Final => (carries, vec![stacked_ys]),
262    })
263}
264
265/// Topological copy of a body node, substituting `ScanPlaceholder{Carry,_}` ->
266/// `carry`, `ScanPlaceholder{Elem,i}` -> `elem[i]`, and keeping any node in
267/// `consts_set` shared (not copied). Memoized in `subst`.
268fn clone_body_node(
269    graph: &mut Graph,
270    id: NodeId,
271    carries: &[NodeId],
272    elem: &[NodeId],
273    consts_set: &std::collections::HashSet<NodeId>,
274    subst: &mut HashMap<NodeId, NodeId>,
275) -> NodeId {
276    if let Some(&m) = subst.get(&id) {
277        return m;
278    }
279    if consts_set.contains(&id) {
280        return id;
281    }
282    let (op, in_ids, shape, dtype) = {
283        let n = graph.node(id);
284        (n.op.clone(), n.inputs.clone(), n.shape.clone(), n.dtype)
285    };
286    let mapped = match op {
287        // GAP-303: the two arms are now SYMMETRIC. `Carry` used to ignore
288        // `index` entirely (it was required to be 0); it now indexes `carries`
289        // exactly as `Elem` indexes `elem`. Both are range-checked up front by
290        // `validate_scan_body_placeholders`, so both index safely.
291        Op::ScanPlaceholder {
292            role: ScanRole::Carry,
293            index,
294        } => carries[index],
295        Op::ScanPlaceholder {
296            role: ScanRole::Elem,
297            index,
298        } => elem[index],
299        _ => {
300            let new_inputs: Vec<NodeId> = in_ids
301                .iter()
302                .map(|&c| clone_body_node(graph, c, carries, elem, consts_set, subst))
303                .collect();
304            graph.push(Node {
305                op,
306                inputs: new_inputs,
307                shape,
308                dtype,
309            })
310        }
311    };
312    subst.insert(id, mapped);
313    mapped
314}
315
316/// The parsed input layout of an [`Op::Scan`] node — the fields the early-exit
317/// step driver needs, extracted from the trailing-input encoding.
318pub struct ScanLayout {
319    /// GAP-303: how many carries the leading/trailing blocks are wide.
320    pub n_carries: usize,
321    pub n_xs: usize,
322    pub bound: usize,
323    pub emit: ScanEmit,
324    pub init_carries: Vec<NodeId>,
325    pub xs: Vec<NodeId>,
326    pub consts: Vec<NodeId>,
327    pub body_new_carries: Vec<NodeId>,
328    pub body_y: NodeId,
329    /// `Some` when `early_exit = Some` — the scalar-`U8` convergence predicate.
330    pub pred_exit: Option<NodeId>,
331}
332
333/// One materialized scan step: the post-step carry, the emitted `y`, and the
334/// (optional) realized `stop` predicate node for this step.
335pub struct ScanStep {
336    pub new_carries: Vec<NodeId>,
337    pub y: NodeId,
338    pub stop: Option<NodeId>,
339}
340
341/// Parse an [`Op::Scan`] node's trailing-input layout. Mirrors the parse in
342/// [`unroll_scan`], including the `early_exit`-aware trailing count, but yields
343/// a struct the step driver can drive one step at a time. Returns a typed
344/// `Err` for a non-`Op::Scan` node or a malformed (too-short) input layout.
345pub fn parse_scan_layout(
346    graph: &Graph,
347    scan_id: NodeId,
348) -> std::result::Result<ScanLayout, fuel_ir::Error> {
349    if scan_id.0 >= graph.len() {
350        return Err(fuel_ir::Error::Msg(format!(
351            "parse_scan_layout: scan_id {} is out of range (graph has {} nodes)",
352            scan_id.0,
353            graph.len(),
354        ))
355        .bt());
356    }
357    let n = graph.node(scan_id);
358    let (n_carries, n_xs, bound, emit, has_exit) = match &n.op {
359        Op::Scan {
360            n_carries,
361            n_xs,
362            bound,
363            emit,
364            early_exit,
365        } => (*n_carries, *n_xs, *bound, *emit, early_exit.is_some()),
366        other => {
367            return Err(fuel_ir::Error::Msg(format!(
368                "parse_scan_layout: node {} is not an Op::Scan ({})",
369                scan_id.0,
370                other.short_name(),
371            ))
372            .bt());
373        }
374    };
375    let inputs = &n.inputs;
376    // GAP-303: mirrors `unroll_scan`'s arithmetic exactly -- the two parses must
377    // agree or a step driver and the oracle disagree about the same node.
378    let n_trailing = n_carries + 1 + usize::from(has_exit);
379    if n_carries == 0 {
380        return Err(fuel_ir::Error::Msg(
381            "parse_scan_layout: Op::Scan with n_carries = 0 has no recurrent state".to_string(),
382        )
383        .bt());
384    }
385    if inputs.len() < n_carries + n_xs + n_trailing {
386        return Err(fuel_ir::Error::Msg(format!(
387            "parse_scan_layout: malformed Op::Scan inputs — need >= {} (n_carries={n_carries} + n_xs={n_xs} + {n_trailing} trailing), got {}",
388            n_carries + n_xs + n_trailing, inputs.len(),
389        )).bt());
390    }
391    let init_carries: Vec<NodeId> = inputs[0..n_carries].to_vec();
392    let xs: Vec<NodeId> = inputs[n_carries..n_carries + n_xs].to_vec();
393    let consts: Vec<NodeId> = inputs[n_carries + n_xs..inputs.len() - n_trailing].to_vec();
394    let body_new_carries: Vec<NodeId> =
395        inputs[inputs.len() - n_trailing..inputs.len() - n_trailing + n_carries].to_vec();
396    let body_y = inputs[inputs.len() - n_trailing + n_carries];
397    let pred_exit = has_exit.then(|| inputs[inputs.len() - 1]);
398    Ok(ScanLayout {
399        n_carries,
400        n_xs,
401        bound,
402        emit,
403        init_carries,
404        xs,
405        consts,
406        body_new_carries,
407        body_y,
408        pred_exit,
409    })
410}
411
412/// Materialize one step of a scan at step index `t` with the given `carry`
413/// NodeId. Slices each `xs[i]` at `[t, t+1)` (squeezed), clones `body_new_carry`
414/// / `body_y` / `pred_exit` with a **single shared** substitution map so that a
415/// `pred_exit` referencing `body_new_carry` resolves to *this step's* new carry
416/// (no double-clone). Returns the post-step carry, `y`, and the optional `stop`
417/// predicate node.
418pub fn build_scan_step(
419    graph: &mut Graph,
420    layout: &ScanLayout,
421    t: usize,
422    carries: &[NodeId],
423) -> std::result::Result<ScanStep, fuel_ir::Error> {
424    // per-step elem slices: xs[i] sliced [t,t+1) on axis 0, squeezed.
425    let mut elem: Vec<NodeId> = Vec::with_capacity(layout.n_xs);
426    for &x in &layout.xs {
427        let (x_shape, x_dtype) = {
428            let n = graph.node(x);
429            (n.shape.clone(), n.dtype)
430        };
431        if x_shape.dims().first().is_none_or(|&d0| d0 <= t) {
432            return Err(fuel_ir::Error::Msg(format!(
433                "build_scan_step: xs slice at t={t} out of range for shape {:?}",
434                x_shape.dims()
435            ))
436            .bt());
437        }
438        let sliced: Vec<usize> = std::iter::once(1usize)
439            .chain(x_shape.dims().iter().skip(1).copied())
440            .collect();
441        let sl = graph.push(Node {
442            op: Op::Slice {
443                dim: 0,
444                start: t,
445                len: 1,
446            },
447            inputs: vec![x],
448            shape: fuel_ir::Shape::from_dims(&sliced),
449            dtype: x_dtype,
450        });
451        let sq_dims: Vec<usize> = x_shape.dims().iter().skip(1).copied().collect();
452        let sq = graph.push(Node {
453            op: Op::Squeeze { dim: 0 },
454            inputs: vec![sl],
455            shape: fuel_ir::Shape::from_dims(&sq_dims),
456            dtype: x_dtype,
457        });
458        elem.push(sq);
459    }
460    let consts_set: std::collections::HashSet<NodeId> = layout.consts.iter().copied().collect();
461    let mut subst: HashMap<NodeId, NodeId> = HashMap::new();
462    // Clone EVERY body_new_carry FIRST so subst records each body_new_carry ->
463    // new_carry, then clone body_y and pred_exit sharing subst (spec "Predicate
464    // referencing body_new_carry" — no double-clone).
465    //
466    // GAP-303: the whole block is built from THIS step's `carries` before any of
467    // it is returned, matching `unroll_scan`'s ordering.
468    let new_carries: Vec<NodeId> = layout
469        .body_new_carries
470        .iter()
471        .map(|&nc| clone_body_node(graph, nc, carries, &elem, &consts_set, &mut subst))
472        .collect();
473    let y = clone_body_node(
474        graph,
475        layout.body_y,
476        carries,
477        &elem,
478        &consts_set,
479        &mut subst,
480    );
481    let stop = layout
482        .pred_exit
483        .map(|p| clone_body_node(graph, p, carries, &elem, &consts_set, &mut subst));
484    Ok(ScanStep {
485        new_carries,
486        y,
487        stop,
488    })
489}
490
491#[cfg(test)]
492mod tests {
493    use crate::opt::lower_to_base_map;
494    use crate::scan::unroll_scan;
495    use crate::{Graph, Node, Op, ScanEmit, ScanPredicate, ScanRole};
496    use fuel_ir::{DType, Shape};
497    use std::sync::{Arc, RwLock};
498
499    /// Tests that build data `Const` tensors need a real device for the
500    /// slot-populating constructors. Singleton CpuBackendDevice via OnceLock
501    /// (mirrors grad.rs:216).
502    fn cpu_dev() -> &'static std::sync::Arc<dyn fuel_backend_contract::DynBackendDevice> {
503        static D: std::sync::OnceLock<std::sync::Arc<dyn fuel_backend_contract::DynBackendDevice>> =
504            std::sync::OnceLock::new();
505        D.get_or_init(|| std::sync::Arc::new(fuel_cpu_backend::dyn_impl::CpuBackendDevice))
506    }
507
508    // Build a trivial scan: carry [1], body new_carry = carry*2, body_y =
509    // new_carry, n_xs = 0, bound = 3, emit = All. Returns (graph_arc, scan_id).
510    fn trivial_scan(
511        bound: usize,
512        emit: ScanEmit,
513        early_exit: Option<ScanPredicate>,
514    ) -> (Arc<RwLock<Graph>>, crate::NodeId) {
515        let graph = Arc::new(RwLock::new(Graph::new()));
516        let scan = {
517            let mut g = graph.write().unwrap();
518            let s = Shape::from_dims(&[1]);
519            let carry = g.push(Node {
520                op: Op::Const,
521                inputs: vec![],
522                shape: s.clone(),
523                dtype: DType::F32,
524            });
525            let hole = g.push(Node {
526                op: Op::ScanPlaceholder {
527                    role: ScanRole::Carry,
528                    index: 0,
529                },
530                inputs: vec![],
531                shape: s.clone(),
532                dtype: DType::F32,
533            });
534            let nc = g.push(Node {
535                op: Op::MulScalar(2.0),
536                inputs: vec![hole],
537                shape: s.clone(),
538                dtype: DType::F32,
539            });
540            g.push(Node {
541                op: Op::Scan {
542                    n_carries: 1,
543                    n_xs: 0,
544                    bound,
545                    emit,
546                    early_exit,
547                },
548                inputs: vec![carry, nc, nc],
549                shape: Shape::from_dims(&[bound, 1]),
550                dtype: DType::F32,
551            })
552        };
553        (graph, scan)
554    }
555
556    /// GAP-303: a TWO-carry scan, exercised through the real internal path.
557    ///
558    /// ⚠️ `NodeHandle::scan` pins `n_carries: 1`, so nothing public can build one
559    /// of these — without this test the multi-carry machinery ships written but
560    /// never RUN, and a green suite says nothing about it.
561    ///
562    /// THE DISCRIMINATING ASSERTION is that carry 1 reads its OWN init, not carry
563    /// 0's. Before this change `clone_body_node`'s `Carry` arm ignored `index`
564    /// entirely (`role: Carry, .. => carry`), so EVERY carry hole substituted to
565    /// the single carry — a 2-carry body would have silently threaded carry 0
566    /// into both slots and produced a graph that runs and is wrong.
567    fn two_carry_scan() -> (
568        Arc<RwLock<Graph>>,
569        crate::NodeId,
570        crate::NodeId,
571        crate::NodeId,
572    ) {
573        let graph = Arc::new(RwLock::new(Graph::new()));
574        let (scan, c0, c1) = {
575            let mut g = graph.write().unwrap();
576            let sh = Shape::from_dims(&[1]);
577            let mk = |g: &mut Graph, op: Op, inputs: Vec<crate::NodeId>| {
578                g.push(Node {
579                    op,
580                    inputs,
581                    shape: sh.clone(),
582                    dtype: DType::F32,
583                })
584            };
585            // Two DISTINCT inits, so "carry 1 read carry 0" is observable.
586            let c0 = mk(&mut g, Op::Const, vec![]);
587            let c1 = mk(&mut g, Op::Const, vec![]);
588            let h0 = mk(
589                &mut g,
590                Op::ScanPlaceholder {
591                    role: ScanRole::Carry,
592                    index: 0,
593                },
594                vec![],
595            );
596            let h1 = mk(
597                &mut g,
598                Op::ScanPlaceholder {
599                    role: ScanRole::Carry,
600                    index: 1,
601                },
602                vec![],
603            );
604            // Different factors, so the two carry chains are distinguishable.
605            let nc0 = mk(&mut g, Op::MulScalar(2.0), vec![h0]);
606            let nc1 = mk(&mut g, Op::MulScalar(3.0), vec![h1]);
607            let y = mk(&mut g, Op::MulScalar(5.0), vec![h0]);
608            // layout = [init_c0, init_c1, | nc0, nc1, y]
609            let scan = g.push(Node {
610                op: Op::Scan {
611                    n_carries: 2,
612                    n_xs: 0,
613                    bound: 1,
614                    emit: ScanEmit::All,
615                    early_exit: None,
616                },
617                inputs: vec![c0, c1, nc0, nc1, y],
618                shape: Shape::from_dims(&[1, 1]),
619                dtype: DType::F32,
620            });
621            (scan, c0, c1)
622        };
623        (graph, scan, c0, c1)
624    }
625
626    /// GAP-303 / Baracuda's requirement: the DUMP must name the carry count.
627    ///
628    /// ⚠️ Their reason, earned the expensive way: they spent a day on a benchmark
629    /// that ran four times, reproducibly, with four valid controls, AGAINST THE
630    /// WRONG KERNEL. "A control tells you the encoding is sound; only the
631    /// subject's identity tells you what it encoded." `short_name()` renders
632    /// every scan as the bare string "Scan", so a body built from the wrong scan
633    /// was indistinguishable in every panic message this graph can emit.
634    #[test]
635    fn gap303_describe_node_names_the_carry_count() {
636        let (graph, scan, _c0, _c1) = two_carry_scan();
637        let g = graph.read().unwrap();
638        let d = g.describe_node(scan);
639        assert!(
640            d.contains("n_carries=2"),
641            "the dump must name the carry count, got: {d}"
642        );
643        assert!(
644            d.contains("n_xs=0") && d.contains("bound=1"),
645            "and its siblings: {d}"
646        );
647        // A placeholder must name WHICH hole it is -- Carry/0 vs Carry/1 is the
648        // distinction a wrong-body bug turns on.
649        let layout = crate::scan::parse_scan_layout(&g, scan).expect("layout");
650        let hole = g.node(layout.body_new_carries[1]).inputs[0];
651        let hd = g.describe_node(hole);
652        assert!(
653            hd.contains("index=1") && hd.contains("Carry"),
654            "a carry hole must name its index, got: {hd}"
655        );
656    }
657
658    #[test]
659    fn gap303_parse_scan_layout_splits_a_two_carry_block() {
660        let (graph, scan, c0, c1) = two_carry_scan();
661        let g = graph.read().unwrap();
662        let l = crate::scan::parse_scan_layout(&g, scan).expect("layout");
663        assert_eq!(l.n_carries, 2);
664        assert_eq!(
665            l.init_carries,
666            vec![c0, c1],
667            "leading block is n_carries wide"
668        );
669        assert_eq!(
670            l.body_new_carries.len(),
671            2,
672            "trailing block is n_carries wide"
673        );
674        assert!(l.xs.is_empty() && l.consts.is_empty());
675        // body_y must be the node AFTER the new-carry block, not inside it.
676        assert!(
677            !l.body_new_carries.contains(&l.body_y),
678            "body_y must not be mistaken for a new-carry slot"
679        );
680    }
681
682    #[test]
683    fn gap303_two_carries_thread_independently_through_unroll() {
684        let (graph, scan, c0, c1) = two_carry_scan();
685        let (_ys, carries) = {
686            let mut g = graph.write().unwrap();
687            unroll_scan(&mut g, scan, 1).expect("unroll a 2-carry scan")
688        };
689        assert_eq!(carries.len(), 2, "one final carry per carry slot");
690        let g = graph.read().unwrap();
691        // ⚠️ THE ASSERTION THAT CATCHES THE OLD BEHAVIOUR: carry 1's chain must
692        // reach ITS OWN init. With the index-ignoring arm it reached c0.
693        assert_eq!(
694            g.node(carries[1]).inputs,
695            vec![c1],
696            "carry 1 must read init_carry 1, NOT init_carry 0"
697        );
698        assert_eq!(
699            g.node(carries[0]).inputs,
700            vec![c0],
701            "carry 0 reads its own init"
702        );
703        assert!(
704            matches!(g.node(carries[1]).op, Op::MulScalar(f) if f == 3.0),
705            "carry 1 keeps its own body op"
706        );
707        assert_ne!(carries[0], carries[1], "the two carries are distinct nodes");
708    }
709
710    #[test]
711    fn gap303_build_scan_step_materialises_every_carry() {
712        let (graph, scan, c0, c1) = two_carry_scan();
713        let layout = {
714            let g = graph.read().unwrap();
715            crate::scan::parse_scan_layout(&g, scan).expect("layout")
716        };
717        let step = {
718            let mut g = graph.write().unwrap();
719            crate::scan::build_scan_step(&mut g, &layout, 0, &[c0, c1]).expect("step")
720        };
721        assert_eq!(step.new_carries.len(), 2, "a step advances every carry");
722        let g = graph.read().unwrap();
723        assert_eq!(
724            g.node(step.new_carries[1]).inputs,
725            vec![c1],
726            "carry 1 stays its own"
727        );
728    }
729
730    #[test]
731    fn gap303_validator_rejects_a_carry_index_past_n_carries() {
732        let (graph, scan, _c0, _c1) = two_carry_scan();
733        let g = graph.read().unwrap();
734        let l = crate::scan::parse_scan_layout(&g, scan).expect("layout");
735        let mut roots = l.body_new_carries.clone();
736        roots.push(l.body_y);
737        // In range for n_carries = 2 ...
738        crate::scan::validate_scan_body_placeholders(&g, &roots, 2, 0)
739            .expect("Carry/0 and Carry/1 are in range when n_carries = 2");
740        // ... and OUT of range once the count shrinks. This is the guard that
741        // used to read `index != 0`; it now range-checks like the Elem arm.
742        let err = crate::scan::validate_scan_body_placeholders(&g, &roots, 1, 0)
743            .expect_err("Carry/1 must be rejected when n_carries = 1");
744        let msg = err.to_string();
745        assert!(
746            msg.contains("Carry, 1") && msg.contains("n_carries = 1"),
747            "the error must name the offending index AND the bound, got: {msg}"
748        );
749    }
750
751    #[test]
752    fn unroll_scan_all_produces_a_concat_of_steps_and_no_scan_nodes() {
753        let (graph, scan) = trivial_scan(3, ScanEmit::All, None);
754        let (ys_v, _carries) = {
755            let mut g = graph.write().unwrap();
756            unroll_scan(&mut g, scan, 3).expect("unroll")
757        };
758        // GAP-303: the ys side is a one-element Vec now (see unroll_scan).
759        let ys = ys_v[0];
760        let g = graph.read().unwrap();
761        // ys root is a Concat over the 3 steps.
762        assert!(
763            matches!(g.node(ys).op, Op::Concat { .. }),
764            "emit=All ys root should be Concat, got {:?}",
765            g.node(ys).op.short_name()
766        );
767        assert_eq!(g.node(ys).inputs.len(), 3, "one input per step");
768        // No Op::Scan / Op::ScanPlaceholder reachable from the unrolled root.
769        let reachable = crate::topo_order_multi(&g, &[ys]);
770        assert!(
771            !reachable
772                .iter()
773                .any(|&n| matches!(g.node(n).op, Op::Scan { .. } | Op::ScanPlaceholder { .. })),
774            "unrolled graph must contain no Scan/ScanPlaceholder nodes"
775        );
776    }
777
778    #[test]
779    fn unroll_scan_early_exit_some_peels_predicate_and_unrolls() {
780        // early_exit = Some layout: [carry, consts=[thr], body_new_carry, body_y, pred_exit].
781        // unroll must PEEL pred_exit, IGNORE it, and emit a 3-step Concat with no scan nodes.
782        let graph = Arc::new(RwLock::new(Graph::new()));
783        let scan = {
784            let mut g = graph.write().unwrap();
785            let s = Shape::from_dims(&[1]);
786            let carry = g.push(Node {
787                op: Op::Const,
788                inputs: vec![],
789                shape: s.clone(),
790                dtype: DType::F32,
791            });
792            let thr = g.push(Node {
793                op: Op::Const,
794                inputs: vec![],
795                shape: s.clone(),
796                dtype: DType::F32,
797            });
798            let hole = g.push(Node {
799                op: Op::ScanPlaceholder {
800                    role: ScanRole::Carry,
801                    index: 0,
802                },
803                inputs: vec![],
804                shape: s.clone(),
805                dtype: DType::F32,
806            });
807            let nc = g.push(Node {
808                op: Op::MulScalar(2.0),
809                inputs: vec![hole],
810                shape: s.clone(),
811                dtype: DType::F32,
812            });
813            // predicate sub-DAG over the post-step carry (ignored by unroll).
814            let pred = g.push(Node {
815                op: Op::Ge,
816                inputs: vec![nc, thr],
817                shape: s.clone(),
818                dtype: DType::Bool,
819            });
820            g.push(Node {
821                op: Op::Scan {
822                    n_carries: 1,
823                    n_xs: 0,
824                    bound: 3,
825                    emit: ScanEmit::All,
826                    early_exit: Some(ScanPredicate),
827                },
828                inputs: vec![carry, thr, nc, nc, pred], // consts=[thr], new_carry=nc, y=nc, pred_exit=pred
829                shape: Shape::from_dims(&[3, 1]),
830                dtype: DType::F32,
831            })
832        };
833        let (ys_v, _carries) = {
834            let mut g = graph.write().unwrap();
835            unroll_scan(&mut g, scan, 3).expect("unroll must peel + ignore the predicate")
836        };
837        // GAP-303: the ys side is a one-element Vec now (see unroll_scan).
838        let ys = ys_v[0];
839        let g = graph.read().unwrap();
840        assert!(
841            matches!(g.node(ys).op, Op::Concat { .. }),
842            "emit=All ys root should be Concat"
843        );
844        assert_eq!(g.node(ys).inputs.len(), 3, "one input per step");
845        let reachable = crate::topo_order_multi(&g, &[ys]);
846        assert!(
847            !reachable
848                .iter()
849                .any(|&n| matches!(g.node(n).op, Op::Scan { .. } | Op::ScanPlaceholder { .. })),
850            "unrolled graph must contain no Scan/ScanPlaceholder nodes"
851        );
852    }
853
854    #[test]
855    fn op_scan_is_a_terminal_in_the_base_map() {
856        // lower_to_base_map must LEAVE Op::Scan in place (no LoweringRule
857        // matches a bare Op variant) — not silently expanded, not errored.
858        let (graph, scan) = trivial_scan(3, ScanEmit::All, None);
859        let roots = lower_to_base_map(&graph, &[scan]);
860        let g = graph.read().unwrap();
861        let reachable = crate::topo_order_multi(&g, &roots);
862        assert!(
863            reachable
864                .iter()
865                .any(|&n| matches!(g.node(n).op, Op::Scan { .. })),
866            "Op::Scan must remain a terminal after lower_to_base_map"
867        );
868    }
869
870    #[test]
871    fn unroll_scan_rejects_malformed_short_inputs() {
872        // n_xs = 0 well-formed minimum is init_carry(1) + body_exits(2) = 3.
873        // Build inputs of length 2 (one short) — must be a typed Err, not a
874        // panic from the `consts = inputs[1+n_xs..inputs.len()-2]` slice
875        // (start=1 > end=0 when inputs.len() == 2).
876        let mut g = Graph::new();
877        let s = Shape::from_dims(&[1]);
878        let carry = g.push(Node {
879            op: Op::Const,
880            inputs: vec![],
881            shape: s.clone(),
882            dtype: DType::F32,
883        });
884        let body_exit = g.push(Node {
885            op: Op::Const,
886            inputs: vec![],
887            shape: s.clone(),
888            dtype: DType::F32,
889        });
890        let scan = g.push(Node {
891            op: Op::Scan {
892                n_carries: 1,
893                n_xs: 0,
894                bound: 1,
895                emit: ScanEmit::All,
896                early_exit: None,
897            },
898            inputs: vec![carry, body_exit],
899            shape: Shape::from_dims(&[1, 1]),
900            dtype: DType::F32,
901        });
902        let r = unroll_scan(&mut g, scan, 1);
903        assert!(
904            r.is_err(),
905            "inputs.len() == n_xs + 2 must be rejected as malformed, not panic"
906        );
907    }
908
909    #[test]
910    fn unroll_scan_rejects_elem_index_out_of_range() {
911        // n_xs = 0 (no xs slots) but the body references ScanPlaceholder{Elem,
912        // 0} — index 0 is out of range since n_xs = 0. Must be a typed Err,
913        // not an `elem[index]` panic inside clone_body_node.
914        let graph = Arc::new(RwLock::new(Graph::new()));
915        let scan = {
916            let mut g = graph.write().unwrap();
917            let s = Shape::from_dims(&[1]);
918            let carry = g.push(Node {
919                op: Op::Const,
920                inputs: vec![],
921                shape: s.clone(),
922                dtype: DType::F32,
923            });
924            let elem_hole = g.push(Node {
925                op: Op::ScanPlaceholder {
926                    role: ScanRole::Elem,
927                    index: 0,
928                },
929                inputs: vec![],
930                shape: s.clone(),
931                dtype: DType::F32,
932            });
933            let nc = g.push(Node {
934                op: Op::MulScalar(2.0),
935                inputs: vec![elem_hole],
936                shape: s.clone(),
937                dtype: DType::F32,
938            });
939            g.push(Node {
940                op: Op::Scan {
941                    n_carries: 1,
942                    n_xs: 0,
943                    bound: 1,
944                    emit: ScanEmit::All,
945                    early_exit: None,
946                },
947                inputs: vec![carry, nc, nc],
948                shape: Shape::from_dims(&[1, 1]),
949                dtype: DType::F32,
950            })
951        };
952        let mut g = graph.write().unwrap();
953        let r = unroll_scan(&mut g, scan, 1);
954        assert!(
955            r.is_err(),
956            "Elem index >= n_xs must be a typed Err, never an elem[index] panic"
957        );
958    }
959
960    #[test]
961    fn unroll_scan_nxs_positive_slices_substitutes_and_shares_consts() {
962        // n_xs = 1, one shared const, bound = steps = 2, emit = All. Body:
963        // new_carry = carry + elem0; y = (carry + elem0) * const — references
964        // BOTH placeholders AND the shared const. xs[0] shape [2, 1] (leading
965        // dim = bound). Locks the slice/substitute/const-sharing semantics
966        // Tasks 6-7 depend on.
967        let graph = Arc::new(RwLock::new(Graph::new()));
968        let (scan, const_id) = {
969            let mut g = graph.write().unwrap();
970            let carry_shape = Shape::from_dims(&[1]);
971            let xs_shape = Shape::from_dims(&[2, 1]);
972            let init_carry = g.push(Node {
973                op: Op::Const,
974                inputs: vec![],
975                shape: carry_shape.clone(),
976                dtype: DType::F32,
977            });
978            let xs0 = g.push(Node {
979                op: Op::Const,
980                inputs: vec![],
981                shape: xs_shape.clone(),
982                dtype: DType::F32,
983            });
984            let const_id = g.push(Node {
985                op: Op::Const,
986                inputs: vec![],
987                shape: carry_shape.clone(),
988                dtype: DType::F32,
989            });
990            let carry_hole = g.push(Node {
991                op: Op::ScanPlaceholder {
992                    role: ScanRole::Carry,
993                    index: 0,
994                },
995                inputs: vec![],
996                shape: carry_shape.clone(),
997                dtype: DType::F32,
998            });
999            let elem_hole = g.push(Node {
1000                op: Op::ScanPlaceholder {
1001                    role: ScanRole::Elem,
1002                    index: 0,
1003                },
1004                inputs: vec![],
1005                shape: carry_shape.clone(),
1006                dtype: DType::F32,
1007            });
1008            let sum = g.push(Node {
1009                op: Op::Add,
1010                inputs: vec![carry_hole, elem_hole],
1011                shape: carry_shape.clone(),
1012                dtype: DType::F32,
1013            });
1014            let new_carry = sum;
1015            let y = g.push(Node {
1016                op: Op::Mul,
1017                inputs: vec![sum, const_id],
1018                shape: carry_shape.clone(),
1019                dtype: DType::F32,
1020            });
1021            let scan = g.push(Node {
1022                op: Op::Scan {
1023                    n_carries: 1,
1024                    n_xs: 1,
1025                    bound: 2,
1026                    emit: ScanEmit::All,
1027                    early_exit: None,
1028                },
1029                inputs: vec![init_carry, xs0, const_id, new_carry, y],
1030                shape: Shape::from_dims(&[2, 1]),
1031                dtype: DType::F32,
1032            });
1033            (scan, const_id)
1034        };
1035        let (ys_v, _carries) = {
1036            let mut g = graph.write().unwrap();
1037            unroll_scan(&mut g, scan, 2).expect("unroll")
1038        };
1039        // GAP-303: the ys side is a one-element Vec now (see unroll_scan).
1040        let ys = ys_v[0];
1041        let g = graph.read().unwrap();
1042        assert!(
1043            matches!(g.node(ys).op, Op::Concat { .. }),
1044            "emit=All ys root should be Concat, got {:?}",
1045            g.node(ys).op.short_name()
1046        );
1047        assert_eq!(g.node(ys).inputs.len(), 2, "one input per step");
1048        let reachable = crate::topo_order_multi(&g, &[ys]);
1049        assert!(
1050            !reachable
1051                .iter()
1052                .any(|&n| matches!(g.node(n).op, Op::Scan { .. } | Op::ScanPlaceholder { .. })),
1053            "unrolled graph must contain no Scan/ScanPlaceholder nodes"
1054        );
1055        // The const NodeId must be SHARED across both step clones — it
1056        // appears exactly once in the reachable set (topo_order_multi
1057        // dedups by NodeId), never re-cloned per step.
1058        let const_occurrences = reachable.iter().filter(|&&n| n == const_id).count();
1059        assert_eq!(
1060            const_occurrences, 1,
1061            "const node must be shared (not cloned) across steps"
1062        );
1063    }
1064
1065    #[test]
1066    fn scan_until_builds_early_exit_node_hashes_distinctly_and_validates() {
1067        use crate::NodeHandle;
1068        // init_carry [1]; body new_carry = carry*2; consts include threshold.
1069        let init = NodeHandle::from_f32(vec![1.0f32], Shape::from_dims(&[1]), cpu_dev()).unwrap();
1070        let graph = init.graph().clone();
1071        // Build the shared body + predicate at graph level, wrap as NodeHandle handles.
1072        let (nc, thr, pred_ok) = {
1073            let mut g = graph.write().unwrap();
1074            let s = Shape::from_dims(&[1]);
1075            let hole = g.push(Node {
1076                op: Op::ScanPlaceholder {
1077                    role: ScanRole::Carry,
1078                    index: 0,
1079                },
1080                inputs: vec![],
1081                shape: s.clone(),
1082                dtype: DType::F32,
1083            });
1084            let nc = g.push(Node {
1085                op: Op::MulScalar(2.0),
1086                inputs: vec![hole],
1087                shape: s.clone(),
1088                dtype: DType::F32,
1089            });
1090            let thr = g.push(Node {
1091                op: Op::Const,
1092                inputs: vec![],
1093                shape: s.clone(),
1094                dtype: DType::F32,
1095            });
1096            let pred = g.push(Node {
1097                op: Op::Ge,
1098                inputs: vec![nc, thr],
1099                shape: s.clone(),
1100                dtype: DType::Bool,
1101            });
1102            (nc, thr, pred)
1103        };
1104        let nc_t = NodeHandle::from_existing(graph.clone(), nc);
1105        let thr_t = NodeHandle::from_existing(graph.clone(), thr);
1106        let pred_t = NodeHandle::from_existing(graph.clone(), pred_ok);
1107
1108        let out = init
1109            .scan_until(
1110                &[],
1111                std::slice::from_ref(&thr_t),
1112                &nc_t,
1113                &nc_t,
1114                &pred_t,
1115                5,
1116                ScanEmit::Final,
1117            )
1118            .expect("well-formed scan_until must build");
1119        // The producer node behind the emit=Final view is an Op::Scan{early_exit: Some}.
1120        let scan_id = {
1121            let g = graph.read().unwrap();
1122            g.node(out.id()).inputs[0]
1123        };
1124        {
1125            let g = graph.read().unwrap();
1126            match &g.node(scan_id).op {
1127                Op::Scan { early_exit, .. } => {
1128                    assert!(early_exit.is_some(), "early_exit must be Some")
1129                }
1130                other => panic!("expected Op::Scan, got {}", other.short_name()),
1131            }
1132            // pred_exit is the LAST input (trailing), so reachability sees it.
1133            assert_eq!(*g.node(scan_id).inputs.last().unwrap(), pred_ok);
1134        }
1135
1136        // base_map_hash distinctness: a scan with the SAME body but a DIFFERENT predicate hashes differently.
1137        let thr2 = {
1138            let mut g = graph.write().unwrap();
1139            g.push(Node {
1140                op: Op::Const,
1141                inputs: vec![],
1142                shape: Shape::from_dims(&[1]),
1143                dtype: DType::F32,
1144            })
1145        };
1146        let pred2 = {
1147            let mut g = graph.write().unwrap();
1148            g.push(Node {
1149                op: Op::Le,
1150                inputs: vec![nc, thr2],
1151                shape: Shape::from_dims(&[1]),
1152                dtype: DType::Bool,
1153            })
1154        };
1155        let pred2_t = NodeHandle::from_existing(graph.clone(), pred2);
1156        let out2 = init
1157            .scan_until(
1158                &[],
1159                &[NodeHandle::from_existing(graph.clone(), thr2)],
1160                &nc_t,
1161                &nc_t,
1162                &pred2_t,
1163                5,
1164                ScanEmit::Final,
1165            )
1166            .expect("second scan_until builds");
1167        let scan2 = {
1168            let g = graph.read().unwrap();
1169            g.node(out2.id()).inputs[0]
1170        };
1171        let (h1, h2) = {
1172            let g = graph.read().unwrap();
1173            (
1174                crate::opt::base_map_hash(&g, scan_id),
1175                crate::opt::base_map_hash(&g, scan2),
1176            )
1177        };
1178        assert_ne!(
1179            h1, h2,
1180            "different predicates must hash distinctly (predicate is a trailing input)"
1181        );
1182
1183        // Rejection: a NON-scalar predicate is a typed Err (never a panic).
1184        let big =
1185            NodeHandle::from_f32(vec![0.0f32, 1.0], Shape::from_dims(&[2]), cpu_dev()).unwrap(); // wrong graph AND non-scalar
1186        assert!(
1187            init.scan_until(
1188                &[],
1189                std::slice::from_ref(&thr_t),
1190                &nc_t,
1191                &nc_t,
1192                &big,
1193                5,
1194                ScanEmit::Final
1195            )
1196            .is_err(),
1197            "non-same-graph / non-scalar predicate must be a typed Err"
1198        );
1199        // Rejection: a non-Bool predicate.
1200        let f32pred = {
1201            let mut g = graph.write().unwrap();
1202            g.push(Node {
1203                op: Op::Sqr,
1204                inputs: vec![nc],
1205                shape: Shape::from_dims(&[1]),
1206                dtype: DType::F32,
1207            })
1208        };
1209        let f32pred_t = NodeHandle::from_existing(graph.clone(), f32pred);
1210        assert!(
1211            init.scan_until(&[], &[thr_t], &nc_t, &nc_t, &f32pred_t, 5, ScanEmit::Final)
1212                .is_err(),
1213            "non-Bool predicate must be a typed Err"
1214        );
1215    }
1216
1217    #[test]
1218    fn build_scan_step_shares_subst_so_predicate_reads_this_steps_new_carry() {
1219        // Scan: carry [1]; new_carry = carry*2; pred = Ge(new_carry, thr). n_xs=0, bound=4.
1220        let graph = Arc::new(RwLock::new(Graph::new()));
1221        let scan = {
1222            let mut g = graph.write().unwrap();
1223            let s = Shape::from_dims(&[1]);
1224            let carry = g.push(Node {
1225                op: Op::Const,
1226                inputs: vec![],
1227                shape: s.clone(),
1228                dtype: DType::F32,
1229            });
1230            let thr = g.push(Node {
1231                op: Op::Const,
1232                inputs: vec![],
1233                shape: s.clone(),
1234                dtype: DType::F32,
1235            });
1236            let hole = g.push(Node {
1237                op: Op::ScanPlaceholder {
1238                    role: ScanRole::Carry,
1239                    index: 0,
1240                },
1241                inputs: vec![],
1242                shape: s.clone(),
1243                dtype: DType::F32,
1244            });
1245            let nc = g.push(Node {
1246                op: Op::MulScalar(2.0),
1247                inputs: vec![hole],
1248                shape: s.clone(),
1249                dtype: DType::F32,
1250            });
1251            let pred = g.push(Node {
1252                op: Op::Ge,
1253                inputs: vec![nc, thr],
1254                shape: s.clone(),
1255                dtype: DType::Bool,
1256            });
1257            g.push(Node {
1258                op: Op::Scan {
1259                    n_carries: 1,
1260                    n_xs: 0,
1261                    bound: 4,
1262                    emit: ScanEmit::Final,
1263                    early_exit: Some(ScanPredicate),
1264                },
1265                inputs: vec![carry, thr, nc, nc, pred],
1266                shape: Shape::from_dims(&[4, 1]),
1267                dtype: DType::F32,
1268            })
1269        };
1270        let layout = {
1271            let g = graph.read().unwrap();
1272            crate::scan::parse_scan_layout(&g, scan).expect("layout")
1273        };
1274        assert!(layout.pred_exit.is_some());
1275        let init = layout.init_carries.clone();
1276        let step = {
1277            let mut g = graph.write().unwrap();
1278            crate::scan::build_scan_step(&mut g, &layout, 0, &init).expect("step")
1279        };
1280        let stop = step.stop.expect("early-exit scan yields a stop node");
1281        // The predicate clone must reach step.new_carry (the shared post-step carry),
1282        // and must reach NO ScanPlaceholder (all substituted away).
1283        let g = graph.read().unwrap();
1284        let reach = crate::topo_order_multi(&g, &[stop]);
1285        assert!(
1286            reach.contains(&step.new_carries[0]),
1287            "pred must read THIS step's new_carry (shared subst)"
1288        );
1289        assert!(
1290            !reach
1291                .iter()
1292                .any(|&n| matches!(g.node(n).op, Op::ScanPlaceholder { .. })),
1293            "no placeholders survive a materialized step"
1294        );
1295        // The Ge's first input IS the step's new_carry — proof there was no double-clone.
1296        let ge_inputs = &g.node(stop).inputs;
1297        assert_eq!(
1298            ge_inputs[0], step.new_carries[0],
1299            "predicate's post-step operand is the shared new_carry"
1300        );
1301    }
1302
1303    #[test]
1304    fn scan_until_rejects_out_of_range_elem_placeholder_at_build_time() {
1305        use crate::NodeHandle;
1306        // n_xs = 1, but the body references ScanPlaceholder{Elem, 5} — out of
1307        // range. Must be a typed BUILD-TIME Err (before this fix the node built
1308        // fine and the forward driver's build_scan_step -> clone_body_node
1309        // panicked with elem[5] index-OOB — a never-panic regression).
1310        let init = NodeHandle::from_f32(vec![0.0f32], Shape::from_dims(&[1]), cpu_dev()).unwrap();
1311        let graph = init.graph().clone();
1312        let (x, nc, thr, pred) = {
1313            let mut g = graph.write().unwrap();
1314            let s = Shape::from_dims(&[1]);
1315            let x = g.push(Node {
1316                op: Op::Const,
1317                inputs: vec![],
1318                shape: Shape::from_dims(&[3, 1]),
1319                dtype: DType::F32,
1320            });
1321            let bad_elem = g.push(Node {
1322                op: Op::ScanPlaceholder {
1323                    role: ScanRole::Elem,
1324                    index: 5,
1325                },
1326                inputs: vec![],
1327                shape: s.clone(),
1328                dtype: DType::F32,
1329            });
1330            let carry_hole = g.push(Node {
1331                op: Op::ScanPlaceholder {
1332                    role: ScanRole::Carry,
1333                    index: 0,
1334                },
1335                inputs: vec![],
1336                shape: s.clone(),
1337                dtype: DType::F32,
1338            });
1339            let nc = g.push(Node {
1340                op: Op::Add,
1341                inputs: vec![carry_hole, bad_elem],
1342                shape: s.clone(),
1343                dtype: DType::F32,
1344            });
1345            let thr = g.push(Node {
1346                op: Op::Const,
1347                inputs: vec![],
1348                shape: s.clone(),
1349                dtype: DType::F32,
1350            });
1351            let pred = g.push(Node {
1352                op: Op::Ge,
1353                inputs: vec![carry_hole, thr],
1354                shape: s.clone(),
1355                dtype: DType::Bool,
1356            });
1357            (x, nc, thr, pred)
1358        };
1359        let x_t = NodeHandle::from_existing(graph.clone(), x);
1360        let thr_t = NodeHandle::from_existing(graph.clone(), thr);
1361        let nc_t = NodeHandle::from_existing(graph.clone(), nc);
1362        let pred_t = NodeHandle::from_existing(graph.clone(), pred);
1363        let r = init.scan_until(&[x_t], &[thr_t], &nc_t, &nc_t, &pred_t, 3, ScanEmit::Final);
1364        assert!(
1365            r.is_err(),
1366            "scan_until must reject body Elem{{index >= n_xs}} at build time, not panic in the driver"
1367        );
1368    }
1369
1370    #[test]
1371    fn scan_rejects_out_of_range_elem_placeholder_at_build_time() {
1372        use crate::NodeHandle;
1373        // Base builder shares the same missing build-time check (Phase 1 only
1374        // avoided the panic because unroll_scan was its sole forward path).
1375        let init = NodeHandle::from_f32(vec![0.0f32], Shape::from_dims(&[1]), cpu_dev()).unwrap();
1376        let graph = init.graph().clone();
1377        let (x, nc) = {
1378            let mut g = graph.write().unwrap();
1379            let s = Shape::from_dims(&[1]);
1380            let x = g.push(Node {
1381                op: Op::Const,
1382                inputs: vec![],
1383                shape: Shape::from_dims(&[3, 1]),
1384                dtype: DType::F32,
1385            });
1386            let bad_elem = g.push(Node {
1387                op: Op::ScanPlaceholder {
1388                    role: ScanRole::Elem,
1389                    index: 5,
1390                },
1391                inputs: vec![],
1392                shape: s.clone(),
1393                dtype: DType::F32,
1394            });
1395            let carry_hole = g.push(Node {
1396                op: Op::ScanPlaceholder {
1397                    role: ScanRole::Carry,
1398                    index: 0,
1399                },
1400                inputs: vec![],
1401                shape: s.clone(),
1402                dtype: DType::F32,
1403            });
1404            let nc = g.push(Node {
1405                op: Op::Add,
1406                inputs: vec![carry_hole, bad_elem],
1407                shape: s.clone(),
1408                dtype: DType::F32,
1409            });
1410            (x, nc)
1411        };
1412        let x_t = NodeHandle::from_existing(graph.clone(), x);
1413        let nc_t = NodeHandle::from_existing(graph.clone(), nc);
1414        let r = init.scan(&[x_t], &[], &nc_t, &nc_t, 3, ScanEmit::Final);
1415        assert!(
1416            r.is_err(),
1417            "base scan must also reject body Elem{{index >= n_xs}} at build time"
1418        );
1419    }
1420
1421    #[test]
1422    fn backward_lowers_op_scan_and_no_longer_panics() {
1423        use crate::NodeHandle;
1424        // Affine scan: carry[1]; consts a,b; new_carry = a*carry + b; emit=Final. bound=3.
1425        let init = NodeHandle::from_f32(vec![1.0f32], Shape::from_dims(&[1]), cpu_dev()).unwrap();
1426        let a = NodeHandle::from_existing(init.graph().clone(), init.id())
1427            .const_f32_like(vec![0.5f32], Shape::from_dims(&[1]))
1428            .unwrap();
1429        let b = NodeHandle::from_existing(init.graph().clone(), init.id())
1430            .const_f32_like(vec![0.1f32], Shape::from_dims(&[1]))
1431            .unwrap();
1432        let graph = init.graph().clone();
1433        let nc = {
1434            let mut g = graph.write().unwrap();
1435            let s = Shape::from_dims(&[1]);
1436            let hole = g.push(Node {
1437                op: Op::ScanPlaceholder {
1438                    role: ScanRole::Carry,
1439                    index: 0,
1440                },
1441                inputs: vec![],
1442                shape: s.clone(),
1443                dtype: DType::F32,
1444            });
1445            let ac = g.push(Node {
1446                op: Op::Mul,
1447                inputs: vec![a.id(), hole],
1448                shape: s.clone(),
1449                dtype: DType::F32,
1450            });
1451            g.push(Node {
1452                op: Op::Add,
1453                inputs: vec![ac, b.id()],
1454                shape: s.clone(),
1455                dtype: DType::F32,
1456            })
1457        };
1458        let nc_t = NodeHandle::from_existing(graph.clone(), nc);
1459        let out = init
1460            .scan(
1461                &[],
1462                &[a.clone(), b.clone()],
1463                &nc_t,
1464                &nc_t,
1465                3,
1466                ScanEmit::Final,
1467            )
1468            .expect("scan");
1469        // backward() must NOT panic (Phase 1 arm panics here) and must yield a grad for init_carry.
1470        let grads = out.backward();
1471        let g_init = grads.get(&init).expect("gradient for init_carry");
1472        // The gradient's subgraph must contain no Op::Scan/ScanPlaceholder (proof of lowering).
1473        let g = graph.read().unwrap();
1474        let reach = crate::topo_order_multi(&g, &[g_init.id()]);
1475        assert!(
1476            !reach
1477                .iter()
1478                .any(|&n| matches!(g.node(n).op, Op::Scan { .. } | Op::ScanPlaceholder { .. })),
1479            "backward must lower the scan before differentiating"
1480        );
1481        assert!(
1482            grads.get(&a).is_some() && grads.get(&b).is_some(),
1483            "consts a,b get gradients (BPTT)"
1484        );
1485    }
1486}
fuel_graph/scan.rs

fuel_graph/
scan.rs
