ruka_mir/

cfg.rs

use std::collections::BTreeMap;

use cranelift_entity::PrimaryMap;
use ruka_types::Ty;

use crate::{
    LocalInfo, LocalKind, MirBlockId, MirCallArg, MirCallArgBinding, MirEnumDecl, MirFuncId,
    MirHeapOwnership, MirInstr, MirLocalId, MirLocalRepr, MirOwnershipMode, MirStructDecl,
};

/// Whole-program MIR container.
///
/// `function_names` provides fast lookup from source-level function name
/// to MIR function id.
#[derive(Debug, Clone)]
pub struct MirProgram {
    /// Functions stored in the MIR program arena.
    pub functions: PrimaryMap<MirFuncId, MirFunction>,
    /// Lookup table from canonical function name to function id.
    pub function_names: BTreeMap<String, MirFuncId>,
    /// Struct declarations referenced by MIR.
    pub structs: Vec<MirStructDecl>,
    /// Enum declarations referenced by MIR.
    pub enums: Vec<MirEnumDecl>,
}

impl MirProgram {
    /// Return the number of functions stored in the MIR program.
    pub fn function_count(&self) -> usize {
        self.functions.len()
    }
}

/// Function-level MIR in CFG form.
///
/// - `params` are function parameter locals.
/// - `param_modes` mirrors source ownership modes for parameters.
/// - `entry` is the first block reached for execution.
#[derive(Debug, Clone)]
pub struct MirFunction {
    /// User-visible function or field name.
    pub name: String,
    /// Number of source parameters accepted by the function.
    pub arity: usize,
    /// Semantic return type of the function.
    pub return_ty: Ty,
    /// Local metadata arena for function values.
    pub locals: PrimaryMap<MirLocalId, LocalInfo>,
    /// Local ids corresponding to incoming parameters.
    pub params: Vec<MirLocalId>,
    /// Ownership modes for each incoming parameter.
    pub param_modes: Vec<MirOwnershipMode>,
    /// Entry block id for the function CFG.
    pub entry: MirBlockId,
    /// CFG blocks stored for the function.
    pub blocks: PrimaryMap<MirBlockId, MirBlock>,
}

/// Shared parameter query object used by lowering and backends.
#[derive(Debug, Clone, Copy)]
pub struct MirParamBinding<'a> {
    /// Positional parameter index in the source signature.
    pub index: usize,
    /// MIR local receiving the parameter.
    pub local_id: MirLocalId,
    /// Source ownership mode for the parameter.
    pub mode: MirOwnershipMode,
    /// Local metadata recorded for the parameter.
    pub local: &'a LocalInfo,
}

impl<'a> MirParamBinding<'a> {
    /// Return whether the parameter expects read-only view access.
    pub fn expects_view(self) -> bool {
        matches!(self.mode, MirOwnershipMode::View)
    }

    /// Return whether the parameter expects mutable borrow access.
    pub fn expects_mut_borrow(self) -> bool {
        matches!(self.mode, MirOwnershipMode::MutBorrow)
    }

    /// Return whether the parameter expects owned value access.
    pub fn expects_owned(self) -> bool {
        matches!(self.mode, MirOwnershipMode::Owned)
    }

    /// Return the semantic item type exposed at the source-language boundary.
    pub fn semantic_ty(self) -> &'a Ty {
        match self.mode {
            MirOwnershipMode::View | MirOwnershipMode::MutBorrow => {
                self.local.place_item_ty().unwrap_or(&self.local.ty)
            }
            MirOwnershipMode::Owned => &self.local.ty,
        }
    }

    /// Return the source-language representation expected at the function boundary.
    pub fn source_repr(self) -> MirLocalRepr {
        match self.mode {
            MirOwnershipMode::View => MirLocalRepr::SharedPlace,
            MirOwnershipMode::MutBorrow => MirLocalRepr::MutablePlace,
            MirOwnershipMode::Owned => MirLocalRepr::Value,
        }
    }

    /// Return the lowered storage/runtime representation of the parameter local.
    pub fn local_repr(self) -> MirLocalRepr {
        self.local.repr
    }

    /// Return whether the lowered local needs materialization from the source boundary shape.
    pub fn requires_materialization(self) -> bool {
        self.source_repr() != self.local_repr()
    }

    /// Return whether this binding is the view-from-owned materialization case.
    pub fn materializes_view_from_owned(self) -> bool {
        self.expects_view() && self.requires_materialization()
    }
}

