joule_mir/

pretty.rs

//! Pretty printing for MIR
//!
//! This module provides human-readable output for MIR structures,
//! useful for debugging and inspection.

use crate::{
    BasicBlock, BasicBlockId, FunctionMIR, MirContext, Operand, Place, Rvalue, Statement,
    Terminator,
};
use std::fmt::Write;

/// Pretty print a MIR context
pub fn print_mir_context(ctx: &MirContext) -> String {
    let mut output = String::new();

    for (_, func) in &ctx.functions {
        let _ = write!(output, "{}\n", print_function(func));
    }

    output
}

/// Pretty print a function
pub fn print_function(func: &FunctionMIR) -> String {
    let mut output = String::new();

    // Function header
    let _ = writeln!(output, "fn {} -> {:?} {{", func.name, func.return_ty);

    // Locals
    if !func.locals.is_empty() {
        let _ = writeln!(output, "  // Locals:");
        for (idx, local) in func.locals.iter().enumerate() {
            let name = if idx == 0 {
                "_ret".to_string()
            } else {
                format!("_{}", idx)
            };
            let mut_str = if local.mutable { "mut " } else { "" };
            let _ = writeln!(output, "  //   {}: {}{:?}", name, mut_str, local.ty);
        }
        let _ = writeln!(output);
    }

    // Basic blocks
    for (idx, block) in func.basic_blocks.iter().enumerate() {
        let _ = write!(
            output,
            "{}",
            print_basic_block(BasicBlockId::new(idx as u32), block)
        );
    }

    let _ = writeln!(output, "}}");

    output
}

/// Pretty print a basic block
pub fn print_basic_block(id: BasicBlockId, block: &BasicBlock) -> String {
    let mut output = String::new();

    let _ = writeln!(output, "  {}:", id);

    for stmt in &block.statements {
        let _ = writeln!(output, "    {}", print_statement(stmt));
    }

    let _ = writeln!(output, "    {}", print_terminator(&block.terminator));
    let _ = writeln!(output);

    output
}

/// Pretty print a statement
pub fn print_statement(stmt: &Statement) -> String {
    match stmt {
        Statement::Assign {
            place,
            rvalue,
            span: _,
        } => {
            format!("{} = {}", print_place(place), print_rvalue(rvalue))
        }
        Statement::StorageLive { local, span: _ } => {
            format!("StorageLive({})", local)
        }
        Statement::StorageDead { local, span: _ } => {
            format!("StorageDead({})", local)
        }
        Statement::Nop => "nop".to_string(),
    }
}

/// Pretty print a terminator
pub fn print_terminator(term: &Terminator) -> String {
    match term {
        Terminator::Return { .. } => "return".to_string(),
        Terminator::Goto { target, .. } => {
            format!("goto {}", target)
        }
        Terminator::SwitchInt {
            discriminant,
            targets,
            ..
        } => {
            let mut s = format!("switchInt({}) -> [", print_operand(discriminant));
            for (value, target) in &targets.branches {
                s.push_str(&format!("{}: {}, ", value, target));
            }
            s.push_str(&format!("otherwise: {}]", targets.otherwise));
            s
        }
        Terminator::Call {
            func,
            args,
            destination,
            target,
            ..
        } => {
            let args_str = args
                .iter()
                .map(print_operand)
                .collect::<Vec<_>>()
                .join(", ");
            format!(
                "{} = {}({}) -> {}",
                print_place(destination),
                print_operand(func),
                args_str,
                target
            )
        }
        Terminator::Abort { .. } => "abort".to_string(),
        Terminator::Unreachable { .. } => "unreachable".to_string(),

        // Concurrency terminators
        Terminator::Spawn {
            func,
            args,
            destination,
            target,
            ..
        } => {
            let args_str = args
                .iter()
                .map(print_operand)
                .collect::<Vec<_>>()
                .join(", ");
            format!(
                "{} = spawn {}({}) -> {}",
                print_place(destination),
                print_operand(func),
                args_str,
                target
            )
        }
        Terminator::TaskAwait {
            task,
            destination,
            target,
            ..
        } => {
            format!(
                "{} = await {} -> {}",
                print_place(destination),
                print_operand(task),
                target
            )
        }
        Terminator::TaskGroupEnter {
            destination,
            body,
            join_block,
            ..
        } => {
            format!(
                "{} = task_group_enter -> [body: {}, join: {}]",
                print_place(destination),
                body,
                join_block
            )
        }
        Terminator::TaskGroupExit { group, target, .. } => {
            format!("task_group_exit {} -> {}", print_operand(group), target)
        }
        Terminator::ChannelRecv {
            channel,
            destination,
            target,
            closed_target,
            ..
        } => {
            format!(
                "{} = recv {} -> [ok: {}, closed: {}]",
                print_place(destination),
                print_operand(channel),
                target,
                closed_target
            )
        }
        Terminator::ChannelSend {
            channel,
            value,
            target,
            closed_target,
            ..
        } => {
            format!(
                "send {} <- {} -> [ok: {}, closed: {}]",
                print_operand(channel),
                print_operand(value),
                target,
                closed_target
            )
        }
        Terminator::Select {
            arms,
            default,
            destination,
            selected_arm,
            ..
        } => {
            let arms_str = arms
                .iter()
                .enumerate()
                .map(|(i, arm)| format!("{}: {:?} -> {}", i, arm.operation, arm.target))
                .collect::<Vec<_>>()
                .join(", ");
            let default_str = default
                .map(|b| format!(", default: {}", b))
                .unwrap_or_default();
            format!(
                "({}, {}) = select [{}{}]",
                print_place(destination),
                print_place(selected_arm),
                arms_str,
                default_str
            )
        }
        Terminator::Cancel { .. } => "cancel".to_string(),
    }
}

