yulang-runtime 0.1.1

//! Runtime thunk construction rules used by core-to-runtime lowering.
//!
//! This module owns the places where lowering deliberately changes a value
//! expression into `Thunk[effect, value]` or forces a thunk with `BindHere`.
//! Later passes may preserve, specialize, or refine these shapes, but they
//! should not rediscover thunk boundaries from erased display types.
//!
//! The intended rules are:
//!
//! - an expected thunk type wraps a compatible value expression in `Thunk`;
//! - an expected value type forces a compatible thunk with `BindHere`;
//! - forced effects from expression forms are attached once with `Thunk`;
//! - function-boundary hygiene protects created thunks with `AddId[peek, effect]`;
//! - pure or empty effects are left as values, not administrative thunks.

use super::*;

pub(super) fn apply_param_allowed_effect(
    param_ty: RuntimeType,
    annotation: Option<&typed_ir::ParamEffectAnnotation>,
    function_allowed_effects: Option<&typed_ir::FunctionSigAllowedEffects>,
) -> RuntimeType {
    if function_allowed_effects.is_none()
        && let Some(allowed) = allowed_effect_for_param(annotation, None, &param_ty)
            .map(|allowed| project_runtime_effect(&allowed))
    {
        return match param_ty {
            RuntimeType::Thunk { .. } => param_ty,
            value => RuntimeType::thunk(allowed, value),
        };
    }
    param_ty
}

pub(super) fn allowed_effect_for_param(
    annotation: Option<&typed_ir::ParamEffectAnnotation>,
    function_allowed_effects: Option<&typed_ir::FunctionSigAllowedEffects>,
    param_ty: &RuntimeType,
) -> Option<typed_ir::Type> {
    if let Some(allowed) = function_allowed_effects {
        return Some(match allowed {
            typed_ir::FunctionSigAllowedEffects::Wildcard => wildcard_effect_type(),
            typed_ir::FunctionSigAllowedEffects::Effects(paths) => typed_ir::Type::Row {
                items: paths
                    .iter()
                    .cloned()
                    .map(|path| typed_ir::Type::Named {
                        path,
                        args: Vec::new(),
                    })
                    .collect(),
                tail: Box::new(typed_ir::Type::Never),
            },
        });
    }
    match annotation {
        Some(typed_ir::ParamEffectAnnotation::Wildcard) => Some(wildcard_effect_type()),
        Some(typed_ir::ParamEffectAnnotation::Region(_)) => thunk_effect(param_ty),
        None if returns_thunk(param_ty) => Some(empty_row()),
        None => None,
    }
}

pub(super) fn runtime_boundary_effect_for_param(
    annotation: Option<&typed_ir::ParamEffectAnnotation>,
    function_allowed_effects: Option<&typed_ir::FunctionSigAllowedEffects>,
    _param_ty: &RuntimeType,
) -> Option<typed_ir::Type> {
    if let Some(allowed) = function_allowed_effects {
        return Some(function_allowed_effect_type(allowed));
    }
    match annotation {
        Some(typed_ir::ParamEffectAnnotation::Wildcard) => Some(wildcard_effect_type()),
        Some(typed_ir::ParamEffectAnnotation::Region(name)) => Some(typed_ir::Type::Row {
            items: vec![typed_ir::Type::Named {
                path: typed_ir::Path::from_name(name.clone()),
                args: Vec::new(),
            }],
            tail: Box::new(typed_ir::Type::Never),
        }),
        None => Some(empty_row()),
    }
}

pub(super) fn returns_thunk(ty: &RuntimeType) -> bool {
    match ty {
        RuntimeType::Thunk { .. } => true,
        RuntimeType::Fun { ret, .. } => returns_thunk(ret),
        _ => false,
    }
}

pub(super) fn empty_row() -> typed_ir::Type {
    typed_ir::Type::Row {
        items: Vec::new(),
        tail: Box::new(typed_ir::Type::Never),
    }
}