impl MirFunction {
    /// Return metadata for one parameter by index.
    pub fn param_binding(&self, index: usize) -> MirParamBinding<'_> {
        let local_id = self.params[index];
        MirParamBinding {
            index,
            local_id,
            mode: self.param_modes[index],
            local: &self.locals[local_id],
        }
    }

    /// Iterate over source parameters with normalized query helpers.
    pub fn param_bindings(&self) -> impl Iterator<Item = MirParamBinding<'_>> + '_ {
        self.params
            .iter()
            .enumerate()
            .map(|(index, _)| self.param_binding(index))
    }

    /// Return metadata for one local.
    pub fn local_info(&self, local_id: MirLocalId) -> &LocalInfo {
        &self.locals[local_id]
    }

    /// Return normalized metadata for one call argument.
    pub fn call_arg_binding<'a>(&'a self, arg: &'a MirCallArg) -> MirCallArgBinding<'a> {
        MirCallArgBinding {
            arg,
            local: &self.locals[arg.local],
        }
    }

    /// Assert internal MIR invariants needed by downstream backends.
    pub fn assert_valid(&self) {
        assert_eq!(
            self.params.len(),
            self.param_modes.len(),
            "function `{}` has mismatched params and param modes",
            self.name,
        );

        for binding in self.param_bindings() {
            let info = binding.local;
            assert!(
                matches!(info.kind, LocalKind::Param),
                "function `{}` has non-param local {:?} in params list",
                self.name,
                binding.local_id,
            );
            match binding.mode {
                MirOwnershipMode::View => assert!(
                    !info.repr.is_mutable_place(),
                    "function `{}` has mutable view param {:?}",
                    self.name,
                    binding.local_id,
                ),
                MirOwnershipMode::MutBorrow => assert!(
                    info.repr.is_mutable_place(),
                    "function `{}` has non-mutable borrowed param {:?}",
                    self.name,
                    binding.local_id,
                ),
                MirOwnershipMode::Owned => assert!(
                    !info.repr.is_place(),
                    "function `{}` has place-shaped owned param {:?}",
                    self.name,
                    binding.local_id,
                ),
            }
        }

        for (local_id, info) in self.locals.iter() {
            match info.heap_ownership {
                MirHeapOwnership::None => {}
                MirHeapOwnership::Owned | MirHeapOwnership::OwnedShallow => assert!(
                    !info.repr.is_place(),
                    "function `{}` marks place local {:?} as heap-owned",
                    self.name,
                    local_id,
                ),
                MirHeapOwnership::BorrowedView => {
                    assert!(
                        info.repr.is_place(),
                        "function `{}` marks non-place local {:?} as BorrowedView",
                        self.name,
                        local_id,
                    );
                    assert!(
                        matches!(info.place_item_ty(), Some(Ty::Slice(_))),
                        "function `{}` marks non-slice place local {:?} as BorrowedView",
                        self.name,
                        local_id,
                    );
                }
            }
        }

        for (_, block) in self.blocks.iter() {
            for local_id in &block.params {
                let _ = &self.locals[*local_id];
            }
            for instr in &block.instrs {
                self.assert_instr_valid(instr);
            }
            match &block.terminator {
                MirTerminator::Jump { target, args } => {
                    self.assert_edge_args_compatible(
                        self.blocks[*target].params.as_slice(),
                        args.as_slice(),
                        "jump",
                    );
                }
                MirTerminator::Branch {
                    cond,
                    then_target,
                    then_args,
                    else_target,
                    else_args,
                    ..
                } => {
                    assert!(
                        !self.locals[*cond].repr.is_place(),
                        "function `{}` branches on place local {:?}",
                        self.name,
                        cond,
                    );
                    let then_params = &self.blocks[*then_target].params;
                    self.assert_edge_args_compatible(
                        then_params.as_slice(),
                        then_args.as_slice(),
                        "then",
                    );
                    let else_params = &self.blocks[*else_target].params;
                    self.assert_edge_args_compatible(
                        else_params.as_slice(),
                        else_args.as_slice(),
                        "else",
                    );
                }
                MirTerminator::Return { value } => {
                    assert!(
                        !self.locals[*value].repr.is_place(),
                        "function `{}` returns place local {:?}",
                        self.name,
                        value,
                    );
                    assert_eq!(
                        self.locals[*value].ty, self.return_ty,
                        "function `{}` returns local {:?} with mismatched type",
                        self.name, value,
                    );
                }
            }
        }
    }

    /// Assert that edge arguments match the target block parameter contract.
    fn assert_edge_args_compatible(
        &self,
        params: &[MirLocalId],
        args: &[MirLocalId],
        edge_kind: &'static str,
    ) {
        assert_eq!(
            params.len(),
            args.len(),
            "function `{}` has {edge_kind} edge with mismatched block arg count",
            self.name,
        );
        for (param, arg) in params.iter().zip(args.iter()) {
            let param_info = &self.locals[*param];
            let arg_info = &self.locals[*arg];
            assert_eq!(
                param_info.repr, arg_info.repr,
                "function `{}` has {edge_kind} edge with repr mismatch into block param {:?}",
                self.name, param,
            );
            assert_eq!(
                param_info.ty, arg_info.ty,
                "function `{}` has {edge_kind} edge with type mismatch into block param {:?}",
                self.name, param,
            );
        }
    }

    /// Assert per-instruction MIR invariants.
    fn assert_instr_valid(&self, instr: &MirInstr) {
        match instr {
            MirInstr::StoreRef { dst_ref, .. } => assert!(
                self.locals[*dst_ref].repr.is_mutable_place(),
                "function `{}` stores through non-mutable place {:?}",
                self.name,
                dst_ref,
            ),
            MirInstr::ReleaseHeap { local } => {
                assert!(
                    !self.locals[*local].repr.is_place(),
                    "function `{}` applies heap ownership op to place local {:?}",
                    self.name,
                    local,
                );
                assert!(
                    self.locals[*local].uses_heap_ops(),
                    "function `{}` applies heap ownership op to non-heap local {:?}",
                    self.name,
                    local,
                );
            }
            MirInstr::DerefCopy { src, dst } => {
                assert!(
                    self.locals[*src].repr.is_place(),
                    "function `{}` dereferences non-place local {:?}",
                    self.name,
                    src,
                );
                assert!(
                    !self.locals[*dst].repr.is_place(),
                    "function `{}` deref-copies into place local {:?}",
                    self.name,
                    dst,
                );
            }
            MirInstr::Call { args, dst, .. }
            | MirInstr::CallExtern { args, dst, .. }
            | MirInstr::CallIntrinsic { args, dst, .. } => {
                assert!(
                    !self.locals[*dst].repr.is_place(),
                    "function `{}` writes call result into place {:?}",
                    self.name,
                    dst,
                );
                self.assert_call_args_valid(args.as_slice());
            }
            MirInstr::EnumIsVariant { value, dst, .. }
            | MirInstr::EnumGetField { value, dst, .. } => {
                assert!(
                    !self.locals[*value].repr.is_place(),
                    "function `{}` reads enum metadata from place local {:?}",
                    self.name,
                    value,
                );
                assert!(
                    !self.locals[*dst].repr.is_place(),
                    "function `{}` writes enum metadata into place {:?}",
                    self.name,
                    dst,
                );
            }
            MirInstr::NumCast { src, dst } | MirInstr::CheckedIntCast { src, dst } => {
                assert!(
                    !self.locals[*src].repr.is_place(),
                    "function `{}` uses place source {:?} in numeric cast",
                    self.name,
                    src,
                );
                assert!(
                    !self.locals[*dst].repr.is_place(),
                    "function `{}` writes numeric cast into place {:?}",
                    self.name,
                    dst,
                );
            }
            MirInstr::IntLt { lhs, rhs, dst }
            | MirInstr::IntGt { lhs, rhs, dst }
            | MirInstr::IntLtEq { lhs, rhs, dst }
            | MirInstr::IntGtEq { lhs, rhs, dst }
            | MirInstr::IntEq { lhs, rhs, dst }
            | MirInstr::IntNeq { lhs, rhs, dst }
            | MirInstr::IntAdd { lhs, rhs, dst }
            | MirInstr::IntSub { lhs, rhs, dst }
            | MirInstr::IntMul { lhs, rhs, dst }
            | MirInstr::IntDiv { lhs, rhs, dst }
            | MirInstr::IntMod { lhs, rhs, dst } => {
                assert!(
                    !self.locals[*lhs].repr.is_place(),
                    "function `{}` uses place lhs {:?} in scalar op",
                    self.name,
                    lhs,
                );
                assert!(
                    !self.locals[*rhs].repr.is_place(),
                    "function `{}` uses place rhs {:?} in scalar op",
                    self.name,
                    rhs,
                );
                assert!(
                    !self.locals[*dst].repr.is_place(),
                    "function `{}` writes scalar op into place {:?}",
                    self.name,
                    dst,
                );
            }
            _ => {}
        }
    }

    /// Assert call argument invariants that are backend-independent.
    fn assert_call_args_valid(&self, args: &[MirCallArg]) {
        for arg in args {
            let binding = self.call_arg_binding(arg);
            if binding.is_mutable_borrow() {
                assert!(
                    !matches!(binding.local.repr, MirLocalRepr::SharedPlace),
                    "function `{}` mutably borrows shared place {:?}",
                    self.name,
                    binding.local_id(),
                );
            }
        }
    }
}

