Skip to main content

vyre_libs/parsing/rust/lower/
mod.rs

1//! Rust AST to Vyre IR lowering (reusable substrate, Tier 3).
2//!
3//! Mirrors `vyre-libs::parsing::c::lower`: lowering a resolved language AST to a
4//! `vyre::ir::Program` is a Tier-3 concern and lives in the library, not the
5//! frontend driver. Lowering borrows the AST plus its resolution.
6//!
7//! ## Model (nano-subset)
8//!
9//! The last function in the module is the entry kernel. In scalar mode each
10//! `i32`/`bool` parameter becomes a one-element `ReadOnly` input buffer and the
11//! return value is stored into `out[0]`. In batched mode each parameter becomes
12//! a same-length input buffer indexed by `global_id.x`, and the return value is
13//! stored into `out[global_id.x]` behind an out-of-range guard. Local bindings
14//! are alpha-renamed to `v{binding_id}` so Rust's legal shadowing never trips
15//! Vyre's no-shadowing validator (V008). `if`/`else` whose arms both return is
16//! a terminal statement. Anything outside the wired subset returns a loud
17//! [`RustLowerError::Unsupported`](crate::parsing::rust::lower::RustLowerError::Unsupported) rather than miscompiling.
18
19use std::collections::HashMap;
20
21use thiserror::Error;
22use vyre::ir::{BufferDecl, DataType, Expr as IrExpr, Node, Program};
23
24use super::lex::tokens::{ANDAND, EQ, GE, GT, LE, LT, MINUS, NE, OROR, PERCENT, PLUS, SLASH, STAR};
25use super::parse::{Expr, Module, Stmt, Type};
26use super::sema::{BindingId, Resolution};
27
28/// Errors from Rust to Vyre IR lowering.
29#[derive(Debug, Clone, Error)]
30pub enum RustLowerError {
31    /// The module declares no function to use as the entry kernel.
32    #[error("Rust lowering needs at least one function to use as the entry kernel")]
33    NoEntryFunction,
34    /// A construct outside the wired lowering subset was encountered.
35    #[error(
36        "Rust to Vyre IR lowering does not support {0} yet; not emitting a miscompiled Program"
37    )]
38    Unsupported(String),
39}
40
41/// Lower a resolved module to a Vyre IR program (the last function is the entry).
42///
43/// # Errors
44/// Returns [`RustLowerError::NoEntryFunction`] for an empty module, or
45/// [`RustLowerError::Unsupported`] for any construct outside the wired subset.
46pub fn lower(module: &Module, resolution: &Resolution) -> Result<Program, RustLowerError> {
47    lower_entry(module, resolution, LowerMode::Scalar)
48}
49
50/// Lower the module entry as a data-parallel map over `lane_count` elements.
51///
52/// The Rust function keeps scalar source semantics per lane: parameter `pN`
53/// reads from `pN[global_id.x]`, and `return e` writes to `out[global_id.x]`.
54/// Extra invocations from workgroup rounding are guarded before any buffer
55/// access. This preserves the scalar [`lower`] entrypoint while exposing the
56/// source frontend as a GPU-native batch kernel.
57///
58/// # Errors
59/// Returns [`RustLowerError::Unsupported`] when `lane_count == 0`, or any error
60/// returned by scalar lowering for unsupported language constructs.
61pub fn lower_batched(
62    module: &Module,
63    resolution: &Resolution,
64    lane_count: u32,
65) -> Result<Program, RustLowerError> {
66    if lane_count == 0 {
67        return Err(RustLowerError::Unsupported(
68            "batched Rust lowering with zero lanes".to_string(),
69        ));
70    }
71    lower_entry(module, resolution, LowerMode::Batched { lane_count })
72}
73
74#[derive(Clone, Copy)]
75enum LowerMode {
76    Scalar,
77    Batched { lane_count: u32 },
78}
79
80impl LowerMode {
81    fn buffer_count(self) -> u32 {
82        match self {
83            Self::Scalar => 1,
84            Self::Batched { lane_count } => lane_count,
85        }
86    }
87
88    fn workgroup_size(self) -> [u32; 3] {
89        match self {
90            Self::Scalar => [1, 1, 1],
91            Self::Batched { .. } => [256, 1, 1],
92        }
93    }
94
95    fn lane_index(self) -> IrExpr {
96        match self {
97            Self::Scalar => IrExpr::u32(0),
98            Self::Batched { .. } => IrExpr::var(BATCH_LANE_VAR),
99        }
100    }
101}
102
103const BATCH_LANE_VAR: &str = "__rust_lane";
104
105fn lower_entry(
106    module: &Module,
107    resolution: &Resolution,
108    mode: LowerMode,
109) -> Result<Program, RustLowerError> {
110    let entry_index = module
111        .functions
112        .len()
113        .checked_sub(1)
114        .ok_or(RustLowerError::NoEntryFunction)?;
115    let func = &module.functions[entry_index];
116
117    let def_to_id: HashMap<u32, BindingId> = resolution
118        .bindings
119        .iter()
120        .enumerate()
121        .map(|(id, b)| (b.def_offset, id))
122        .collect();
123
124    let mut buffers = Vec::with_capacity(func.params.len() + 1);
125    let mut entry_nodes = Vec::new();
126    for (i, (offset, ty)) in func.params.iter().enumerate() {
127        let dtype = scalar_dtype(ty)?;
128        let buf = format!("p{i}");
129        buffers.push(BufferDecl::read(&buf, i as u32, dtype).with_count(mode.buffer_count()));
130        let binding = def_to_id
131            .get(offset)
132            .copied()
133            .ok_or_else(|| RustLowerError::Unsupported("unresolved parameter".to_string()))?;
134        entry_nodes.push(Node::let_bind(
135            format!("v{binding}"),
136            IrExpr::load(buf, mode.lane_index()),
137        ));
138    }
139    let out_dtype = scalar_dtype(&func.ret)?;
140    buffers.push(
141        BufferDecl::output("out", func.params.len() as u32, out_dtype)
142            .with_count(mode.buffer_count()),
143    );
144
145    let ctx = LowerCtx {
146        module,
147        resolution,
148        def_to_id: &def_to_id,
149        output_index: mode.lane_index(),
150    };
151    entry_nodes.extend(ctx.lower_stmts(&func.body, Subst::Local(None))?);
152    let entry_nodes = match mode {
153        LowerMode::Scalar => entry_nodes,
154        LowerMode::Batched { lane_count } => vec![
155            Node::let_bind(BATCH_LANE_VAR, IrExpr::gid_x()),
156            Node::if_then(
157                IrExpr::lt(IrExpr::var(BATCH_LANE_VAR), IrExpr::u32(lane_count)),
158                entry_nodes,
159            ),
160        ],
161    };
162    Ok(Program::wrapped(
163        buffers,
164        mode.workgroup_size(),
165        entry_nodes,
166    ))
167}
168
169/// The Vyre element type for a nano scalar type (only `i32`/`bool` so far).
170fn scalar_dtype(ty: &Type) -> Result<DataType, RustLowerError> {
171    match ty {
172        Type::I32 => Ok(DataType::I32),
173        Type::Bool => Ok(DataType::Bool),
174        Type::Unit => Err(RustLowerError::Unsupported(
175            "unit-typed parameter or return".to_string(),
176        )),
177        // A reference parameter carries its pointee value: in the
178        // assignment-free nano-subset a `&T` is a pure read-alias, so it lowers
179        // to a buffer of the pointee's element type.
180        Type::Ref { inner, .. } => scalar_dtype(inner),
181    }
182}
183
184struct LowerCtx<'a> {
185    module: &'a Module,
186    resolution: &'a Resolution,
187    def_to_id: &'a HashMap<u32, BindingId>,
188    output_index: IrExpr,
189}
190
191/// How a variable use lowers in the current scope.
192///
193/// The Rust lowering walks two structurally different scopes, and a variable
194/// that is *not* named in the scope's map means opposite things in each. A
195/// single `Option<&HashMap>` cannot express both, so the scope is explicit:
196///
197/// - [`Subst::Local`] is the entry-function scope and every loop body: a
198///   variable lowers to its local `v{binding}` unless the overlay renames it
199///   (e.g. a loop induction variable rebound to its loop variable). A binding
200///   absent from the overlay is a normal local, not an error.
201/// - [`Subst::Inline`] is the scope while a straight-line callee is being
202///   inlined: every callee binding must be present in the map (mapped to its
203///   caller-scope argument or folded `let`). A binding absent from the map is a
204///   lowering bug (an un-substituted callee variable), not a local.
205#[derive(Clone, Copy)]
206enum Subst<'a> {
207    /// Entry/loop-body scope with an optional rename overlay.
208    Local(Option<&'a HashMap<BindingId, IrExpr>>),
209    /// Callee-inline scope: every variable must be substituted.
210    Inline(&'a HashMap<BindingId, IrExpr>),
211}
212
213impl LowerCtx<'_> {
214    fn lower_stmts(&self, stmts: &[Stmt], subst: Subst<'_>) -> Result<Vec<Node>, RustLowerError> {
215        let mut nodes = Vec::new();
216        for stmt in stmts {
217            match stmt {
218                Stmt::Let { name, init, .. } => {
219                    let binding = self.def_to_id.get(name).copied().ok_or_else(|| {
220                        RustLowerError::Unsupported("unresolved let binding".to_string())
221                    })?;
222                    nodes.push(Node::let_bind(
223                        format!("v{binding}"),
224                        self.lower_value(init, subst)?,
225                    ));
226                }
227                Stmt::Return(Some(expr)) => {
228                    nodes.push(Node::store(
229                        "out",
230                        self.output_index.clone(),
231                        self.lower_value(expr, subst)?,
232                    ));
233                    return Ok(nodes);
234                }
235                Stmt::Return(None) => return Ok(nodes),
236                Stmt::Assign { name, value } => {
237                    let binding = self.resolution.uses.get(name).copied().ok_or_else(|| {
238                        RustLowerError::Unsupported("unresolved assignment target".to_string())
239                    })?;
240                    nodes.push(Node::assign(
241                        format!("v{binding}"),
242                        self.lower_value(value, subst)?,
243                    ));
244                }
245                Stmt::Expr(Expr::If {
246                    cond,
247                    then_block,
248                    else_block,
249                }) => {
250                    let then_nodes = self.lower_stmts(block_stmts(then_block), subst)?;
251                    let else_nodes = match else_block {
252                        Some(block) => self.lower_stmts(block_stmts(block), subst)?,
253                        None => Vec::new(),
254                    };
255                    nodes.push(Node::if_then_else(
256                        self.lower_value(cond, subst)?,
257                        then_nodes,
258                        else_nodes,
259                    ));
260                    let then_div = stmts_diverge(block_stmts(then_block));
261                    let else_div = else_block
262                        .as_ref()
263                        .is_some_and(|b| stmts_diverge(block_stmts(b)));
264                    if then_div && else_div {
265                        return Ok(nodes);
266                    }
267                }
268                Stmt::While { cond, body } => {
269                    nodes.extend(self.lower_while(cond, body, subst)?);
270                }
271                Stmt::For {
272                    name,
273                    start,
274                    end,
275                    body,
276                } => {
277                    nodes.extend(self.lower_for_range(*name, start, end, body, subst)?);
278                }
279                // Pure expression statements have no observable effect; drop them.
280                Stmt::Expr(_) => {}
281            }
282        }
283        Ok(nodes)
284    }
285
286    /// Half-open `[0, trip)` u32 trip count for a counted loop over the signed
287    /// range `[lo, hi)`: `hi - lo` when `hi > lo`, else 0. The subtraction is a
288    /// *u32* wrapping subtraction of the two's-complement bit patterns, so a
289    /// range exceeding `i32::MAX` stays exact and a negative `lo` does not wrap
290    /// to ~4.29e9 iterations; the signed `hi > lo` guard (comparisons never
291    /// overflow) zeroes the empty/negative-length case, matching Rust. This is
292    /// the single source of the counted-loop bound invariant, both `while` and
293    /// `for` lowering call it, so the correctness fix lives in exactly one place.
294    fn counted_loop_trip(lo: &IrExpr, hi: &IrExpr) -> IrExpr {
295        let span_u32 = IrExpr::sub(
296            IrExpr::cast(DataType::U32, hi.clone()),
297            IrExpr::cast(DataType::U32, lo.clone()),
298        );
299        IrExpr::select(IrExpr::gt(hi.clone(), lo.clone()), span_u32, IrExpr::u32(0))
300    }
301
302    /// Reconstruct the signed induction variable inside a counted-loop body:
303    /// `lo + lv`, where `lv` is the fresh 0-based u32 loop counter read back as
304    /// `i32` so body arithmetic stays well-typed (and correct for negative `lo`).
305    fn counted_loop_induction(lo: &IrExpr, loop_var: &str) -> IrExpr {
306        IrExpr::add(
307            lo.clone(),
308            IrExpr::cast(DataType::I32, IrExpr::var(loop_var.to_string())),
309        )
310    }
311
312    /// Lower `while i < BOUND { ...; i = i + 1; }` to a counted `Node::Loop`.
313    /// Only this exact counting form is supported; anything else (data-dependent
314    /// exit, mutated bound, `i` assigned outside the trailing increment) returns
315    /// a loud `Unsupported` rather than a miscompiled loop.
316    fn lower_while(
317        &self,
318        cond: &Expr,
319        body: &[Stmt],
320        subst: Subst<'_>,
321    ) -> Result<Vec<Node>, RustLowerError> {
322        let bad = || {
323            RustLowerError::Unsupported(
324                "while loop that is not a canonical `while i < BOUND { ...; i = i + 1; }` counting loop"
325                    .to_string(),
326            )
327        };
328        // cond must be `i < BOUND`.
329        let (i_off, bound) = match cond {
330            Expr::Binary { op, lhs, rhs } if *op == LT => match lhs.as_ref() {
331                Expr::Var(off) => (*off, rhs.as_ref()),
332                _ => return Err(bad()),
333            },
334            _ => return Err(bad()),
335        };
336        let b_i = self.resolution.uses.get(&i_off).copied().ok_or_else(bad)?;
337        // Trailing statement must be `i = i + 1`.
338        let Some((last, init_stmts)) = body.split_last() else {
339            return Err(bad());
340        };
341        let inc_ok = matches!(last, Stmt::Assign { name, value }
342            if self.resolution.uses.get(name).copied() == Some(b_i)
343            && matches!(value, Expr::Binary { op, lhs, rhs }
344                if *op == PLUS
345                && matches!(lhs.as_ref(), Expr::Var(o) if self.resolution.uses.get(o).copied() == Some(b_i))
346                && matches!(rhs.as_ref(), Expr::LiteralInt(_, 1))));
347        if !inc_ok {
348            return Err(bad());
349        }
350        // `i` must not be assigned anywhere except the trailing increment.
351        if stmts_assign_binding(init_stmts, b_i, self.resolution) {
352            return Err(bad());
353        }
354        // BOUND must be loop-invariant: none of its free variables are assigned
355        // in the body (so evaluating it once for the loop bound matches Rust).
356        for v in expr_var_bindings(bound, self.resolution) {
357            if stmts_assign_binding(body, v, self.resolution) {
358                return Err(bad());
359            }
360        }
361        let loop_var = format!("v{b_i}__w");
362        // Extend the current scope's overlay with the induction-variable rename,
363        // preserving the scope kind (a `while` only appears in `Local` scope
364        // today, since inlined callees are straight-line; the `Inline` arm keeps
365        // the invariant sound if that ever changes).
366        let mut inner: HashMap<BindingId, IrExpr> = match subst {
367            Subst::Local(Some(m)) => m.clone(),
368            Subst::Local(None) => HashMap::new(),
369            Subst::Inline(m) => m.clone(),
370        };
371        // The IR loop variable is `u32` by contract (`Node::Loop` bounds and
372        // index are u32), but the Rust induction variable is `i32`. Inside the
373        // body every use of `i` must read back as `i32`, so the overlay maps the
374        // binding to `cast(i32, loop_var)`; otherwise `acc + i` would mix u32 and
375        // i32 operands and fail IR validation.
376        // The induction variable `i` is reconstructed as `i0 + lv`, where `lv`
377        // is a fresh 0-based u32 loop counter and `i0` is the (possibly negative)
378        // signed initial value held in `v{b_i}`. Reading `i` back as i32 keeps
379        // body arithmetic well-typed and stays correct even when `i0 < 0`: a
380        // plain `cast(u32, i0)` loop bound would wrap a negative start to ~4.29e9
381        // and iterate billions of times instead of matching Rust.
382        inner.insert(
383            b_i,
384            Self::counted_loop_induction(&IrExpr::var(format!("v{b_i}")), &loop_var),
385        );
386        let inner_subst = match subst {
387            Subst::Inline(_) => Subst::Inline(&inner),
388            Subst::Local(_) => Subst::Local(Some(&inner)),
389        };
390        // The IR loop is a half-open `[0, trip)` u32 range. Computing the trip
391        // count in *signed* space and clamping to zero before the u32 cast makes
392        // a zero- or negative-length range run the body zero times, exactly like
393        // Rust (`while i < n` with `i >= n` never enters), instead of wrapping a
394        // negative `n - i0` into billions of iterations (a DoS + miscompile).
395        let from_i32 = IrExpr::var(format!("v{b_i}"));
396        let to_i32 = self.lower_value(bound, subst)?;
397        let trip = Self::counted_loop_trip(&from_i32, &to_i32);
398        let from = IrExpr::u32(0);
399        let to = trip;
400        let loop_body = self.lower_stmts(init_stmts, inner_subst)?;
401        // After the loop `i == max(i0, n)`: the bound `n` when the loop ran at
402        // least once, or the unchanged start `i0` for a zero-trip loop. The old
403        // `i := n` was wrong for the zero-trip case (Rust leaves `i` untouched).
404        let post = IrExpr::select(
405            IrExpr::gt(to_i32.clone(), from_i32.clone()),
406            to_i32,
407            from_i32,
408        );
409        Ok(vec![
410            Node::loop_for(loop_var, from, to, loop_body),
411            Node::assign(format!("v{b_i}"), post),
412        ])
413    }
414
415    /// Lower `for i in START..END { body }` to a signed, half-open counted
416    /// `Node::Loop`. Range bounds are snapshotted before the loop so mutation in
417    /// the body cannot change how many iterations run, matching Rust's iterator
418    /// construction semantics for `Range<i32>`.
419    fn lower_for_range(
420        &self,
421        name: u32,
422        start: &Expr,
423        end: &Expr,
424        body: &[Stmt],
425        subst: Subst<'_>,
426    ) -> Result<Vec<Node>, RustLowerError> {
427        let b_i = self.def_to_id.get(&name).copied().ok_or_else(|| {
428            RustLowerError::Unsupported("unresolved for-loop binding".to_string())
429        })?;
430        let start_name = format!("v{b_i}__for_start");
431        let end_name = format!("v{b_i}__for_end");
432        let loop_var = format!("v{b_i}__for");
433
434        let start_i32 = IrExpr::var(start_name.clone());
435        let end_i32 = IrExpr::var(end_name.clone());
436        let trip = Self::counted_loop_trip(&start_i32, &end_i32);
437
438        let mut inner: HashMap<BindingId, IrExpr> = match subst {
439            Subst::Local(Some(m)) => m.clone(),
440            Subst::Local(None) => HashMap::new(),
441            Subst::Inline(m) => m.clone(),
442        };
443        inner.insert(b_i, Self::counted_loop_induction(&start_i32, &loop_var));
444        let inner_subst = match subst {
445            Subst::Inline(_) => Subst::Inline(&inner),
446            Subst::Local(_) => Subst::Local(Some(&inner)),
447        };
448        let loop_body = self.lower_stmts(body, inner_subst)?;
449
450        Ok(vec![
451            Node::let_bind(start_name, self.lower_value(start, subst)?),
452            Node::let_bind(end_name, self.lower_value(end, subst)?),
453            Node::loop_for(loop_var, IrExpr::u32(0), trip, loop_body),
454        ])
455    }
456
457    /// Lower a value expression in the current scope. In a [`Subst::Local`]
458    /// scope a variable lowers to its alpha-renamed local `v{id}` unless the
459    /// overlay renames it (a loop induction variable); in a [`Subst::Inline`]
460    /// scope every callee variable must be present in the map (mapped to its
461    /// caller-scope argument or a folded `let`), and an absent one is a bug.
462    fn lower_value(&self, expr: &Expr, subst: Subst<'_>) -> Result<IrExpr, RustLowerError> {
463        match expr {
464            Expr::LiteralInt(_, value) => Ok(IrExpr::i32(*value as i32)),
465            Expr::LiteralBool(_, value) => Ok(IrExpr::bool(*value)),
466            Expr::Var(offset) => {
467                let binding = self.resolution.uses.get(offset).copied().ok_or_else(|| {
468                    RustLowerError::Unsupported("unresolved variable use".to_string())
469                })?;
470                match subst {
471                    // Local scope: an overlay rename wins, otherwise the binding
472                    // is its own local. An absent binding is a normal local.
473                    Subst::Local(Some(map)) => Ok(map
474                        .get(&binding)
475                        .cloned()
476                        .unwrap_or_else(|| IrExpr::var(format!("v{binding}")))),
477                    Subst::Local(None) => Ok(IrExpr::var(format!("v{binding}"))),
478                    // Inline scope: every callee binding must be substituted.
479                    Subst::Inline(map) => map.get(&binding).cloned().ok_or_else(|| {
480                        RustLowerError::Unsupported("callee variable not substituted".to_string())
481                    }),
482                }
483            }
484            Expr::Binary { op, lhs, rhs } => {
485                let l = self.lower_value(lhs, subst)?;
486                let r = self.lower_value(rhs, subst)?;
487                Ok(match *op {
488                    PLUS => IrExpr::add(l, r),
489                    MINUS => IrExpr::sub(l, r),
490                    STAR => IrExpr::mul(l, r),
491                    SLASH => IrExpr::div(l, r),
492                    // Vyre types `Mod`'s result as u32 even for i32 operands;
493                    // the value is signed-correct, so cast back to i32 to keep
494                    // composition and the i32 store well-typed.
495                    PERCENT => IrExpr::cast(DataType::I32, IrExpr::rem(l, r)),
496                    EQ => IrExpr::eq(l, r),
497                    NE => IrExpr::ne(l, r),
498                    LT => IrExpr::lt(l, r),
499                    GT => IrExpr::gt(l, r),
500                    LE => IrExpr::le(l, r),
501                    GE => IrExpr::ge(l, r),
502                    ANDAND => IrExpr::and(l, r),
503                    OROR => IrExpr::or(l, r),
504                    other => {
505                        return Err(RustLowerError::Unsupported(format!(
506                            "binary operator {other}"
507                        )))
508                    }
509                })
510            }
511            Expr::Call { name, args } => self.lower_call(name, args, subst),
512            // In the assignment-free nano-subset a reference is a pure read
513            // alias: `&e` evaluates to e's value and `*r` reads that value back,
514            // so borrow and dereference are value-transparent.
515            Expr::Borrow { expr, .. } => self.lower_value(expr, subst),
516            Expr::Deref(inner) => self.lower_value(inner, subst),
517            Expr::Not(inner) => Ok(IrExpr::not(self.lower_value(inner, subst)?)),
518            // Rust unary `-x` on i32 is wrapping negation. The IR's total
519            // `Negate` is illegal on i32 (the i32::MIN overflow case), so lower
520            // to `0 - x`, which is value-identical for all in-range x and
521            // wrapping-correct at i32::MIN (matching Rust release semantics).
522            Expr::Neg(inner) => Ok(IrExpr::sub(IrExpr::i32(0), self.lower_value(inner, subst)?)),
523            Expr::Block(_) | Expr::If { .. } => Err(RustLowerError::Unsupported(
524                "block/if used as a value".to_string(),
525            )),
526        }
527    }
528
529    /// Inline a call to a straight-line single-return callee: substitute its
530    /// parameters with the (caller-scope) argument expressions, fold its `let`
531    /// bindings, and return its lowered return expression. A callee with control
532    /// flow or no terminal return is a loud `Unsupported` (never miscompiled).
533    fn lower_call(
534        &self,
535        name: &u32,
536        args: &[Expr],
537        caller_subst: Subst<'_>,
538    ) -> Result<IrExpr, RustLowerError> {
539        let callee_index = self
540            .resolution
541            .calls
542            .get(name)
543            .copied()
544            .ok_or_else(|| RustLowerError::Unsupported("unresolved call".to_string()))?;
545        let callee = &self.module.functions[callee_index];
546        if args.len() != callee.params.len() {
547            return Err(RustLowerError::Unsupported(
548                "call arity mismatch".to_string(),
549            ));
550        }
551        let mut subst: HashMap<BindingId, IrExpr> = HashMap::new();
552        for (i, (offset, _)) in callee.params.iter().enumerate() {
553            let binding = self.def_to_id.get(offset).copied().ok_or_else(|| {
554                RustLowerError::Unsupported("unresolved callee parameter".to_string())
555            })?;
556            subst.insert(binding, self.lower_value(&args[i], caller_subst)?);
557        }
558        for stmt in &callee.body {
559            match stmt {
560                Stmt::Let {
561                    name: offset, init, ..
562                } => {
563                    let value = self.lower_value(init, Subst::Inline(&subst))?;
564                    let binding = self.def_to_id.get(offset).copied().ok_or_else(|| {
565                        RustLowerError::Unsupported("unresolved callee binding".to_string())
566                    })?;
567                    subst.insert(binding, value);
568                }
569                Stmt::Return(Some(expr)) => return self.lower_value(expr, Subst::Inline(&subst)),
570                _ => {
571                    return Err(RustLowerError::Unsupported(
572                        "call to a callee with control flow or no terminal return".to_string(),
573                    ))
574                }
575            }
576        }
577        Err(RustLowerError::Unsupported(
578            "call to a callee with no return".to_string(),
579        ))
580    }
581}
582
583/// Whether any statement assigns binding `b` (recursing into if/while bodies).
584fn stmts_assign_binding(stmts: &[Stmt], b: BindingId, res: &Resolution) -> bool {
585    stmts.iter().any(|s| match s {
586        Stmt::Assign { name, .. } => res.uses.get(name).copied() == Some(b),
587        Stmt::Expr(Expr::If {
588            then_block,
589            else_block,
590            ..
591        }) => {
592            stmts_assign_binding(block_stmts(then_block), b, res)
593                || else_block
594                    .as_ref()
595                    .is_some_and(|e| stmts_assign_binding(block_stmts(e), b, res))
596        }
597        Stmt::While { body, .. } => stmts_assign_binding(body, b, res),
598        Stmt::For { body, .. } => stmts_assign_binding(body, b, res),
599        _ => false,
600    })
601}
602
603/// Collect the binding ids of every variable read in `expr`.
604fn expr_var_bindings(expr: &Expr, res: &Resolution) -> Vec<BindingId> {
605    let mut out = Vec::new();
606    collect_var_bindings(expr, res, &mut out);
607    out
608}
609
610fn collect_var_bindings(expr: &Expr, res: &Resolution, out: &mut Vec<BindingId>) {
611    match expr {
612        Expr::Var(off) => {
613            if let Some(&id) = res.uses.get(off) {
614                out.push(id);
615            }
616        }
617        Expr::Binary { lhs, rhs, .. } => {
618            collect_var_bindings(lhs, res, out);
619            collect_var_bindings(rhs, res, out);
620        }
621        Expr::Borrow { expr, .. } => collect_var_bindings(expr, res, out),
622        Expr::Deref(inner) => collect_var_bindings(inner, res, out),
623        Expr::Not(inner) => collect_var_bindings(inner, res, out),
624        Expr::Neg(inner) => collect_var_bindings(inner, res, out),
625        Expr::Call { args, .. } => {
626            for a in args {
627                collect_var_bindings(a, res, out);
628            }
629        }
630        _ => {}
631    }
632}
633
634/// Statement list of a block expression (empty for anything else).
635fn block_stmts(expr: &Expr) -> &[Stmt] {
636    match expr {
637        Expr::Block(stmts) => stmts,
638        _ => &[],
639    }
640}
641
642/// Whether a statement sequence returns on every path (so trailing code is
643/// unreachable). Mirrors the typeck divergence check.
644fn stmts_diverge(stmts: &[Stmt]) -> bool {
645    stmts.iter().any(|stmt| match stmt {
646        Stmt::Return(_) => true,
647        Stmt::Expr(Expr::If {
648            then_block,
649            else_block: Some(else_block),
650            ..
651        }) => stmts_diverge(block_stmts(then_block)) && stmts_diverge(block_stmts(else_block)),
652        _ => false,
653    })
654}