lift_core/

printer.rs

use crate::attributes::Attribute;
use crate::blocks::BlockKey;
use crate::context::Context;
use crate::operations::OpKey;
use crate::regions::RegionKey;
use crate::values::ValueKey;
use std::fmt::Write;

pub struct Printer<'a> {
    ctx: &'a Context,
    indent: usize,
    output: String,
    value_names: std::collections::HashMap<ValueKey, String>,
    next_value_id: u32,
    next_block_id: u32,
    block_names: std::collections::HashMap<BlockKey, String>,
}

impl<'a> Printer<'a> {
    pub fn new(ctx: &'a Context) -> Self {
        Self {
            ctx,
            indent: 0,
            output: String::new(),
            value_names: std::collections::HashMap::new(),
            next_value_id: 0,
            next_block_id: 0,
            block_names: std::collections::HashMap::new(),
        }
    }

    pub fn print_all(&mut self) -> &str {
        for module in &self.ctx.modules {
            let name = self.ctx.strings.resolve(module.name);
            self.write_line(&format!("module @{} {{", name));
            self.indent += 1;

            for func in &module.functions {
                self.print_function(func);
            }

            self.indent -= 1;
            self.write_line("}");
            self.write_line("");
        }

        // Print any standalone blocks/ops not in modules
        for (block_key, _block) in &self.ctx.blocks {
            if self.ctx.blocks[block_key].parent_region.is_none() {
                self.print_block(block_key, true);
            }
        }

        &self.output
    }

    fn print_function(&mut self, func: &crate::functions::FunctionData) {
        let name = self.ctx.strings.resolve(func.name);
        let mut sig = format!("func @{}(", name);

        // The entry block's arguments are the actual values the body refers
        // to (see lift-ast's builder, which creates them directly from the
        // parameter list). Print those real names instead of disconnected
        // fresh ones, so the signature and body agree and the printed .lif
        // can be parsed back in.
        let mut entry_args: Option<Vec<ValueKey>> = None;
        if let Some(region) = func.body {
            if let Some(r) = self.ctx.get_region(region) {
                if let Some(&block_key) = r.blocks.first() {
                    if let Some(block) = self.ctx.get_block(block_key) {
                        entry_args = Some(block.args.clone());
                    }
                }
            }
        }

        if let Some(args) = entry_args {
            for (i, arg) in args.into_iter().enumerate() {
                if i > 0 {
                    sig.push_str(", ");
                }
                let vname = self.get_value_name(arg);
                if let Some(val) = self.ctx.get_value(arg) {
                    let _ = write!(sig, "%{}: {}", vname, self.format_type(val.ty));
                }
            }
        } else {
            for (i, &param_ty) in func.params.iter().enumerate() {
                if i > 0 {
                    sig.push_str(", ");
                }
                let pname = self.fresh_value_name();
                let _ = write!(sig, "%{}: {}", pname, self.format_type(param_ty));
            }
        }

        sig.push_str(") -> ");
        if func.returns.len() == 1 {
            let _ = write!(sig, "{}", self.format_type(func.returns[0]));
        } else {
            sig.push('(');
            for (i, &ret_ty) in func.returns.iter().enumerate() {
                if i > 0 {
                    sig.push_str(", ");
                }
                let _ = write!(sig, "{}", self.format_type(ret_ty));
            }
            sig.push(')');
        }

        if let Some(body) = func.body {
            sig.push_str(" {");
            self.write_line(&sig);
            self.indent += 1;
            self.print_function_body(body);
            self.indent -= 1;
            self.write_line("}");
        } else {
            self.write_line(&sig);
        }
        self.write_line("");
    }

    fn print_region(&mut self, region_key: RegionKey) {
        if let Some(region) = self.ctx.get_region(region_key) {
            for &block_key in &region.blocks {
                self.print_block(block_key, true);
            }
        }
    }

    /// Prints a function's body region. The entry block's args were already
    /// declared in the function signature (`print_function`), so its
    /// `^bb0(...):` header would be redundant — and unparseable, since the
    /// current `.lif` grammar has no rule for block labels (every function
    /// is single-block). Any further blocks (not producible by the parser
    /// today, but structurally possible) still get a normal header.
    fn print_function_body(&mut self, region_key: RegionKey) {
        if let Some(region) = self.ctx.get_region(region_key) {
            for (i, &block_key) in region.blocks.iter().enumerate() {
                self.print_block(block_key, i != 0);
            }
        }
    }

    fn print_block(&mut self, block_key: BlockKey, print_header: bool) {
        if print_header {
            let block_name = self.get_block_name(block_key);
            if let Some(block) = self.ctx.get_block(block_key) {
                if !block.args.is_empty() {
                    let mut args = String::new();
                    for (i, &arg) in block.args.iter().enumerate() {
                        if i > 0 {
                            args.push_str(", ");
                        }
                        let vname = self.get_value_name(arg);
                        if let Some(val) = self.ctx.get_value(arg) {
                            let _ = write!(args, "%{}: {}", vname, self.format_type(val.ty));
                        }
                    }
                    self.write_line(&format!("^{}({}):", block_name, args));
                } else {
                    self.write_line(&format!("^{}:", block_name));
                }
            }
            self.indent += 1;
        }

        if let Some(block) = self.ctx.get_block(block_key) {
            for &op_key in &block.ops {
                self.print_op(op_key);
            }
        }

        if print_header {
            self.indent -= 1;
        }
    }