/// Pretty print a place
pub fn print_place(place: &Place) -> String {
    let mut s = format!("{}", place.local);
    for elem in &place.projection {
        match elem {
            crate::PlaceElem::Deref => s.push_str(".*"),
            crate::PlaceElem::Field(field) => {
                s.push_str(&format!(".{}", field.0));
            }
            crate::PlaceElem::Index(local) => {
                s.push_str(&format!("[{}]", local));
            }
            crate::PlaceElem::Downcast(name) => {
                s.push_str(&format!(" as {}", name.as_str()));
            }
        }
    }
    s
}

/// Pretty print an rvalue
pub fn print_rvalue(rvalue: &Rvalue) -> String {
    match rvalue {
        Rvalue::Use(op) => print_operand(op),
        Rvalue::BinaryOp { op, left, right } => {
            format!(
                "{:?}({}, {})",
                op,
                print_operand(left),
                print_operand(right)
            )
        }
        Rvalue::UnaryOp { op, operand } => {
            format!("{:?}({})", op, print_operand(operand))
        }
        Rvalue::Ref { mutable, place } => {
            let mut_str = if *mutable { "mut " } else { "" };
            format!("&{}{}", mut_str, print_place(place))
        }
        Rvalue::Aggregate { kind, operands } => {
            let ops = operands
                .iter()
                .map(print_operand)
                .collect::<Vec<_>>()
                .join(", ");
            format!("{:?}({})", kind, ops)
        }
        Rvalue::Cast {
            operand,
            target_ty,
            kind,
        } => {
            format!("{} as {:?} ({:?})", print_operand(operand), target_ty, kind)
        }
        Rvalue::Discriminant { place } => {
            format!("discriminant({})", print_place(place))
        }
        Rvalue::Len { place } => {
            format!("len({})", print_place(place))
        }
        Rvalue::SimdBinaryOp {
            op,
            element_ty,
            width,
            left,
            right,
        } => {
            format!(
                "simd_{:?}_{:?}x{:?}({}, {})",
                op,
                element_ty,
                width,
                print_operand(left),
                print_operand(right)
            )
        }
        Rvalue::SimdLoad {
            element_ty,
            width,
            source,
        } => {
            format!(
                "simd_load_{:?}x{:?}({})",
                element_ty,
                width,
                print_place(source)
            )
        }
        Rvalue::SimdStore {
            element_ty,
            width,
            value,
            dest,
        } => {
            format!(
                "simd_store_{:?}x{:?}({}, {})",
                element_ty,
                width,
                print_operand(value),
                print_place(dest)
            )
        }
        Rvalue::SimdSplat {
            element_ty,
            width,
            value,
        } => {
            format!(
                "simd_splat_{:?}x{:?}({})",
                element_ty,
                width,
                print_operand(value)
            )
        }

        // Concurrency rvalues
        Rvalue::ChannelCreate {
            element_ty,
            capacity,
        } => {
            format!("channel::<{:?}>({})", element_ty, capacity)
        }
        Rvalue::ChannelTryRecv { channel } => {
            format!("try_recv({})", print_operand(channel))
        }
        Rvalue::ChannelTrySend { channel, value } => {
            format!(
                "try_send({}, {})",
                print_operand(channel),
                print_operand(value)
            )
        }
        Rvalue::ChannelSender { channel } => {
            format!("sender({})", print_operand(channel))
        }
        Rvalue::ChannelReceiver { channel } => {
            format!("receiver({})", print_operand(channel))
        }
        Rvalue::ChannelClose { channel } => {
            format!("close({})", print_operand(channel))
        }
        Rvalue::IsCancelled => "is_cancelled()".to_string(),
        Rvalue::CurrentTask => "current_task()".to_string(),
        Rvalue::Try(operand) => format!("try({})", print_operand(operand)),
        Rvalue::EnumVariant { variant, fields } => {
            if fields.is_empty() {
                format!("enum_variant({})", variant)
            } else {
                let field_strs: Vec<String> = fields.iter().map(print_operand).collect();
                format!("enum_variant({}, {})", variant, field_strs.join(", "))
            }
        }
    }
}

/// Pretty print an operand
pub fn print_operand(operand: &Operand) -> String {
    match operand {
        Operand::Copy(place) => format!("copy {}", print_place(place)),
        Operand::Move(place) => format!("move {}", print_place(place)),
        Operand::Constant(c) => format!("const {:?}", c.literal),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::Local;
    use joule_common::Span;

    #[test]
    fn test_print_place() {
        let place = Place::from_local(Local::new(1));
        assert_eq!(print_place(&place), "_1");

        let deref_place = place.deref();
        assert_eq!(print_place(&deref_place), "_1.*");
    }

    #[test]
    fn test_print_statement() {
        let stmt = Statement::StorageLive {
            local: Local::new(1),
            span: Span::dummy(),
        };
        assert_eq!(print_statement(&stmt), "StorageLive(_1)");

        let stmt = Statement::Nop;
        assert_eq!(print_statement(&stmt), "nop");
    }

    #[test]
    fn test_print_terminator() {
        let term = Terminator::Return {
            span: Span::dummy(),
        };
        assert_eq!(print_terminator(&term), "return");

        let term = Terminator::Goto {
            target: BasicBlockId::new(1),
            span: Span::dummy(),
        };
        assert_eq!(print_terminator(&term), "goto bb1");
    }

    #[test]
    fn test_print_operand() {
        let op = Operand::Copy(Place::from_local(Local::new(1)));
        assert_eq!(print_operand(&op), "copy _1");

        let op = Operand::Move(Place::from_local(Local::new(2)));
        assert_eq!(print_operand(&op), "move _2");
    }
}