developing-compilers/src/ast/mod.rs

pub mod typed;
pub mod untyped;

use core::fmt;
use std::str::FromStr;

use typed::{Type, TypedExpr};
pub use untyped::UntypedExpr;

use crate::error::{self, SemanticError, SemanticErrorKind};
use crate::symbols::SymbolsTable;

#[derive(Debug, Copy, Clone)]
pub struct Span {
    pub start: usize,
    pub end: usize,
}

impl Span {
    pub fn new(start: usize, end: usize) -> Self {
        Self { start, end }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BinOp {
    Add,
    Sub,
    Mul,
    Div,
}

impl BinOp {
    pub fn precedence(&self) -> u8 {
        match self {
            BinOp::Add => 1,
            BinOp::Sub => 1,
            BinOp::Mul => 2,
            BinOp::Div => 2,
        }
    }
}

impl From<&BinOp> for u8 {
    fn from(value: &BinOp) -> Self {
        match value {
            BinOp::Add => 1,
            BinOp::Sub => 2,
            BinOp::Mul => 3,
            BinOp::Div => 4,
        }
    }
}

impl TryFrom<u8> for BinOp {
    type Error = u8;

    fn try_from(value: u8) -> Result<Self, Self::Error> {
        match value {
            1 => Ok(BinOp::Add),
            2 => Ok(BinOp::Sub),
            3 => Ok(BinOp::Mul),
            4 => Ok(BinOp::Div),
            b => Err(b),
        }
    }
}

impl fmt::Display for BinOp {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if f.alternate() {
            match self {
                BinOp::Add => f.write_str("add"),
                BinOp::Sub => f.write_str("sub"),
                BinOp::Mul => f.write_str("mul"),
                BinOp::Div => f.write_str("div"),
            }
        } else {
            match self {
                BinOp::Add => f.write_str("+"),
                BinOp::Sub => f.write_str("-"),
                BinOp::Mul => f.write_str("*"),
                BinOp::Div => f.write_str("/"),
            }
        }
    }
}

pub fn to_typed_expr(
    expr: UntypedExpr,
    optimize: bool,
    symbols: &mut SymbolsTable,
) -> error::Result<TypedExpr> {
    let expr = convert_to_typed_expr(expr, symbols)?;
    let expr = coerce_types(expr, symbols)?;
    let expr = if optimize { optimize_expr(expr) } else { expr };

    Ok(expr)
}

fn optimize_expr(expr: TypedExpr) -> TypedExpr {
    let expr = match expr {
        TypedExpr::BinOp { lhs, rhs, op, span } => {
            let lhs = optimize_expr(*lhs);
            let rhs = optimize_expr(*rhs);

            match (lhs, rhs) {
                (TypedExpr::Int { value: lhs, .. }, TypedExpr::Int { value: rhs, .. }) => {
                    TypedExpr::Int {
                        span,
                        value: match op {
                            BinOp::Add => lhs + rhs,
                            BinOp::Sub => lhs - rhs,
                            BinOp::Mul => lhs * rhs,
                            BinOp::Div => lhs / rhs,
                        },
                    }
                },
                (TypedExpr::Float { value: lhs, .. }, TypedExpr::Float { value: rhs, .. }) => {
                    TypedExpr::Float {
                        span,
                        value: match op {
                            BinOp::Add => lhs + rhs,
                            BinOp::Sub => lhs - rhs,
                            BinOp::Mul => lhs * rhs,
                            BinOp::Div => lhs / rhs,
                        },
                    }
                },
                (lhs, TypedExpr::Int { value: 0, .. })
                | (TypedExpr::Float { value: 0.0, .. }, lhs)
                    if matches!(op, BinOp::Add | BinOp::Sub) =>
                {
                    lhs
                },
                (lhs, TypedExpr::Int { value: 1, .. })
                | (lhs, TypedExpr::Float { value: 1.0, .. })
                    if matches!(op, BinOp::Mul | BinOp::Div) =>
                {
                    lhs
                },
                (TypedExpr::Int { value: 0, .. }, rhs)
                | (TypedExpr::Float { value: 0.0, .. }, rhs)
                    if matches!(op, BinOp::Add) =>
                {
                    rhs
                },
                (TypedExpr::Int { value: 1, .. }, rhs)
                | (TypedExpr::Float { value: 1.0, .. }, rhs)
                    if matches!(op, BinOp::Mul) =>
                {
                    rhs
                },
                (TypedExpr::Int { value: 0, .. }, _) | (_, TypedExpr::Int { value: 0, .. })
                    if matches!(op, BinOp::Mul) =>
                {
                    TypedExpr::Int { span, value: 0 }
                },
                (TypedExpr::Float { value: 0.0, .. }, _)
                | (_, TypedExpr::Float { value: 0.0, .. })
                    if matches!(op, BinOp::Mul) =>
                {
                    TypedExpr::Float { span, value: 0.0 }
                },
                (lhs, rhs) => TypedExpr::BinOp {
                    span,
                    lhs: Box::new(lhs),
                    op,
                    rhs: Box::new(rhs),
                },
            }
        },
        TypedExpr::IntToFloat { value } => {
            let value = optimize_expr(*value);
            if let TypedExpr::Int { value, span } = value {
                TypedExpr::Float {
                    span,
                    value: value as f64,
                }
            } else {
                TypedExpr::IntToFloat {
                    value: Box::new(value),
                }
            }
        },
        expr => expr,
    };

    expr
}

fn convert_to_typed_expr(
    expr: UntypedExpr,
    symbols: &mut SymbolsTable,
) -> error::Result<TypedExpr> {
    let expr = match expr {
        UntypedExpr::Int { span, value } => TypedExpr::Int { span, value },
        UntypedExpr::Float { span, value } => TypedExpr::Float { span, value },
        UntypedExpr::Var {
            span,
            name,
            typename,
        } => {
            let ty = typename
                .and_then(|t| symbols.resolve(&t))
                .map(|data| data.name.as_str())
                .map(Type::from_str)
                .transpose()
                .map_err(|e| SemanticError::new(span, e))?;
            {
                let symbol = symbols.resolve_mut(&name).unwrap();
                match (symbol.ty, ty) {
                    (Some(ty), Some(ty2)) if ty != ty2 => {
                        return Err(SemanticError::new(
                            span,
                            SemanticErrorKind::DuplicateSymbol(symbol.name.clone()),
                        )
                        .into());
                    },
                    (None, Some(ty)) => {
                        symbol.ty = Some(ty);
                    },
                    _ => {},
                }
            }

            TypedExpr::Var { span, name }
        },
        UntypedExpr::BinOp { span, lhs, op, rhs } => {
            let rhs = *rhs;
            let lhs = *lhs;

            match op {
                BinOp::Div
                    if matches!(rhs, UntypedExpr::Int { .. } | UntypedExpr::Float { .. }) =>
                {
                    match &rhs {
                        UntypedExpr::Int { span, value } if *value == 0 => {
                            return Err(SemanticError::new(
                                *span,
                                SemanticErrorKind::DivisionByZero,
                            )
                            .into());
                        },
                        UntypedExpr::Float { span, value } if *value == 0.0 => {
                            return Err(SemanticError::new(
                                *span,
                                SemanticErrorKind::DivisionByZero,
                            )
                            .into());
                        },
                        _ => {},
                    }
                },
                _ => {},
            }

            let lhs = convert_to_typed_expr(lhs, symbols)?;
            let rhs = convert_to_typed_expr(rhs, symbols)?;

            TypedExpr::BinOp {
                span,
                lhs: Box::new(lhs),
                op,
                rhs: Box::new(rhs),
            }
        },
    };

    Ok(expr)
}

fn coerce_types(expr: TypedExpr, symbols: &mut SymbolsTable) -> error::Result<TypedExpr> {
    let expr = match expr {
        TypedExpr::Int { .. } => expr,
        TypedExpr::Float { .. } => expr,
        TypedExpr::Var { name, .. } => {
            let symbol = symbols.resolve_mut(&name).unwrap();

            if symbol.ty.is_none() {
                symbol.ty = Some(Type::Int);
            }

            expr
        },
        TypedExpr::BinOp { lhs, rhs, op, span } => {
            let lhs = coerce_types(*lhs, symbols)?;
            let rhs = coerce_types(*rhs, symbols)?;

            let (lhs, rhs) = match (lhs.ty(symbols), rhs.ty(symbols)) {
                (Type::Int, Type::Int) => (lhs, rhs),
                (Type::Float, Type::Float) => (lhs, rhs),
                (Type::Int, Type::Float) => {
                    let lhs = TypedExpr::cast_to_float(lhs);
                    (lhs, rhs)
                },
                (Type::Float, Type::Int) => {
                    let rhs = TypedExpr::cast_to_float(rhs);
                    (lhs, rhs)
                },
            };

            TypedExpr::BinOp {
                span,
                lhs: Box::new(lhs),
                op,
                rhs: Box::new(rhs),
            }
        },
        TypedExpr::IntToFloat { .. } => expr,
    };

    Ok(expr)
}