fn function_allowed_effect_type(allowed: &typed_ir::FunctionSigAllowedEffects) -> typed_ir::Type {
    match allowed {
        typed_ir::FunctionSigAllowedEffects::Wildcard => wildcard_effect_type(),
        typed_ir::FunctionSigAllowedEffects::Effects(paths) => typed_ir::Type::Row {
            items: paths
                .iter()
                .cloned()
                .map(|path| typed_ir::Type::Named {
                    path,
                    args: Vec::new(),
                })
                .collect(),
            tail: Box::new(typed_ir::Type::Never),
        },
    }
}

#[cfg_attr(not(test), allow(dead_code))]
pub(super) fn prepare_expr_for_expected_profiled(
    expr: Expr,
    expected: &RuntimeType,
    source: TypeSource,
    profile: &mut RuntimeAdapterProfile,
) -> RuntimeResult<Expr> {
    prepare_expr_for_expected_with_adapter_source_profiled(expr, expected, source, profile, None)
}

pub(super) fn prepare_expr_for_expected_with_adapter_source_profiled(
    expr: Expr,
    expected: &RuntimeType,
    source: TypeSource,
    profile: &mut RuntimeAdapterProfile,
    adapter_source: Option<RuntimeAdapterSource>,
) -> RuntimeResult<Expr> {
    if hir_type_contains_unknown(expected) {
        return Ok(expr);
    }
    match expected {
        RuntimeType::Unknown | RuntimeType::Core(typed_ir::Type::Unknown) => Ok(expr),
        RuntimeType::Thunk { effect, value } => match &expr.ty {
            RuntimeType::Thunk { .. } => {
                require_apply_arg_compatible(expected, &expr.ty, source)?;
                profile.reused_thunk += 1;
                Ok(expr)
            }
            _ => {
                if matches!(expr.ty, RuntimeType::Fun { .. })
                    && !effect_is_empty(effect)
                    && !can_delay_function_value_for_expected_thunk(expected, source)
                {
                    return Err(RuntimeError::ExpectedThunk {
                        ty: runtime_core_type(&expr.ty),
                    });
                }
                require_same_hir_type(value, &expr.ty, source)?;
                let value = more_informative_hir_type(value, &expr.ty);
                let ty = RuntimeType::thunk(effect.clone(), value.clone());
                profile.value_to_thunk += 1;
                if apply_adapter_source(source) {
                    profile.apply_evidence_value_to_thunk += 1;
                    record_apply_adapter(
                        profile,
                        RuntimeAdapterEventKind::ValueToThunk,
                        source,
                        adapter_source.clone(),
                        &expr.ty,
                        expected,
                    );
                    if apply_arg_source_edge(source)
                        || adapter_source.as_ref().is_some_and(|source| {
                            source.has_apply_callee_source_edge || source.has_apply_arg_source_edge
                        })
                    {
                        profile.apply_evidence_value_to_thunk_with_source_edge += 1;
                    }
                }
                Ok(Expr::typed(
                    ExprKind::Thunk {
                        effect: effect.clone(),
                        value,
                        expr: Box::new(expr),
                    },
                    ty,
                ))
            }
        },
        // A join is a semantic value boundary: even an open result slot must
        // force effectful branch arms here. Render/discard sinks wrap outside it.
        RuntimeType::Core(typed_ir::Type::Any | typed_ir::Type::Var(_))
            if source != TypeSource::JoinEvidence
                && matches!(expr.ty, RuntimeType::Thunk { .. }) =>
        {
            Ok(expr)
        }
        _ => {
            if matches!(expr.ty, RuntimeType::Thunk { .. }) && apply_adapter_source(source) {
                profile.apply_evidence_thunk_to_value += 1;
                profile.apply_evidence_bind_here += 1;
                record_apply_adapter(
                    profile,
                    RuntimeAdapterEventKind::ThunkToValue,
                    source,
                    adapter_source.clone(),
                    &expr.ty,
                    expected,
                );
                record_apply_adapter(
                    profile,
                    RuntimeAdapterEventKind::BindHere,
                    source,
                    adapter_source.clone(),
                    &expr.ty,
                    expected,
                );
                if apply_arg_source_edge(source)
                    || adapter_source.as_ref().is_some_and(|source| {
                        source.has_apply_callee_source_edge || source.has_apply_arg_source_edge
                    })
                {
                    profile.apply_evidence_thunk_to_value_with_source_edge += 1;
                    profile.apply_evidence_bind_here_with_source_edge += 1;
                }
            }
            let (expr, actual) = force_value_expr_profiled(expr, profile);
            let expected_core = runtime_core_type(expected);
            let actual_core = runtime_core_type(&actual);
            if needs_runtime_coercion(&expected_core, &actual_core) {
                return Ok(Expr::typed(
                    ExprKind::Coerce {
                        from: actual_core,
                        to: expected_core,
                        expr: Box::new(expr),
                    },
                    expected.clone(),
                ));
            }
            require_same_hir_type(expected, &actual, source)?;
            Ok(expr)
        }
    }
}

