developing-compilers/src/util.rs

use std::collections::HashSet;
use std::io;

use crate::ast::UntypedExpr;
use crate::ast::typed::TypedExpr;
use crate::representation::intermediate::{IntermediateExpr, IntermediateValue};
use crate::symbols::{Symbol, SymbolsTable};

pub fn collect_used_symbols(expr: &[IntermediateExpr]) -> Vec<Symbol> {
    let mut used_symbols = HashSet::new();

    for expr in expr {
        match expr {
            IntermediateExpr::IntToFloat { result, value } => {
                used_symbols.insert(*result);

                if let IntermediateValue::Var { name } = value {
                    used_symbols.insert(*name);
                }
            },
            IntermediateExpr::BinOp {
                result, lhs, rhs, ..
            } => {
                used_symbols.insert(*result);

                if let IntermediateValue::Var { name } = lhs {
                    used_symbols.insert(*name);
                }

                if let IntermediateValue::Var { name } = rhs {
                    used_symbols.insert(*name);
                }
            },
            IntermediateExpr::Identity { result, value } => {
                used_symbols.insert(*result);

                if let IntermediateValue::Var { name } = value {
                    used_symbols.insert(*name);
                }
            },
        }
    }

    used_symbols.into_iter().collect()
}

pub fn print_intermediate_exprs(
    exprs: &[IntermediateExpr],
    writer: &mut impl io::Write,
) -> io::Result<()> {
    for expr in exprs {
        writeln!(writer, "{}", expr)?;
    }

    Ok(())
}

pub fn print_postfix_expr(expr: &TypedExpr, writer: &mut impl io::Write) -> io::Result<()> {
    match expr {
        TypedExpr::Int { value, .. } => write!(writer, "{} ", value)?,
        TypedExpr::Float { value, .. } => write!(writer, "{} ", value)?,
        TypedExpr::Var { name, .. } => write!(writer, "<id,{}> ", name)?,
        TypedExpr::BinOp { lhs, op, rhs, .. } => {
            print_postfix_expr(lhs, writer)?;
            print_postfix_expr(rhs, writer)?;
            write!(writer, "{} ", op)?;
        },
        TypedExpr::IntToFloat { value, .. } => {
            print_postfix_expr(value, writer)?;
            write!(writer, "i2f ")?;
        },
    }

    Ok(())
}

fn write_typed_expr(
    expr: &TypedExpr,
    symbols: &SymbolsTable,
    writer: &mut impl io::Write,
    prefix: &str,
    is_last: bool,
) -> io::Result<()> {
    let branch = if is_last { "└──" } else { "├──" };

    write!(writer, "{}{}", prefix, branch)?;

    match expr {
        TypedExpr::Int { value, .. } => writeln!(writer, "<{}>", value),
        TypedExpr::Float { value, .. } => {
            if value == &value.trunc() {
                writeln!(writer, "<{}.0>", value)
            } else {
                writeln!(writer, "<{}>", value)
            }
        },
        TypedExpr::Var { name, .. } => {
            let ty = symbols.resolve(name).unwrap().ty.unwrap();
            writeln!(writer, "<id,{},{}>", name, ty)
        },
        TypedExpr::BinOp { lhs, op, rhs, .. } => {
            writeln!(writer, "<{}>", op)?;

            let new_prefix = if is_last {
                format!("{}    ", prefix)
            } else {
                format!("{}│   ", prefix)
            };

            write_typed_expr(lhs, symbols, writer, &new_prefix, false)?;
            write_typed_expr(rhs, symbols, writer, &new_prefix, true)
        },
        TypedExpr::IntToFloat { value, .. } => {
            writeln!(writer, "i2f")?;

            let new_prefix = if is_last {
                format!("{}    ", prefix)
            } else {
                format!("{}│   ", prefix)
            };

            write_typed_expr(value, symbols, writer, &new_prefix, true)
        },
    }
}

pub fn print_typed_expr(
    expr: &TypedExpr,
    symbols: &SymbolsTable,
    writer: &mut impl io::Write,
) -> io::Result<()> {
    write_typed_expr(expr, symbols, writer, "", true)
}

fn write_untyped_expr(
    expr: &UntypedExpr,
    writer: &mut impl io::Write,
    prefix: &str,
    is_last: bool,
) -> io::Result<()> {
    let branch = if is_last { "└──" } else { "├──" };

    write!(writer, "{}{}", prefix, branch)?;

    match expr {
        UntypedExpr::Int { value, .. } => writeln!(writer, "<{}>", value),
        UntypedExpr::Float { value, .. } => {
            if value == &value.trunc() {
                writeln!(writer, "<{}.0>", value)
            } else {
                writeln!(writer, "<{}>", value)
            }
        },
        UntypedExpr::Var {
            name: id, typename, ..
        } => {
            write!(writer, "<id,{}", id)?;
            if let Some(typename) = typename {
                write!(writer, ",{}", typename)?;
            }
            writeln!(writer, ">")
        },
        UntypedExpr::BinOp { lhs, op, rhs, .. } => {
            writeln!(writer, "<{}>", op)?;

            let new_prefix = if is_last {
                format!("{}    ", prefix)
            } else {
                format!("{}│   ", prefix)
            };

            write_untyped_expr(lhs, writer, &new_prefix, false)?;
            write_untyped_expr(rhs, writer, &new_prefix, true)
        },
    }
}

pub fn print_untyped_expr(expr: &UntypedExpr, writer: &mut impl io::Write) -> io::Result<()> {
    write_untyped_expr(expr, writer, "", true)
}