slonik/src/ir/write.rs

//! Text formatting for Slonik IR.

use std::fmt;

use crate::ir::{Block, BlockCall, Body, Stmt, StmtData, Value};

/// Writes `body` in textual form.
pub fn write_body(f: &mut fmt::Formatter<'_>, body: &Body) -> fmt::Result {
    Printer { body }.write_body(f)
}

impl fmt::Display for Body {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write_body(f, self)
    }
}

struct Printer<'a> {
    body: &'a Body,
}

impl Printer<'_> {
    fn write_body(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(f, "body {{")?;

        let mut first = true;
        for block in self.body.blocks() {
            if !first {
                writeln!(f)?;
            }
            first = false;
            self.write_block(f, block)?;
        }

        writeln!(f, "}}")
    }

    fn write_block(&self, f: &mut fmt::Formatter<'_>, block: Block) -> fmt::Result {
        write!(f, "  ^{block}")?;

        let params = self.body.block_params(block);
        if !params.is_empty() {
            write!(f, "(")?;
            for (i, &param) in params.iter().enumerate() {
                if i != 0 {
                    write!(f, ", ")?;
                }
                write!(f, "%{}: {}", param, self.body.value_type(param))?;
            }
            write!(f, ")")?;
        }

        writeln!(f, ":")?;

        for &stmt in self.body.block_stmts(block) {
            self.write_stmt(f, stmt)?;
        }

        Ok(())
    }

    fn write_stmt(&self, f: &mut fmt::Formatter<'_>, stmt: Stmt) -> fmt::Result {
        write!(f, "    ")?;

        match *self.body.stmt_data(stmt) {
            StmtData::IConst { imm, result } => {
                writeln!(f, "%{result} = iconst {imm} : {}", self.body.value_type(result))
            }

            StmtData::F32Const { bits, result } => {
                writeln!(f, "%{result} = f32const 0x{bits:08x} : f32")
            }

            StmtData::F64Const { bits, result } => {
                writeln!(f, "%{result} = f64const 0x{bits:016x} : f64")
            }

            StmtData::BConst { value, result } => {
                writeln!(f, "%{result} = bconst {value} : bool")
            }

            StmtData::Unary { op, arg, result } => {
                writeln!(
                    f,
                    "%{result} = {op} %{arg} : {}",
                    self.body.value_type(result)
                )
            }

            StmtData::Binary { op, lhs, rhs, result } => {
                writeln!(
                    f,
                    "%{result} = {op} %{lhs}, %{rhs} : {}",
                    self.body.value_type(result)
                )
            }

            StmtData::Cast { op, arg, result } => {
                writeln!(
                    f,
                    "%{result} = {op} %{arg} : {} -> {}",
                    self.body.value_type(arg),
                    self.body.value_type(result)
                )
            }

            StmtData::Icmp { cc, lhs, rhs, result } => {
                writeln!(
                    f,
                    "%{result} = icmp {cc} %{lhs}, %{rhs} : {}",
                    self.body.value_type(lhs)
                )
            }

            StmtData::Fcmp { cc, lhs, rhs, result } => {
                writeln!(
                    f,
                    "%{result} = fcmp {cc} %{lhs}, %{rhs} : {}",
                    self.body.value_type(lhs)
                )
            }

            StmtData::Select {
                cond,
                if_true,
                if_false,
                result,
            } => {
                writeln!(
                    f,
                    "%{result} = select %{cond}, %{if_true}, %{if_false} : {}",
                    self.body.value_type(result)
                )
            }

            StmtData::Load { addr, size, result } => {
                writeln!(
                    f,
                    "%{result} = load %{addr} : {size} -> {}",
                    self.body.value_type(result)
                )
            }

            StmtData::Store { addr, value, size } => {
                writeln!(f, "store %{value}, %{addr} : {size}")
            }

            StmtData::Call { callee, args } => {
                write!(f, "call %{callee}")?;
                let args = args.as_slice(&self.body.value_lists);
                if !args.is_empty() {
                    write!(f, "(")?;
                    self.write_value_list(f, args)?;
                    write!(f, ")")?;
                }
                writeln!(f)
            }

            StmtData::Jump { dst } => {
                write!(f, "br ")?;
                self.write_block_call(f, dst)?;
                writeln!(f)
            }

            StmtData::BrIf {
                cond,
                then_dst,
                else_dst,
            } => {
                write!(f, "cond_br %{cond}, ")?;
                self.write_block_call(f, then_dst)?;
                write!(f, ", ")?;
                self.write_block_call(f, else_dst)?;
                writeln!(f)
            }

            StmtData::Return { values } => {
                write!(f, "return")?;
                let values = values.as_slice(&self.body.value_lists);
                if !values.is_empty() {
                    write!(f, " ")?;
                    self.write_value_list(f, values)?;
                }
                writeln!(f)
            }
        }
    }

    fn write_block_call(&self, f: &mut fmt::Formatter<'_>, call: BlockCall) -> fmt::Result {
        write!(f, "^{}", call.block)?;
        let args = call.args.as_slice(&self.body.value_lists);
        if !args.is_empty() {
            write!(f, "(")?;
            self.write_value_list(f, args)?;
            write!(f, ")")?;
        }
        Ok(())
    }

    fn write_value_list(&self, f: &mut fmt::Formatter<'_>, values: &[Value]) -> fmt::Result {
        for (i, value) in values.iter().enumerate() {
            if i != 0 {
                write!(f, ", ")?;
            }
            write!(f, "%{value}")?;
        }
        Ok(())
    }
}