    fn print_op(&mut self, op_key: OpKey) {
        if let Some(op) = self.ctx.get_op(op_key) {
            let op_name = self.ctx.strings.resolve(op.name).to_string();
            let mut line = String::new();

            // Result values
            if !op.results.is_empty() {
                for (i, &result) in op.results.iter().enumerate() {
                    if i > 0 {
                        line.push_str(", ");
                    }
                    let vname = self.get_value_name(result);
                    let _ = write!(line, "%{}", vname);
                }
                line.push_str(" = ");
            }

            // Operation name
            let _ = write!(line, "\"{}\"", op_name);

            // Inputs
            line.push('(');
            for (i, &input) in op.inputs.iter().enumerate() {
                if i > 0 {
                    line.push_str(", ");
                }
                let vname = self.get_value_name(input);
                let _ = write!(line, "%{}", vname);
            }
            line.push(')');

            // Attributes
            if !op.attrs.is_empty() {
                line.push_str(" {");
                for (i, (key, val)) in op.attrs.iter().enumerate() {
                    if i > 0 {
                        line.push_str(", ");
                    }
                    let _ = write!(line, "{} = {}", key, self.format_attr(val));
                }
                line.push('}');
            }

            // Type signature
            line.push_str(" : (");
            for (i, &input) in op.inputs.iter().enumerate() {
                if i > 0 {
                    line.push_str(", ");
                }
                if let Some(val) = self.ctx.get_value(input) {
                    let _ = write!(line, "{}", self.format_type(val.ty));
                }
            }
            line.push_str(") -> ");

            if op.results.len() == 1 {
                if let Some(val) = self.ctx.get_value(op.results[0]) {
                    let _ = write!(line, "{}", self.format_type(val.ty));
                }
            } else if op.results.is_empty() {
                line.push_str("()");
            } else {
                line.push('(');
                for (i, &result) in op.results.iter().enumerate() {
                    if i > 0 {
                        line.push_str(", ");
                    }
                    if let Some(val) = self.ctx.get_value(result) {
                        let _ = write!(line, "{}", self.format_type(val.ty));
                    }
                }
                line.push(')');
            }

            self.write_line(&line);

            // Print nested regions
            if !op.regions.is_empty() {
                for &region in &op.regions {
                    self.indent += 1;
                    self.print_region(region);
                    self.indent -= 1;
                }
            }
        }
    }

    fn format_type(&self, ty_id: crate::types::TypeId) -> String {
        let ty = self.ctx.resolve_type(ty_id);
        format!("{}", ty)
    }

    fn format_attr(&self, attr: &Attribute) -> String {
        match attr {
            Attribute::Integer(v) => format!("{}", v),
            Attribute::Float(v) => format!("{:.6}", v),
            Attribute::String(s) => {
                let resolved = self.ctx.strings.resolve(*s);
                format!("\"{}\"", resolved)
            }
            Attribute::Bool(b) => format!("{}", b),
            Attribute::Type(_) => "type".to_string(),
            Attribute::Array(arr) => {
                let inner: Vec<String> = arr.iter().map(|a| self.format_attr(a)).collect();
                format!("[{}]", inner.join(", "))
            }
            Attribute::Dict(map) => {
                let inner: Vec<String> = map
                    .iter()
                    .map(|(k, v)| format!("{}: {}", k, self.format_attr(v)))
                    .collect();
                format!("{{{}}}", inner.join(", "))
            }
        }
    }

    fn get_value_name(&mut self, key: ValueKey) -> String {
        if let Some(name) = self.value_names.get(&key) {
            return name.clone();
        }

        // Try to use the debug name if available
        let name = if let Some(val) = self.ctx.get_value(key) {
            if let Some(name_id) = val.name {
                self.ctx.strings.resolve(name_id).to_string()
            } else {
                self.fresh_value_name()
            }
        } else {
            self.fresh_value_name()
        };

        self.value_names.insert(key, name.clone());
        name
    }

    fn fresh_value_name(&mut self) -> String {
        let name = format!("v{}", self.next_value_id);
        self.next_value_id += 1;
        name
    }

    fn get_block_name(&mut self, key: BlockKey) -> String {
        if let Some(name) = self.block_names.get(&key) {
            return name.clone();
        }
        let name = format!("bb{}", self.next_block_id);
        self.next_block_id += 1;
        self.block_names.insert(key, name.clone());
        name
    }

    fn write_line(&mut self, text: &str) {
        for _ in 0..self.indent {
            self.output.push_str("    ");
        }
        self.output.push_str(text);
        self.output.push('\n');
    }

    pub fn into_string(self) -> String {
        self.output
    }
}

pub fn print_ir(ctx: &Context) -> String {
    let mut printer = Printer::new(ctx);
    printer.print_all();
    printer.into_string()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_print_empty_context() {
        let ctx = Context::new();
        let output = print_ir(&ctx);
        assert!(output.is_empty() || output.trim().is_empty());
    }
}