runmat-mir 0.5.0

use crate::{
    MirOutputTarget, MirOutputTargetList, MirPlaceMutation, MirRvalue, MirStmt, MirStmtKind,
};
use runmat_hir::{
    AssignmentCreationPolicy, AssignmentShapePolicy, EnvironmentEffect, ExprId, HirError, HirExpr,
    HirExprKind, HirStmt, HirStmtKind, OutputTarget, PlaceMutationKind, RequestedOutputCount, Span,
    WorkspaceEffect,
};
use std::collections::HashMap;

use super::{expr::lower_expr_with_replacements, place::lower_place, MirLoweringContext};

pub(crate) fn lower_stmt_with_replacements(
    ctx: &MirLoweringContext,
    stmt: &HirStmt,
    await_replacements: &HashMap<ExprId, crate::MirOperand>,
) -> Result<Vec<MirStmt>, HirError> {
    Ok(match &stmt.kind {
        HirStmtKind::Assign(place, expr, _) => {
            let mut stmts = Vec::new();
            let value = lower_expr_with_replacements(ctx, expr, &mut stmts, await_replacements)?;
            stmts.extend(effect_stmts_for_rvalue(&value, stmt.span));
            let place = lower_place(ctx, place, &mut stmts)?;
            let deletion = is_empty_array_deletion_place(&place, expr);
            if !matches!(place, crate::MirPlace::Local(_)) || deletion {
                stmts.push(MirStmt {
                    kind: MirStmtKind::PlaceMutation(place_mutation(place.clone(), deletion)),
                    span: stmt.span,
                });
            }
            stmts.push(MirStmt {
                kind: MirStmtKind::Assign { place, value },
                span: stmt.span,
            });
            stmts
        }
        HirStmtKind::MultiAssign(targets, expr, _) => {
            let mut stmts = Vec::new();
            let value = lower_expr_with_replacements(ctx, expr, &mut stmts, await_replacements)?;
            stmts.extend(effect_stmts_for_rvalue(&value, stmt.span));
            let lowered_targets = lower_output_targets(ctx, &targets.targets, &mut stmts)?;
            let lowered_target_list = MirOutputTargetList {
                requested_outputs: targets.requested_outputs.clone(),
                targets: lowered_targets,
            };
            let requested_outputs = lowered_target_list
                .validate_fixed_arity("MIR multi-assign")
                .map_err(HirError::new)?;
            if let MirRvalue::Call(call) = &value {
                let call_outputs = fixed_requested_output_count(
                    &call.requested_outputs,
                    "MIR call requested output count",
                );
                if call_outputs != requested_outputs {
                    return Err(HirError::new(format!(
                        "MIR call requested outputs ({call_outputs}) must match multi-assign targets ({requested_outputs})"
                    )));
                }
            }
            stmts.push(MirStmt {
                kind: MirStmtKind::MultiAssign {
                    targets: lowered_target_list,
                    value,
                },
                span: stmt.span,
            });
            stmts
        }
        HirStmtKind::ExprStmt(expr, _) => {
            let mut stmts = Vec::new();
            let value = lower_expr_with_replacements(ctx, expr, &mut stmts, await_replacements)?;
            stmts.extend(effect_stmts_for_rvalue(&value, stmt.span));
            stmts.push(MirStmt {
                kind: MirStmtKind::Expr(value),
                span: stmt.span,
            });
            stmts
        }
        HirStmtKind::Global(bindings) => vec![MirStmt {
            kind: MirStmtKind::WorkspaceEffect {
                effect: WorkspaceEffect::MutatesGlobal,
                bindings: bindings
                    .iter()
                    .map(|binding| ctx.local_for_binding(*binding))
                    .collect::<Result<_, _>>()?,
            },
            span: stmt.span,
        }],
        HirStmtKind::Persistent(bindings) => vec![MirStmt {
            kind: MirStmtKind::WorkspaceEffect {
                effect: WorkspaceEffect::MutatesPersistent,
                bindings: bindings
                    .iter()
                    .map(|binding| ctx.local_for_binding(*binding))
                    .collect::<Result<_, _>>()?,
            },
            span: stmt.span,
        }],
        HirStmtKind::Return | HirStmtKind::Import(_) => Vec::new(),
        _ => {
            return Err(HirError::new(format!(
                "unexpected statement in stmt-level MIR lowering; control-flow lowering should have handled: {:?}",
                stmt.kind
            )))
        }
    })
}

fn effect_stmts_for_rvalue(value: &MirRvalue, span: Span) -> Vec<MirStmt> {
    let MirRvalue::Call(call) = value else {
        return Vec::new();
    };
    let mut stmts = Vec::new();
    if let Some(effect) = call.workspace_effect.map(workspace_effect_from_builtin) {
        stmts.push(MirStmt {
            kind: MirStmtKind::WorkspaceEffect {
                effect,
                bindings: Vec::new(),
            },
            span,
        });
    }
    if let Some(effect) = call.environment_effect.map(environment_effect_from_builtin) {
        stmts.push(MirStmt {
            kind: MirStmtKind::EnvironmentEffect(effect),
            span,
        });
    }
    stmts
}