fn can_delay_function_value_for_expected_thunk(expected: &RuntimeType, source: TypeSource) -> bool {
    apply_arg_source_edge(source)
        && matches!(
            expected,
            RuntimeType::Thunk { value, .. } if matches!(value.as_ref(), RuntimeType::Fun { .. })
        )
}

fn apply_adapter_source(source: TypeSource) -> bool {
    matches!(
        source,
        TypeSource::ApplyEvidence
            | TypeSource::ApplyCalleeEvidence
            | TypeSource::ApplyArgumentEvidence
            | TypeSource::ApplyArgumentSourceEdge
    )
}

fn hir_type_contains_unknown(ty: &RuntimeType) -> bool {
    match ty {
        RuntimeType::Unknown => true,
        RuntimeType::Core(ty) => core_type_contains_unknown(ty),
        RuntimeType::Fun { param, ret } => {
            hir_type_contains_unknown(param) || hir_type_contains_unknown(ret)
        }
        RuntimeType::Thunk { effect, value } => {
            core_type_contains_unknown(effect) || hir_type_contains_unknown(value)
        }
    }
}

fn core_type_contains_unknown(ty: &typed_ir::Type) -> bool {
    match ty {
        typed_ir::Type::Unknown => true,
        typed_ir::Type::Never | typed_ir::Type::Any | typed_ir::Type::Var(_) => false,
        typed_ir::Type::Named { args, .. } => args.iter().any(type_arg_contains_unknown),
        typed_ir::Type::Fun {
            param,
            param_effect,
            ret_effect,
            ret,
        } => {
            core_type_contains_unknown(param)
                || core_type_contains_unknown(param_effect)
                || core_type_contains_unknown(ret_effect)
                || core_type_contains_unknown(ret)
        }
        typed_ir::Type::Tuple(items)
        | typed_ir::Type::Union(items)
        | typed_ir::Type::Inter(items) => items.iter().any(core_type_contains_unknown),
        typed_ir::Type::Record(record) => {
            record
                .fields
                .iter()
                .any(|field| core_type_contains_unknown(&field.value))
                || record.spread.as_ref().is_some_and(|spread| match spread {
                    typed_ir::RecordSpread::Head(ty) | typed_ir::RecordSpread::Tail(ty) => {
                        core_type_contains_unknown(ty)
                    }
                })
        }
        typed_ir::Type::Variant(variant) => {
            variant
                .cases
                .iter()
                .any(|case| case.payloads.iter().any(core_type_contains_unknown))
                || variant
                    .tail
                    .as_deref()
                    .is_some_and(core_type_contains_unknown)
        }
        typed_ir::Type::Row { items, tail } => {
            items.iter().any(core_type_contains_unknown) || core_type_contains_unknown(tail)
        }
        typed_ir::Type::Recursive { body, .. } => core_type_contains_unknown(body),
    }
}

fn type_arg_contains_unknown(arg: &typed_ir::TypeArg) -> bool {
    match arg {
        typed_ir::TypeArg::Type(ty) => core_type_contains_unknown(ty),
        typed_ir::TypeArg::Bounds(bounds) => {
            bounds
                .lower
                .as_deref()
                .is_some_and(core_type_contains_unknown)
                || bounds
                    .upper
                    .as_deref()
                    .is_some_and(core_type_contains_unknown)
        }
    }
}

fn apply_arg_source_edge(source: TypeSource) -> bool {
    matches!(source, TypeSource::ApplyArgumentSourceEdge)
}

