formawasm 0.0.1-beta

//! Lowering of `Optional<T>` value coercions.
//!
//! Phase 2 mc2 already lowers the `nil` literal end-to-end. This
//! module supplies the matching `Some` direction: when an expression
//! of static type `T` flows into a slot typed `Optional<T>`, allocate
//! a full-size `Optional<T>` cell, store the `OPTIONAL_TAG_SOME`
//! discriminant at offset 0 and the payload at `payload_offset`, and
//! leave the cell's pointer on the stack. Aggregate payloads (struct,
//! enum, tuple, array) ride later mcs — for now `wrap_some` only
//! handles primitive payload types via [`super::aggregate::store_primitive`].

use formalang::ir::{IrExpr, ResolvedType};
use wasm_encoder::{InstructionSink, MemArg, ValType};

use super::aggregate::{allocate_aggregate, primitive_of, store_primitive};
use super::block::{ScratchCounts, bump_count};
use super::{LowerContext, LowerError, lower_expr};
use crate::layout::{FieldLayout, OPTIONAL_TAG_ALIGN, OPTIONAL_TAG_SOME, plan_optional};
use crate::module::MEMORY_INDEX;

/// Whether `target_ty` is an `Optional<U>` and `value_ty` is exactly
/// `U` (the Some-wrap shape). Returns `Some(U)` for the wrap case so
/// callers can plan layouts and scratch slots ahead of emission;
/// returns `None` for every other type combination — including the
/// already-Optional `value_ty == target_ty` case and the
/// `Optional<Never>` "nil widens to Optional<T>" shortcut, both of
/// which travel through the existing pointer pass-through path.
#[must_use]
pub(super) fn some_wrap_payload<'a>(
    target_ty: &'a ResolvedType,
    value_ty: &ResolvedType,
    module: &formalang::ir::IrModule,
) -> Option<&'a ResolvedType> {
    let inner = crate::compound::optional_inner(target_ty, module)?;
    if value_ty == inner { Some(inner) } else { None }
}

/// Lower a Some-wrap of `value_expr`'s value into a fresh
/// `Optional<payload_ty>` cell.
///
/// The value is evaluated first and stashed in a typed scratch local
/// so the bump-allocator call doesn't disturb the operand stack. The
/// emitted sequence is:
///
/// ```text
/// <value_expr>          ; stack: [value]
/// local.set value_scratch
/// i32.const size
/// call __alloc          ; stack: [ptr]
/// local.set ptr_scratch
/// local.get ptr_scratch
/// i32.const OPTIONAL_TAG_SOME
/// i32.store offset=0
/// local.get ptr_scratch
/// local.get value_scratch
/// <store_primitive>     ; payload at payload_offset
/// local.get ptr_scratch ; stack: [ptr]
/// ```
pub(super) fn lower_some_wrap(
    value_expr: &IrExpr,
    payload_ty: &ResolvedType,
    sink: &mut InstructionSink<'_>,
    ctx: &LowerContext<'_>,
) -> Result<(), LowerError> {
    let prim = primitive_of(payload_ty).map_err(|_| LowerError::NotYetImplemented {
        what: format!("Some-wrap of non-primitive payload type {payload_ty:?}"),
    })?;
    let module = ctx.module()?;
    let layout = plan_optional(payload_ty, module)?;

    let value_vt = primitive_to_valtype(payload_ty)?;
    let value_scratch = ctx.next_scratch_local(value_vt)?;

    lower_expr(value_expr, sink, ctx)?;
    sink.local_set(value_scratch);

    let base_local = allocate_aggregate(layout.size, sink, ctx)?;

    sink.local_get(base_local);
    sink.i32_const(i32::try_from(OPTIONAL_TAG_SOME).unwrap_or(i32::MAX));
    sink.i32_store(MemArg {
        offset: u64::from(layout.tag_offset),
        align: OPTIONAL_TAG_ALIGN.trailing_zeros(),
        memory_index: MEMORY_INDEX,
    });

    sink.local_get(base_local);
    sink.local_get(value_scratch);
    let payload_field = FieldLayout {
        offset: layout.payload_offset,
        size: layout.payload_size,
        align: layout.payload_align,
    };
    store_primitive(prim, payload_field, sink);

    sink.local_get(base_local);
    Ok(())
}