fn workspace_effect_from_builtin(
    effect: runmat_builtins::BuiltinWorkspaceEffect,
) -> WorkspaceEffect {
    match effect {
        runmat_builtins::BuiltinWorkspaceEffect::ReadsWorkspace => WorkspaceEffect::ReadsWorkspace,
        runmat_builtins::BuiltinWorkspaceEffect::CreatesBinding => WorkspaceEffect::CreatesBinding,
        runmat_builtins::BuiltinWorkspaceEffect::ClearsBinding => WorkspaceEffect::ClearsBinding,
        runmat_builtins::BuiltinWorkspaceEffect::ClearsFunctionCache => {
            WorkspaceEffect::ClearsFunctionCache
        }
        runmat_builtins::BuiltinWorkspaceEffect::LoadsExternalBindings => {
            WorkspaceEffect::LoadsExternalBindings
        }
        runmat_builtins::BuiltinWorkspaceEffect::DynamicEval => WorkspaceEffect::DynamicEval,
    }
}

fn environment_effect_from_builtin(
    effect: runmat_builtins::BuiltinEnvironmentEffect,
) -> EnvironmentEffect {
    match effect {
        runmat_builtins::BuiltinEnvironmentEffect::PathMutation => EnvironmentEffect::PathMutation,
        runmat_builtins::BuiltinEnvironmentEffect::WorkingDirectoryMutation => {
            EnvironmentEffect::WorkingDirectoryMutation
        }
        runmat_builtins::BuiltinEnvironmentEffect::FunctionCacheInvalidation => {
            EnvironmentEffect::FunctionCacheInvalidation
        }
        runmat_builtins::BuiltinEnvironmentEffect::DynamicLookupInvalidation => {
            EnvironmentEffect::DynamicLookupInvalidation
        }
    }
}

fn place_mutation(place: crate::MirPlace, deletion: bool) -> MirPlaceMutation {
    MirPlaceMutation {
        kind: mutation_kind_for_place(&place, deletion),
        creation_policy: creation_policy_for_place(&place, deletion),
        shape_policy: AssignmentShapePolicy::MatlabCompatible,
        place,
    }
}

fn mutation_kind_for_place(place: &crate::MirPlace, deletion: bool) -> PlaceMutationKind {
    if deletion {
        return PlaceMutationKind::Delete;
    }
    match place {
        crate::MirPlace::Local(_) | crate::MirPlace::Binding(_) => PlaceMutationKind::BindOrAssign,
        crate::MirPlace::Index(_, _) => PlaceMutationKind::IndexedAssign,
        crate::MirPlace::Member(_, _) | crate::MirPlace::DynamicMember(_, _) => {
            PlaceMutationKind::MemberAssign
        }
    }
}

fn creation_policy_for_place(place: &crate::MirPlace, deletion: bool) -> AssignmentCreationPolicy {
    if deletion {
        return AssignmentCreationPolicy::ExistingOnly;
    }
    match place {
        crate::MirPlace::Local(_) | crate::MirPlace::Binding(_) => {
            AssignmentCreationPolicy::CreateBinding
        }
        crate::MirPlace::Index(_, _) => AssignmentCreationPolicy::CreateArrayByIndex,
        crate::MirPlace::Member(_, _) | crate::MirPlace::DynamicMember(_, _) => {
            AssignmentCreationPolicy::CreateStructFieldPath
        }
    }
}

fn is_empty_array_expr(expr: &HirExpr) -> bool {
    matches!(&expr.kind, HirExprKind::Tensor(rows) if rows.is_empty() || rows.iter().all(Vec::is_empty))
}

fn is_empty_array_deletion_place(place: &crate::MirPlace, expr: &HirExpr) -> bool {
    if !is_empty_array_expr(expr) {
        return false;
    }
    matches!(
        place,
        crate::MirPlace::Index(
            _,
            crate::MirIndexing {
                kind: runmat_hir::IndexKind::Paren,
                ..
            }
        )
    )
}

fn lower_output_targets(
    ctx: &MirLoweringContext,
    targets: &[OutputTarget],
    stmts: &mut Vec<MirStmt>,
) -> Result<Vec<MirOutputTarget>, HirError> {
    targets
        .iter()
        .map(|target| lower_output_target(ctx, target, stmts))
        .collect()
}

fn lower_output_target(
    ctx: &MirLoweringContext,
    target: &OutputTarget,
    stmts: &mut Vec<MirStmt>,
) -> Result<MirOutputTarget, HirError> {
    Ok(match target {
        OutputTarget::Place(place) => {
            let mut temps = Vec::new();
            let place = lower_place(ctx, place, &mut temps)?;
            stmts.extend(temps);
            MirOutputTarget::Place(place)
        }
        OutputTarget::Discard => MirOutputTarget::Discard,
    })
}

fn fixed_requested_output_count(requested_outputs: &RequestedOutputCount, _context: &str) -> usize {
    requested_outputs.fixed_count()
}