lamina-mir 0.2.0

//! Basic blocks for LUMIR (Lamina Unified Machine Intermediate Representation).
//!
//! A basic block is a sequence of instructions with a single entry point and
//! a single exit point (terminator instruction). Basic blocks are the fundamental
//! unit of control flow in LUMIR.
use crate::instruction::Instruction;
use crate::loc::{Located, SourceLoc};
use std::fmt;

/// Block in LUMIR
#[derive(Debug, Clone)]
pub struct Block {
    /// Block label/name
    pub label: String,

    /// Instructions in this block
    pub instructions: Vec<Located>,
}

impl fmt::Display for Block {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(f, "{}:", self.label)?;
        for instr in &self.instructions {
            writeln!(f, "  {instr}")?;
        }
        Ok(())
    }
}

impl PartialEq for Block {
    fn eq(&self, other: &Self) -> bool {
        self.label == other.label
    }
}

impl Block {
    /// Create a new basic block with the given label
    pub fn new(label: impl Into<String>) -> Self {
        Self {
            label: label.into(),
            instructions: Vec::new(),
        }
    }

    /// Add an instruction to this block
    pub fn push(&mut self, instr: impl Into<Located>) {
        self.instructions.push(instr.into());
    }

    /// Append an instruction that came from `loc` in the source.
    pub fn push_at(&mut self, instr: Instruction, loc: Option<SourceLoc>) {
        self.instructions.push(Located::at(instr, loc));
    }

    /// Get the terminator instruction (last instruction)
    pub fn terminator(&self) -> Option<&Instruction> {
        self.instructions.last().map(|i| &i.instr)
    }

    /// Check if this block has a terminator
    pub fn has_terminator(&self) -> bool {
        self.terminator().is_some_and(Instruction::is_terminator)
    }

    /// Get all non-terminator instructions
    pub fn body(&self) -> &[Located] {
        if self.has_terminator() {
            &self.instructions[..self.instructions.len() - 1]
        } else {
            &self.instructions
        }
    }

    /// Number of instructions in this block
    pub fn len(&self) -> usize {
        self.instructions.len()
    }

    /// Check if this block is empty
    pub fn is_empty(&self) -> bool {
        self.instructions.is_empty()
    }

    /// Get successor block labels, borrowing from the terminator instruction.
    pub fn successors(&self) -> Vec<&str> {
        match self.terminator() {
            Some(Instruction::Jmp { target }) => vec![target.as_str()],
            Some(Instruction::Br {
                true_target,
                false_target,
                ..
            }) => vec![true_target.as_str(), false_target.as_str()],
            Some(Instruction::Switch { cases, default, .. }) => {
                let mut targets = vec![default.as_str()];
                targets.extend(cases.iter().map(|(_, t)| t.as_str()));
                targets
            }
            _ => vec![],
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        instruction::{Immediate, Instruction, IntBinOp, Operand},
        register::{Register, VirtualReg},
        types::{MirType, ScalarType},
    };

    #[test]
    fn test_basic_block_push() {
        let mut bb = Block::new("bb0");

        let add = Instruction::IntBinary {
            op: IntBinOp::Add,
            ty: MirType::Scalar(ScalarType::I64),
            dst: Register::Virtual(VirtualReg::gpr(0)),
            lhs: Operand::Immediate(Immediate::I64(1)),
            rhs: Operand::Immediate(Immediate::I64(2)),
        };

        bb.push(add);
        assert_eq!(bb.len(), 1);
    }

    #[test]
    fn test_basic_block_terminator() {
        let mut bb = Block::new("bb0");

        let add = Instruction::IntBinary {
            op: IntBinOp::Add,
            ty: MirType::Scalar(ScalarType::I64),
            dst: Register::Virtual(VirtualReg::gpr(0)),
            lhs: Operand::Immediate(Immediate::I64(1)),
            rhs: Operand::Immediate(Immediate::I64(2)),
        };

        let ret = Instruction::Ret { value: None };

        bb.push(add.clone());
        assert!(!bb.has_terminator());

        bb.push(ret);
        assert!(bb.has_terminator());
        assert_eq!(bb.body().len(), 1);
    }
}