/// Map a primitive payload type to the wasm value type its scratch
/// slot needs. Aggregates won't reach this helper — `lower_some_wrap`
/// rejects non-primitive payloads with `NotYetImplemented` ahead of
/// any scratch-local allocation.
fn primitive_to_valtype(ty: &ResolvedType) -> Result<ValType, LowerError> {
    let prim = primitive_of(ty).map_err(|_| LowerError::NotYetImplemented {
        what: format!("Some-wrap scratch slot for non-primitive payload {ty:?}"),
    })?;
    match prim {
        formalang::ast::PrimitiveType::I32 | formalang::ast::PrimitiveType::Boolean => {
            Ok(ValType::I32)
        }
        formalang::ast::PrimitiveType::I64 => Ok(ValType::I64),
        formalang::ast::PrimitiveType::F32 => Ok(ValType::F32),
        formalang::ast::PrimitiveType::F64 => Ok(ValType::F64),
        // String / Path / Regex are stored as `i32` header pointers
        // — `primitive_size_align` lays them out at `(4, 4)` and
        // `lower_string_literal` materialises them onto the stack
        // as i32 pointer values. The Some-wrap path stashes that
        // pointer in an `i32` scratch and stores it at the
        // optional payload offset.
        formalang::ast::PrimitiveType::String
        | formalang::ast::PrimitiveType::Path
        | formalang::ast::PrimitiveType::Regex => Ok(ValType::I32),
        // Never has no instance.
        formalang::ast::PrimitiveType::Never | _ => Err(LowerError::NotYetImplemented {
            what: format!("Some-wrap scratch slot for {prim:?} payload"),
        }),
    }
}

/// Wasm value-type used for the typed scratch slot a Some-wrap
/// reserves to stash its payload across the bump-allocator call. The
/// pre-walk in `block::walk_count` consults this so the per-type
/// scratch counts match the lowering walker's consumption.
pub(super) fn some_wrap_scratch_valtype(payload_ty: &ResolvedType) -> Result<ValType, LowerError> {
    primitive_to_valtype(payload_ty)
}

/// Lower `value_expr` into a slot of static type `target_ty`,
/// inserting an `Optional<T>` Some-wrap when the slot widens the
/// value's type. Sites that own a known target type (function
/// returns, if branches, match arm bodies) call this in place of the
/// bare [`lower_expr`] so a primitive `T` flowing into an
/// `Optional<T>` slot materializes as a tagged-Some cell rather than
/// a bare `T`. Other type combinations (exact match,
/// Optional<Never> -> Optional<T>) flow through unchanged via the
/// regular lowering path.
pub(super) fn lower_coerced(
    value_expr: &IrExpr,
    target_ty: &ResolvedType,
    sink: &mut InstructionSink<'_>,
    ctx: &LowerContext<'_>,
) -> Result<(), LowerError> {
    // Some-wrap recognition needs the prelude's Optional id, which
    // lives on the IrModule. When the lowering context doesn't carry
    // a module (the legacy primitive-only `lower_function_body`
    // entry point), an Optional Some-wrap can't apply by definition
    // — there's no Optional<T> declaration to recognise. Fall
    // through to the plain lowering rather than tripping
    // `MissingContext`.
    if let Some(module) = ctx.module_opt()
        && let Some(payload_ty) = some_wrap_payload(target_ty, value_expr.ty(), module)
    {
        return lower_some_wrap(value_expr, payload_ty, sink, ctx);
    }
    // Trait-erasure widening: `let s: Shape = Square(...)`. The
    // value's static type is a concrete `Struct(_)` / `Enum(_)`;
    // the binding's declared type is `Trait(_)`. Materialise an
    // 8-byte fat-pointer cell `(vtable_offset, data_ptr)` so the
    // surrounding code can dispatch through the trait without
    // knowing the concrete impl. Vtable offsets come from
    // `ctx.vtable_offset(trait_id, impl_target_key(target))` —
    // populated by `module_lowering::build_vtable_plumbing`.
    //
    // For control-flow values (If, Match, Block) whose branches each
    // produce a *different* concrete type, the static vtable offset
    // would only be right for one branch. Push the coercion into
    // each branch instead so every leaf materialises with its own
    // matching vtable.
    if let ResolvedType::Trait(trait_id) = target_ty {
        return lower_trait_coercion(*trait_id, value_expr, sink, ctx);
    }
    lower_expr(value_expr, sink, ctx)
}

/// Resolve a concrete `(Struct/Enum)` `ImplTarget` for a value's
/// type so the trait-erasure path can pick the matching vtable.
/// Returns `None` for primitive / closure / trait-typed values
/// (the latter is already a fat pointer; no re-materialisation
/// needed).
const fn trait_dispatch_target(ty: &ResolvedType) -> Option<formalang::ir::ImplTarget> {
    use formalang::ir::ImplTarget;
    if let Some(sid) = crate::compound::struct_id_of(ty) {
        return Some(ImplTarget::Struct(sid));
    }
    if let Some(eid) = crate::compound::enum_id_of(ty) {
        return Some(ImplTarget::Enum(eid));
    }
    None
}