fn record_apply_adapter(
    profile: &mut RuntimeAdapterProfile,
    kind: RuntimeAdapterEventKind,
    source: TypeSource,
    adapter_source: Option<RuntimeAdapterSource>,
    actual: &RuntimeType,
    expected: &RuntimeType,
) {
    let phase = adapter_source
        .as_ref()
        .map(|source| {
            source.profile_apply_adapter(profile);
            source.phase
        })
        .or_else(|| apply_adapter_phase(source));
    let Some(phase) = phase else {
        return;
    };
    record_apply_adapter_phase(profile, phase, kind);
    if profile.collect_events
        && let Some(adapter_source) = &adapter_source
    {
        profile.events.push(RuntimeAdapterEvent {
            kind,
            phase,
            owner: adapter_source.owner.clone(),
            apply_target: adapter_source.apply_target.clone(),
            callee_source_edge: adapter_source.callee_source_edge,
            arg_source_edge: adapter_source.arg_source_edge,
            actual: actual.clone(),
            expected: expected.clone(),
        });
    }
    if adapter_source.is_none() && apply_arg_source_edge(source) {
        profile.apply_evidence_adapter_with_source_edge += 1;
    }
}

fn apply_adapter_phase(source: TypeSource) -> Option<RuntimeApplyAdapterPhase> {
    match source {
        TypeSource::ApplyCalleeEvidence => Some(RuntimeApplyAdapterPhase::LowerCallee),
        TypeSource::ApplyArgumentEvidence | TypeSource::ApplyArgumentSourceEdge => {
            Some(RuntimeApplyAdapterPhase::LowerArgument)
        }
        TypeSource::ApplyEvidence => None,
        _ => None,
    }
}