/// Basic block in MIR CFG.
///
/// `params` are edge-passed block parameters (phi-like interface).
#[derive(Debug, Clone)]
pub struct MirBlock {
    /// Local ids corresponding to incoming parameters.
    pub params: Vec<MirLocalId>,
    /// Instructions executed before the block terminator.
    pub instrs: Vec<MirInstr>,
    /// Terminator controlling the outgoing flow edge(s).
    pub terminator: MirTerminator,
}

/// Block terminator defining control-flow edges.
#[derive(Debug, Clone)]
pub enum MirTerminator {
    /// Unconditional jump to `target` with block argument values.
    Jump {
        /// Block id to jump to.
        target: MirBlockId,
        /// Block arguments passed to the target.
        args: Vec<MirLocalId>,
    },
    /// Conditional branch with explicit then/else targets and arguments.
    Branch {
        /// Local id holding the branch condition.
        cond: MirLocalId,
        /// Target block when the condition is true.
        then_target: MirBlockId,
        /// Block arguments passed to the then target.
        then_args: Vec<MirLocalId>,
        /// Target block when the condition is false.
        else_target: MirBlockId,
        /// Block arguments passed to the else target.
        else_args: Vec<MirLocalId>,
    },
    /// Return from current function with value local.
    /// Structured return used during CFG construction.
    Return {
        /// Local id returned from the function.
        value: MirLocalId,
    },
}