/// Coerce `value_expr` into a trait fat pointer. For control-flow
/// shapes whose branches each produce a different concrete type
/// (If / Match / tail-expression Block), recurse so the
/// materialisation happens at every leaf with its branch's matching
/// vtable. Otherwise emit a single materialisation against the
/// value's static type.
fn lower_trait_coercion(
    trait_id: formalang::ir::TraitId,
    value_expr: &IrExpr,
    sink: &mut InstructionSink<'_>,
    ctx: &LowerContext<'_>,
) -> Result<(), LowerError> {
    // Already a trait fat pointer — passthrough.
    if matches!(value_expr.ty(), ResolvedType::Trait(_)) {
        return lower_expr(value_expr, sink, ctx);
    }
    match value_expr {
        IrExpr::If {
            condition,
            then_branch,
            else_branch,
            ..
        } => {
            // Mirror `lower_if`: condition + framed if/else, but with
            // each branch coerced into the trait shape (i32 cell ptr).
            lower_expr(condition, sink, ctx)?;
            sink.if_(wasm_encoder::BlockType::Result(ValType::I32));
            lower_trait_coercion(trait_id, then_branch, sink, ctx)?;
            if let Some(else_b) = else_branch {
                sink.else_();
                lower_trait_coercion(trait_id, else_b, sink, ctx)?;
            }
            sink.end();
            Ok(())
        }
        IrExpr::Match { .. } => {
            // Match arms each produce a concrete branch value too.
            // Falling back to a single materialisation here would
            // miscompile cross-impl dispatch — surface explicitly so
            // a future pass adds the per-arm wiring.
            Err(LowerError::NotYetImplemented {
                what: "trait coercion through `match` arms (push lower_coerced into each arm)"
                    .to_owned(),
            })
        }
        IrExpr::Block {
            statements, result, ..
        } => {
            // Lower preceding statements normally; coerce the result expression.
            for stmt in statements {
                super::block::lower_block_statement(stmt, sink, ctx)?;
            }
            lower_trait_coercion(trait_id, result, sink, ctx)
        }
        IrExpr::Literal { .. }
        | IrExpr::StructInst { .. }
        | IrExpr::EnumInst { .. }
        | IrExpr::Array { .. }
        | IrExpr::Tuple { .. }
        | IrExpr::Reference { .. }
        | IrExpr::SelfFieldRef { .. }
        | IrExpr::FieldAccess { .. }
        | IrExpr::LetRef { .. }
        | IrExpr::BinaryOp { .. }
        | IrExpr::UnaryOp { .. }
        | IrExpr::For { .. }
        | IrExpr::FunctionCall { .. }
        | IrExpr::CallClosure { .. }
        | IrExpr::MethodCall { .. }
        | IrExpr::Closure { .. }
        | IrExpr::ClosureRef { .. }
        | IrExpr::DictLiteral { .. }
        | IrExpr::DictAccess { .. } => {
            let target = trait_dispatch_target(value_expr.ty()).ok_or_else(|| {
                LowerError::NotYetImplemented {
                    what: format!(
                        "trait coercion of value with non-aggregate type {:?}",
                        value_expr.ty()
                    ),
                }
            })?;
            materialize_trait_fat_pointer(trait_id, target, value_expr, sink, ctx)
        }
    }
}

/// Allocate an 8-byte fat-pointer cell `(vtable_offset:i32,
/// data_ptr:i32)` and leave the cell pointer on the wasm stack.
/// `vtable_offset` is resolved from the lowering context; the
/// data pointer is the concrete value the inner `value_expr`
/// produces.
fn materialize_trait_fat_pointer(
    trait_id: formalang::ir::TraitId,
    target: formalang::ir::ImplTarget,
    value_expr: &IrExpr,
    sink: &mut InstructionSink<'_>,
    ctx: &LowerContext<'_>,
) -> Result<(), LowerError> {
    use crate::layout::POINTER_SIZE;
    use crate::module::MEMORY_INDEX;
    use wasm_encoder::MemArg;
    let vtable_offset =
        ctx.vtable_offset(trait_id, crate::module_lowering::impl_target_key(target))?;
    // Stash the data pointer in a scratch local so we can store
    // it after evaluating the allocator call.
    let data_scratch = ctx.next_scratch_local(ValType::I32)?;
    super::lower_expr(value_expr, sink, ctx)?;
    sink.local_set(data_scratch);

    let cell_size = i32::try_from(POINTER_SIZE.saturating_mul(2)).unwrap_or(8);
    let bump_idx = ctx.bump_allocator()?;
    let cell_scratch = ctx.next_scratch_local(ValType::I32)?;
    sink.i32_const(cell_size)
        .call(bump_idx)
        .local_set(cell_scratch);

    let mem_arg = |off: u64| MemArg {
        offset: off,
        align: 2,
        memory_index: MEMORY_INDEX,
    };
    let vtable_off_signed = i32::try_from(vtable_offset).unwrap_or(i32::MAX);
    sink.local_get(cell_scratch);
    sink.i32_const(vtable_off_signed);
    sink.i32_store(mem_arg(0));
    sink.local_get(cell_scratch);
    sink.local_get(data_scratch);
    sink.i32_store(mem_arg(u64::from(POINTER_SIZE)));
    sink.local_get(cell_scratch);
    Ok(())
}

