developing-compilers/src/representation/intermediate.rs

use std::{fmt, io};

use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};

use crate::ast::BinOp;
use crate::ast::typed::TypedExpr;
use crate::representation::util::TempVarContext;
use crate::symbols::{Symbol, SymbolsTable};

#[derive(Debug)]
pub enum IntermediateValue {
    Int { value: i64 },
    Float { value: f64 },
    Var { name: Symbol },
}

impl IntermediateValue {
    pub fn as_u8(&self) -> u8 {
        match self {
            IntermediateValue::Int { .. } => b'i',
            IntermediateValue::Float { .. } => b'f',
            IntermediateValue::Var { .. } => b'v',
        }
    }

    pub fn to_writer(&self, writer: &mut impl io::Write) -> io::Result<()> {
        writer.write_u8(self.as_u8())?;

        match self {
            IntermediateValue::Int { value } => writer.write_i64::<BigEndian>(*value)?,
            IntermediateValue::Float { value } => writer.write_f64::<BigEndian>(*value)?,
            IntermediateValue::Var { name } => writer.write_u64::<BigEndian>(name.into())?,
        }

        Ok(())
    }

    pub fn from_reader(reader: &mut impl io::Read) -> io::Result<Self> {
        let b = reader.read_u8()?;
        match b {
            b'i' => Ok(IntermediateValue::Int {
                value: reader.read_i64::<BigEndian>()?,
            }),
            b'f' => Ok(IntermediateValue::Float {
                value: reader.read_f64::<BigEndian>()?,
            }),
            b'v' => Ok(IntermediateValue::Var {
                name: reader.read_u64::<BigEndian>()?.into(),
            }),
            b => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("invalid byte: {b}"),
            )),
        }
    }
}

impl fmt::Display for IntermediateValue {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            IntermediateValue::Int { value } => write!(f, "{}", value),
            IntermediateValue::Float { value } => {
                if *value == value.trunc() {
                    write!(f, "{value}.0")
                } else {
                    write!(f, "{value}")
                }
            },
            IntermediateValue::Var { name } => write!(f, "<id,{}>", name),
        }
    }
}

#[derive(Debug)]
pub enum IntermediateExpr {
    IntToFloat {
        result: Symbol,
        value: IntermediateValue,
    },
    BinOp {
        result: Symbol,
        lhs: IntermediateValue,
        op: BinOp,
        rhs: IntermediateValue,
    },
    Identity {
        result: Symbol,
        value: IntermediateValue,
    },
}

impl fmt::Display for IntermediateExpr {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            IntermediateExpr::IntToFloat { result, value } => {
                write!(f, "i2f <id,{}>, {}", result, value)
            },
            IntermediateExpr::BinOp {
                result,
                lhs,
                op,
                rhs,
            } => {
                write!(f, "{:#} <id,{}>, {}, {}", op, result, lhs, rhs)
            },
            IntermediateExpr::Identity { result, value } => {
                write!(f, "idt <id,{}>, {}", result, value)
            },
        }
    }
}

impl IntermediateExpr {
    pub fn as_u8(&self) -> u8 {
        match self {
            IntermediateExpr::IntToFloat { .. } => 0x01,
            IntermediateExpr::BinOp { .. } => 0x02,
            IntermediateExpr::Identity { .. } => 0x03,
        }
    }

    pub fn to_writer(&self, writer: &mut impl io::Write) -> io::Result<()> {
        writer.write_u8(self.as_u8())?;

        match self {
            IntermediateExpr::IntToFloat { result, value } => {
                writer.write_u64::<BigEndian>(result.into())?;
                value.to_writer(writer)?;
            },
            IntermediateExpr::BinOp {
                result,
                lhs,
                op,
                rhs,
            } => {
                writer.write_u64::<BigEndian>(result.into())?;
                lhs.to_writer(writer)?;
                writer.write_u8(op.into())?;
                rhs.to_writer(writer)?;
            },
            IntermediateExpr::Identity { result, value } => {
                writer.write_u64::<BigEndian>(result.into())?;
                value.to_writer(writer)?;
            },
        }

        Ok(())
    }

    pub fn from_reader(reader: &mut impl io::Read) -> io::Result<Self> {
        let b = reader.read_u8()?;
        match b {
            0x01 => Ok(IntermediateExpr::IntToFloat {
                result: reader.read_u64::<BigEndian>()?.into(),
                value: IntermediateValue::from_reader(reader)?,
            }),
            0x02 => Ok(IntermediateExpr::BinOp {
                result: reader.read_u64::<BigEndian>()?.into(),
                lhs: IntermediateValue::from_reader(reader)?,
                op: BinOp::try_from(reader.read_u8()?)
                    .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "invalid bin op"))?,
                rhs: IntermediateValue::from_reader(reader)?,
            }),
            0x03 => Ok(IntermediateExpr::Identity {
                result: reader.read_u64::<BigEndian>()?.into(),
                value: IntermediateValue::from_reader(reader)?,
            }),
            b => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("invalid byte: {b}"),
            )),
        }
    }

    pub fn from_typed_expr(
        expr: TypedExpr,
        reuse_symbols: bool,
        symbols: &mut SymbolsTable,
    ) -> Vec<IntermediateExpr> {
        let ty = expr.ty(symbols);

        let mut exprs = Vec::new();
        let mut temp_var_context = TempVarContext::new(reuse_symbols, symbols);

        let value = Self::from_typed_expr_inner(expr, &mut exprs, &mut temp_var_context);

        if exprs.is_empty() {
            return vec![IntermediateExpr::Identity {
                result: temp_var_context.get_temp_var(ty),
                value,
            }];
        }

        exprs
    }

    fn from_typed_expr_inner(
        expr: TypedExpr,
        exprs: &mut Vec<IntermediateExpr>,
        temp_var_context: &mut TempVarContext,
    ) -> IntermediateValue {
        let ty = expr.ty(temp_var_context.symbols());

        match expr {
            TypedExpr::Int { value, .. } => IntermediateValue::Int { value },
            TypedExpr::Float { value, .. } => IntermediateValue::Float { value },
            TypedExpr::Var { name, .. } => IntermediateValue::Var { name },
            TypedExpr::BinOp { lhs, op, rhs, .. } => {
                let lhs = Self::from_typed_expr_inner(*lhs, exprs, temp_var_context);
                let rhs = Self::from_typed_expr_inner(*rhs, exprs, temp_var_context);

                if let IntermediateValue::Var { name } = lhs {
                    temp_var_context.free_temp_var(name)
                }

                if let IntermediateValue::Var { name } = rhs {
                    temp_var_context.free_temp_var(name)
                }

                let result = temp_var_context.get_temp_var(ty);

                exprs.push(IntermediateExpr::BinOp {
                    result,
                    lhs,
                    op,
                    rhs,
                });

                IntermediateValue::Var { name: result }
            },
            TypedExpr::IntToFloat { value, .. } => {
                let value = Self::from_typed_expr_inner(*value, exprs, temp_var_context);

                if let IntermediateValue::Var { name: value } = value {
                    temp_var_context.free_temp_var(value)
                }

                let result = temp_var_context.get_temp_var(ty);

                exprs.push(IntermediateExpr::IntToFloat { result, value });

                IntermediateValue::Var { name: result }
            },
        }
    }
}