fn record_apply_adapter_phase(
    profile: &mut RuntimeAdapterProfile,
    phase: RuntimeApplyAdapterPhase,
    kind: RuntimeAdapterEventKind,
) {
    match (phase, kind) {
        (RuntimeApplyAdapterPhase::LowerCallee, RuntimeAdapterEventKind::ValueToThunk) => {
            profile.apply_lower_callee_value_to_thunk += 1;
        }
        (RuntimeApplyAdapterPhase::LowerCallee, RuntimeAdapterEventKind::ThunkToValue) => {
            profile.apply_lower_callee_thunk_to_value += 1;
        }
        (RuntimeApplyAdapterPhase::LowerCallee, RuntimeAdapterEventKind::BindHere) => {
            profile.apply_lower_callee_bind_here += 1;
        }
        (RuntimeApplyAdapterPhase::LowerArgument, RuntimeAdapterEventKind::ValueToThunk) => {
            profile.apply_lower_argument_value_to_thunk += 1;
        }
        (RuntimeApplyAdapterPhase::LowerArgument, RuntimeAdapterEventKind::ThunkToValue) => {
            profile.apply_lower_argument_thunk_to_value += 1;
        }
        (RuntimeApplyAdapterPhase::LowerArgument, RuntimeAdapterEventKind::BindHere) => {
            profile.apply_lower_argument_bind_here += 1;
        }
        (RuntimeApplyAdapterPhase::PrepareFinalArgument, RuntimeAdapterEventKind::ValueToThunk) => {
            profile.apply_prepare_final_argument_value_to_thunk += 1;
        }
        (RuntimeApplyAdapterPhase::PrepareFinalArgument, RuntimeAdapterEventKind::ThunkToValue) => {
            profile.apply_prepare_final_argument_thunk_to_value += 1;
        }
        (RuntimeApplyAdapterPhase::PrepareFinalArgument, RuntimeAdapterEventKind::BindHere) => {
            profile.apply_prepare_final_argument_bind_here += 1;
        }
        (
            RuntimeApplyAdapterPhase::PrepareEffectOperationArgument,
            RuntimeAdapterEventKind::ValueToThunk,
        ) => {
            profile.apply_prepare_effect_operation_argument_value_to_thunk += 1;
        }
        (
            RuntimeApplyAdapterPhase::PrepareEffectOperationArgument,
            RuntimeAdapterEventKind::ThunkToValue,
        ) => {
            profile.apply_prepare_effect_operation_argument_thunk_to_value += 1;
        }
        (
            RuntimeApplyAdapterPhase::PrepareEffectOperationArgument,
            RuntimeAdapterEventKind::BindHere,
        ) => {
            profile.apply_prepare_effect_operation_argument_bind_here += 1;
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) struct RuntimeAdapterSource {
    pub phase: RuntimeApplyAdapterPhase,
    pub has_apply_evidence: bool,
    pub has_apply_callee_source_edge: bool,
    pub has_apply_arg_source_edge: bool,
    pub callee_source_edge: Option<u32>,
    pub arg_source_edge: Option<u32>,
    pub owner: Option<typed_ir::Path>,
    pub apply_target: Option<typed_ir::Path>,
}

impl RuntimeAdapterSource {
    fn profile_apply_adapter(&self, profile: &mut RuntimeAdapterProfile) {
        if self.has_apply_evidence {
            profile.apply_evidence_adapter_with_evidence += 1;
        } else {
            profile.apply_evidence_adapter_without_evidence += 1;
        }
        if self.has_apply_callee_source_edge || self.has_apply_arg_source_edge {
            profile.apply_evidence_adapter_with_source_edge += 1;
        }
    }
}

pub(super) fn finalize_effectful_expr_profiled(
    expr: Expr,
    expected: Option<&RuntimeType>,
    source: TypeSource,
    profile: &mut RuntimeAdapterProfile,
    adapter_source: Option<RuntimeAdapterSource>,
) -> RuntimeResult<Expr> {
    let expr = attach_forced_effect_profiled(expr, profile);
    match expected {
        Some(expected) => prepare_expr_for_expected_with_adapter_source_profiled(
            expr,
            expected,
            source,
            profile,
            adapter_source,
        ),
        None => Ok(expr),
    }
}

pub(super) fn finalize_handler_expr_profiled(
    expr: Expr,
    expected: Option<&RuntimeType>,
    source: TypeSource,
    profile: &mut RuntimeAdapterProfile,
    adapter_source: Option<RuntimeAdapterSource>,
) -> RuntimeResult<Expr> {
    let expr = attach_forced_effect_profiled(expr, profile);
    match (expected, &expr.ty) {
        (
            Some(RuntimeType::Thunk {
                value: expected_value,
                ..
            }),
            actual,
        ) if hir_type_compatible(expected_value, actual) => {
            require_same_hir_type(expected_value, actual, source)?;
            Ok(expr)
        }
        (Some(expected), _) => prepare_expr_for_expected_with_adapter_source_profiled(
            expr,
            expected,
            source,
            profile,
            adapter_source,
        ),
        (None, _) => Ok(expr),
    }
}

pub(super) fn attach_forced_effect_profiled(
    expr: Expr,
    profile: &mut RuntimeAdapterProfile,
) -> Expr {
    match expr_forced_effect(&expr) {
        Some(effect) => {
            let effect = project_runtime_effect(&effect);
            if should_thunk_effect(&effect) {
                profile.forced_effect_thunk += 1;
                attach_expr_effect(expr, effect)
            } else {
                expr
            }
        }
        _ => expr,
    }
}

pub(super) fn attach_expr_effect(expr: Expr, effect: typed_ir::Type) -> Expr {
    match expr.ty.clone() {
        RuntimeType::Thunk {
            effect: existing,
            value,
        } => {
            let effect = project_runtime_effect(&merge_effect_rows(effect, existing));
            let ty = RuntimeType::thunk(effect.clone(), (*value).clone());
            let kind = match expr.kind {
                ExprKind::Thunk {
                    value, expr: inner, ..
                } => ExprKind::Thunk {
                    effect,
                    value,
                    expr: inner,
                },
                other => other,
            };
            Expr { ty, kind }
        }
        value => Expr::typed(
            ExprKind::Thunk {
                effect: effect.clone(),
                value: value.clone(),
                expr: Box::new(expr),
            },
            RuntimeType::thunk(effect, value),
        ),
    }
}

pub(super) fn add_id_to_created_thunks(expr: Expr) -> Expr {
    let ty = expr.ty;
    let kind = match expr.kind {
        ExprKind::Thunk {
            effect,
            value,
            expr,
        } => {
            let inner = add_id_to_created_thunks(*expr);
            let thunk = Expr::typed(
                ExprKind::Thunk {
                    effect: effect.clone(),
                    value,
                    expr: Box::new(inner),
                },
                ty,
            );
            return add_id_with_peek_if_needed(thunk, effect);
        }
        ExprKind::Lambda {
            param,
            param_effect_annotation,
            param_function_allowed_effects,
            body,
        } => ExprKind::Lambda {
            param,
            param_effect_annotation,
            param_function_allowed_effects,
            body,
        },
        ExprKind::Apply {
            callee,
            arg,
            evidence,
            instantiation,
        } => ExprKind::Apply {
            callee: Box::new(add_id_to_created_thunks(*callee)),
            arg: Box::new(add_id_to_created_thunks(*arg)),
            evidence,
            instantiation,
        },
        ExprKind::If {
            cond,
            then_branch,
            else_branch,
            evidence,
        } => ExprKind::If {
            cond: Box::new(add_id_to_created_thunks(*cond)),
            then_branch: Box::new(add_id_to_created_thunks(*then_branch)),
            else_branch: Box::new(add_id_to_created_thunks(*else_branch)),
            evidence,
        },
        ExprKind::Tuple(items) => {
            ExprKind::Tuple(items.into_iter().map(add_id_to_created_thunks).collect())
        }
        ExprKind::Record { fields, spread } => ExprKind::Record {
            fields: fields
                .into_iter()
                .map(|field| RecordExprField {
                    name: field.name,
                    value: add_id_to_created_thunks(field.value),
                })
                .collect(),
            spread: spread.map(|spread| match spread {
                RecordSpreadExpr::Head(expr) => {
                    RecordSpreadExpr::Head(Box::new(add_id_to_created_thunks(*expr)))
                }
                RecordSpreadExpr::Tail(expr) => {
                    RecordSpreadExpr::Tail(Box::new(add_id_to_created_thunks(*expr)))
                }
            }),
        },
        ExprKind::Variant { tag, value } => ExprKind::Variant {
            tag,
            value: value.map(|value| Box::new(add_id_to_created_thunks(*value))),
        },
        ExprKind::Select { base, field } => ExprKind::Select {
            base: Box::new(add_id_to_created_thunks(*base)),
            field,
        },
        ExprKind::Match {
            scrutinee,
            arms,
            evidence,
        } => ExprKind::Match {
            scrutinee: Box::new(add_id_to_created_thunks(*scrutinee)),
            arms: arms
                .into_iter()
                .map(|arm| MatchArm {
                    pattern: arm.pattern,
                    guard: arm.guard.map(add_id_to_created_thunks),
                    body: add_id_to_created_thunks(arm.body),
                })
                .collect(),
            evidence,
        },
        ExprKind::Block { stmts, tail } => ExprKind::Block {
            stmts: stmts
                .into_iter()
                .map(|stmt| match stmt {
                    Stmt::Let { pattern, value } => Stmt::Let {
                        pattern,
                        value: add_id_to_created_thunks(value),
                    },
                    Stmt::Expr(expr) => Stmt::Expr(add_id_to_created_thunks(expr)),
                    Stmt::Module { def, body } => Stmt::Module {
                        def,
                        body: add_id_to_created_thunks(body),
                    },
                })
                .collect(),
            tail: tail.map(|tail| Box::new(add_id_to_created_thunks(*tail))),
        },
        ExprKind::Handle {
            body,
            arms,
            evidence,
            handler,
        } => ExprKind::Handle {
            body: Box::new(add_id_to_created_thunks(*body)),
            arms: arms
                .into_iter()
                .map(|arm| HandleArm {
                    effect: arm.effect,
                    payload: arm.payload,
                    resume: arm.resume,
                    guard: arm.guard.map(add_id_to_created_thunks),
                    body: add_id_to_created_thunks(arm.body),
                })
                .collect(),
            evidence,
            handler,
        },
        ExprKind::BindHere { expr } => ExprKind::BindHere {
            expr: Box::new(add_id_to_created_thunks(*expr)),
        },
        ExprKind::LocalPushId { id, body } => ExprKind::LocalPushId {
            id,
            body: Box::new(add_id_to_created_thunks(*body)),
        },
        ExprKind::FindId { id } => ExprKind::FindId { id },
        ExprKind::AddId {
            id,
            allowed,
            active,
            thunk,
        } => ExprKind::AddId {
            id,
            allowed,
            active,
            thunk: Box::new(add_id_to_created_thunks(*thunk)),
        },
        ExprKind::Coerce { from, to, expr } => ExprKind::Coerce {
            from,
            to,
            expr: Box::new(add_id_to_created_thunks(*expr)),
        },
        ExprKind::Pack { var, expr } => ExprKind::Pack {
            var,
            expr: Box::new(add_id_to_created_thunks(*expr)),
        },
        kind @ (ExprKind::Var(_)
        | ExprKind::EffectOp(_)
        | ExprKind::PrimitiveOp(_)
        | ExprKind::Lit(_)
        | ExprKind::PeekId) => kind,
    };
    Expr { ty, kind }
}

pub(super) fn add_id_with_peek_if_needed(thunk: Expr, allowed: typed_ir::Type) -> Expr {
    let allowed = project_runtime_effect(&allowed);
    if !should_thunk_effect(&allowed) {
        return thunk;
    }
    add_id_with_peek(thunk, allowed, false)
}

pub(super) fn add_boundary_id_with_peek(thunk: Expr, allowed: typed_ir::Type) -> Expr {
    add_boundary_id(thunk, EffectIdRef::Peek, allowed)
}

pub(super) fn add_boundary_id(thunk: Expr, id: EffectIdRef, allowed: typed_ir::Type) -> Expr {
    let allowed = project_runtime_effect(&allowed);
    add_id_required(thunk, id, allowed, true)
}

fn add_id_with_peek(thunk: Expr, allowed: typed_ir::Type, active: bool) -> Expr {
    if !matches!(thunk.ty, RuntimeType::Thunk { .. }) {
        return thunk;
    }
    Expr::typed(
        ExprKind::AddId {
            id: EffectIdRef::Peek,
            allowed,
            active,
            thunk: Box::new(thunk.clone()),
        },
        thunk.ty,
    )
}

fn add_id_required(thunk: Expr, id: EffectIdRef, allowed: typed_ir::Type, active: bool) -> Expr {
    if !matches!(thunk.ty, RuntimeType::Thunk { .. }) {
        return thunk;
    }
    Expr::typed(
        ExprKind::AddId {
            id,
            allowed,
            active,
            thunk: Box::new(thunk.clone()),
        },
        thunk.ty,
    )
}

pub(super) fn contains_peek_add_id(expr: &Expr) -> bool {
    match &expr.kind {
        ExprKind::AddId {
            id: EffectIdRef::Peek,
            ..
        } => true,
        ExprKind::Lambda { .. } => false,
        ExprKind::Apply { callee, arg, .. } => {
            contains_peek_add_id(callee) || contains_peek_add_id(arg)
        }
        ExprKind::If {
            cond,
            then_branch,
            else_branch,
            ..
        } => {
            contains_peek_add_id(cond)
                || contains_peek_add_id(then_branch)
                || contains_peek_add_id(else_branch)
        }
        ExprKind::Tuple(items) => items.iter().any(contains_peek_add_id),
        ExprKind::Record { fields, spread } => {
            fields
                .iter()
                .any(|field| contains_peek_add_id(&field.value))
                || spread.as_ref().is_some_and(|spread| match spread {
                    RecordSpreadExpr::Head(expr) | RecordSpreadExpr::Tail(expr) => {
                        contains_peek_add_id(expr)
                    }
                })
        }
        ExprKind::Variant { value, .. } => value.as_deref().is_some_and(contains_peek_add_id),
        ExprKind::Select { base, .. } => contains_peek_add_id(base),
        ExprKind::Match {
            scrutinee, arms, ..
        } => {
            contains_peek_add_id(scrutinee)
                || arms.iter().any(|arm| {
                    arm.guard.as_ref().is_some_and(contains_peek_add_id)
                        || contains_peek_add_id(&arm.body)
                })
        }
        ExprKind::Block { stmts, tail } => {
            stmts.iter().any(|stmt| match stmt {
                Stmt::Let { value, .. } | Stmt::Expr(value) => contains_peek_add_id(value),
                Stmt::Module { body, .. } => contains_peek_add_id(body),
            }) || tail.as_deref().is_some_and(contains_peek_add_id)
        }
        ExprKind::Handle { body, arms, .. } => {
            contains_peek_add_id(body)
                || arms.iter().any(|arm| {
                    arm.guard.as_ref().is_some_and(contains_peek_add_id)
                        || contains_peek_add_id(&arm.body)
                })
        }
        ExprKind::BindHere { expr }
        | ExprKind::Thunk { expr, .. }
        | ExprKind::Coerce { expr, .. }
        | ExprKind::Pack { expr, .. } => contains_peek_add_id(expr),
        ExprKind::LocalPushId { body, .. } => contains_peek_add_id(body),
        ExprKind::AddId { thunk, .. } => contains_peek_add_id(thunk),
        ExprKind::Var(_)
        | ExprKind::EffectOp(_)
        | ExprKind::PrimitiveOp(_)
        | ExprKind::Lit(_)
        | ExprKind::PeekId
        | ExprKind::FindId { .. } => false,
    }
}

pub(super) fn contains_effect_id_var(expr: &Expr, target: EffectIdVar) -> bool {
    match &expr.kind {
        ExprKind::AddId {
            id: EffectIdRef::Var(id),
            ..
        }
        | ExprKind::FindId {
            id: EffectIdRef::Var(id),
        } if *id == target => true,
        ExprKind::Lambda { body, .. } => contains_effect_id_var(body, target),
        ExprKind::Apply { callee, arg, .. } => {
            contains_effect_id_var(callee, target) || contains_effect_id_var(arg, target)
        }
        ExprKind::If {
            cond,
            then_branch,
            else_branch,
            ..
        } => {
            contains_effect_id_var(cond, target)
                || contains_effect_id_var(then_branch, target)
                || contains_effect_id_var(else_branch, target)
        }
        ExprKind::Tuple(items) => items
            .iter()
            .any(|item| contains_effect_id_var(item, target)),
        ExprKind::Record { fields, spread } => {
            fields
                .iter()
                .any(|field| contains_effect_id_var(&field.value, target))
                || spread.as_ref().is_some_and(|spread| match spread {
                    RecordSpreadExpr::Head(expr) | RecordSpreadExpr::Tail(expr) => {
                        contains_effect_id_var(expr, target)
                    }
                })
        }
        ExprKind::Variant { value, .. } => value
            .as_deref()
            .is_some_and(|value| contains_effect_id_var(value, target)),
        ExprKind::Select { base, .. } => contains_effect_id_var(base, target),
        ExprKind::Match {
            scrutinee, arms, ..
        } => {
            contains_effect_id_var(scrutinee, target)
                || arms.iter().any(|arm| {
                    arm.guard
                        .as_ref()
                        .is_some_and(|guard| contains_effect_id_var(guard, target))
                        || contains_effect_id_var(&arm.body, target)
                })
        }
        ExprKind::Block { stmts, tail } => {
            stmts.iter().any(|stmt| match stmt {
                Stmt::Let { value, .. } | Stmt::Expr(value) => {
                    contains_effect_id_var(value, target)
                }
                Stmt::Module { body, .. } => contains_effect_id_var(body, target),
            }) || tail
                .as_deref()
                .is_some_and(|tail| contains_effect_id_var(tail, target))
        }
        ExprKind::Handle { body, arms, .. } => {
            contains_effect_id_var(body, target)
                || arms.iter().any(|arm| {
                    arm.guard
                        .as_ref()
                        .is_some_and(|guard| contains_effect_id_var(guard, target))
                        || contains_effect_id_var(&arm.body, target)
                })
        }
        ExprKind::BindHere { expr }
        | ExprKind::Thunk { expr, .. }
        | ExprKind::Coerce { expr, .. }
        | ExprKind::Pack { expr, .. } => contains_effect_id_var(expr, target),
        ExprKind::LocalPushId { body, .. } => contains_effect_id_var(body, target),
        ExprKind::AddId { thunk, .. } => contains_effect_id_var(thunk, target),
        ExprKind::Var(_)
        | ExprKind::EffectOp(_)
        | ExprKind::PrimitiveOp(_)
        | ExprKind::Lit(_)
        | ExprKind::PeekId
        | ExprKind::FindId { .. } => false,
    }
}