/// Add the scratch-slot reservations a Some-wrap coercion at this
/// site requires, on top of whatever the inner expression already
/// counts. Callers that own a known target type pair this with the
/// regular `walk_count` recursion so the pre-walk's totals match the
/// lowering walker's consumption.
pub(super) fn coercion_scratch_counts(
    target_ty: &ResolvedType,
    value_expr: &IrExpr,
    out: &mut ScratchCounts,
    module: Option<&formalang::ir::IrModule>,
) -> Result<(), LowerError> {
    let value_ty = value_expr.ty();
    // Without module access we can't classify Optional<T> from
    // Generic — over-reserving is harmless to correctness, but
    // skipping when None mirrors the previous behaviour where
    // unrelated types didn't reserve scratch slots either.
    let Some(module) = module else {
        return Ok(());
    };
    // Trait-erasure widening: target=Trait(_), value coerces into a
    // fat-pointer cell. The lowering path pushes the coercion into
    // each branch of an If/Match/Block, so we count two i32 scratch
    // locals per *leaf* materialisation, not per outer call.
    if matches!(target_ty, ResolvedType::Trait(_)) {
        return count_trait_coercion_leaves(value_expr, out);
    }
    let Some(payload_ty) = some_wrap_payload(target_ty, value_ty, module) else {
        return Ok(());
    };
    bump_count(&mut out.i32)?;
    let vt = some_wrap_scratch_valtype(payload_ty)?;
    match vt {
        ValType::I32 => bump_count(&mut out.i32)?,
        ValType::I64 => bump_count(&mut out.i64)?,
        ValType::F32 => bump_count(&mut out.f32)?,
        ValType::F64 => bump_count(&mut out.f64)?,
        ValType::V128 | ValType::Ref(_) => {
            return Err(LowerError::NotYetImplemented {
                what: format!("Some-wrap scratch slot of value type {vt:?}"),
            });
        }
    }
    Ok(())
}

/// Walk `value_expr` and reserve 2 i32 scratch locals per
/// materialisation leaf, mirroring the recursion in
/// [`lower_trait_coercion`]. An If contributes leaves through both
/// arms; a Block contributes through its result expression. Other
/// shapes count as a single leaf.
fn count_trait_coercion_leaves(
    value_expr: &IrExpr,
    out: &mut ScratchCounts,
) -> Result<(), LowerError> {
    if matches!(value_expr.ty(), ResolvedType::Trait(_)) {
        // Already a trait fat pointer — passthrough, no extra scratch.
        return Ok(());
    }
    match value_expr {
        IrExpr::If {
            then_branch,
            else_branch,
            ..
        } => {
            count_trait_coercion_leaves(then_branch, out)?;
            if let Some(else_b) = else_branch {
                count_trait_coercion_leaves(else_b, out)?;
            }
            Ok(())
        }
        IrExpr::Block { result, .. } => count_trait_coercion_leaves(result, out),
        IrExpr::Literal { .. }
        | IrExpr::StructInst { .. }
        | IrExpr::EnumInst { .. }
        | IrExpr::Array { .. }
        | IrExpr::Tuple { .. }
        | IrExpr::Reference { .. }
        | IrExpr::SelfFieldRef { .. }
        | IrExpr::FieldAccess { .. }
        | IrExpr::LetRef { .. }
        | IrExpr::BinaryOp { .. }
        | IrExpr::UnaryOp { .. }
        | IrExpr::For { .. }
        | IrExpr::Match { .. }
        | IrExpr::FunctionCall { .. }
        | IrExpr::CallClosure { .. }
        | IrExpr::MethodCall { .. }
        | IrExpr::Closure { .. }
        | IrExpr::ClosureRef { .. }
        | IrExpr::DictLiteral { .. }
        | IrExpr::DictAccess { .. } => {
            bump_count(&mut out.i32)?;
            bump_count(&mut out.i32)?;
            Ok(())
        }
    }
}