path: root/matrix-lang/src/ast.rs

use std::{ops::{Neg, Not}, fmt::{Debug, Display}};

use crate::prelude::*;

pub type AstName = (Rc<str>, Position);
pub type InlineList = Vec<Expr>;
pub type InlineMatrix = (usize, usize, Vec<Expr>);
pub type InlineTable = Vec<(Expr, Expr)>;

#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub enum UnaryOp {
    // normal math
    Negate = 1,
    // equality
    Not = 2,
}

impl TryFrom<u8> for UnaryOp {
    type Error = Exception;

    fn try_from(value: u8) -> std::prelude::v1::Result<Self, Self::Error> {
        if value < 1 || value > 2 {
            Err(exception!(BINARY_EXCEPTION, "cannot convert {value} to UnaryOp"))
        } else {
            unsafe { Ok(std::mem::transmute(value)) }
        }
    }
}

#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub enum BinaryOp {
    // normal math
    Add = 1,
    Subtract = 2,
    Multiply = 3,
    Divide = 4,
    Modulo = 5,
    Power = 6,
    // binary math
    BitwiseAnd = 7,
    BitwiseOr = 8,
    BitwiseXor = 9,
    BitwiseShiftLeft = 10,
    BitwiseShiftRight = 11,
    // equality
    Equals = 12,
    NotEquals = 13,
    GreaterEquals = 14,
    LessEquals = 15,
    GreaterThan = 16,
    LessThan = 17,
    // iter
    Range = 18,
    RangeEq = 19
}

impl TryFrom<u8> for BinaryOp {
    type Error = Exception;

    fn try_from(value: u8) -> std::prelude::v1::Result<Self, Self::Error> {
        if value < 1 || value > 19 {
            Err(exception!(BINARY_EXCEPTION, "cannot convert {value} to BinaryOp"))
        } else {
            unsafe { Ok(std::mem::transmute(value)) }
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub enum ExprData {
    NoOp,

    Literal(Value),
    Ident(Rc<str>),

    UnaryOp(Box<Expr>, UnaryOp),
    BinaryOp(Box<Expr>, Box<Expr>, BinaryOp),

    Index(Box<Expr>, Vec<Expr>),
    FnCall(Box<Expr>, Vec<Expr>),
    FieldAccess(Box<Expr>, AstName),

    List(InlineList),
    Matrix(InlineMatrix),
    Table(InlineTable),

    And(Box<Expr>, Box<Expr>),
    Or(Box<Expr>, Box<Expr>),

    Assign(Box<Expr>, Box<Expr>),

    If(Box<Expr>, Box<Expr>, Option<Box<Expr>>),
    Function(AstName, Vec<AstName>, Box<Expr>, bool),
    Lambda(Vec<AstName>, Box<Expr>, bool),

    Loop(Box<Expr>),
    While(Box<Expr>, Box<Expr>),
    DoWhile(Box<Expr>, Box<Expr>),
    For(AstName, Box<Expr>, Box<Expr>),

    Block(Vec<Expr>),

    Try(Box<Expr>, AstName, Box<Expr>),

    Let(AstName, Box<Expr>),
    Const(AstName, Box<Expr>),

    Pipeline(Box<Expr>, Box<Expr>),

    Continue,
    Break,
    Return(Box<Expr>),
}

impl Display for Expr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{self:?}")
    }
}

impl Debug for Expr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{:?}", self.data)
    }
}

#[derive(Clone, PartialEq)]
pub struct Expr {
    pub data: ExprData,
    pub pos: Position
}

impl From<(ExprData, Position)> for Expr {
    fn from(value: (ExprData, Position)) -> Self {
        Self { data: value.0, pos: value.1 }
    }
}

impl Neg for BinaryOp {
    type Output = Option<Self>;
    fn neg(self) -> Self::Output {
        use BinaryOp::*;
        Some(match self {
            Add => Subtract,
            Subtract => Add,
            _ => return None
        })
    }
}

impl Not for BinaryOp {
    type Output = Option<Self>;
    fn not(self) ->  Self::Output {
        use BinaryOp::*;
        Some(match self {
            Equals => NotEquals,
            NotEquals => Equals,
            GreaterEquals => LessThan,
            LessEquals => GreaterThan,
            GreaterThan => LessEquals,
            LessThan => GreaterEquals,
            _ => return None
        })
    }
}

impl Neg for Expr {
    type Output = std::result::Result<Self, Self>;
    fn neg(mut self) -> Self::Output {
        use ExprData as E;
        let mut pos = self.pos;
        let data = match self.data {
            E::Literal(v) => E::Literal(-v),
            E::BinaryOp(lhs, rhs, op) => {
                let Some(op_n) = -op else {
                    return Err((E::BinaryOp(lhs, rhs, op), pos).into());
                };
                if let Ok(lhs) = -*lhs.clone() {
                    pos = lhs.pos;
                    E::BinaryOp(Box::new(lhs), rhs, op_n)
                } else if let Ok(rhs) = -*rhs.clone() {
                    pos = rhs.pos;
                    E::BinaryOp(lhs, Box::new(rhs), op_n)
                } else {
                    return Err((E::BinaryOp(lhs, rhs, op), pos).into());
                }
            },
            E::UnaryOp(expr, op) => {
                match op {
                    UnaryOp::Negate => {
                        pos = expr.pos;
                        expr.data
                    },
                    _ => return Err((E::UnaryOp(expr, op), pos).into())
                }
            }
            data => return Err((data, pos).into())
        };
        self.data = data;
        self.pos = pos;
        Ok(self)
    }
}

impl Not for Expr {
    type Output = std::result::Result<Self, Self>;
    fn not(mut self) -> Self::Output {
        use ExprData as E;
        let mut pos = self.pos;
        let data = match self.data {
            E::Literal(v) => E::Literal(Value::Bool(!v)),
            E::UnaryOp(expr, op) => {
                match op {
                    self::UnaryOp::Not => {
                        pos = expr.pos;
                        expr.data
                    },
                    _ => return Err((E::UnaryOp(expr, op), pos).into())
                }
            }
            data => return Err((data, pos).into())
        };
        self.data = data;
        self.pos = pos;
        Ok(self)
    }
}

pub fn optimize(mut expr: Expr) -> Result<Expr> {
    use ExprData as E;
    let mut pos = expr.pos;
    let data: ExprData = match expr.data {
        E::UnaryOp(expr, op) => {
            let expr = optimize(*expr)?;
            match match op {
                UnaryOp::Negate => -expr,
                UnaryOp::Not => !expr,
            } {
                Ok(expr) => {
                    pos = expr.pos;
                    expr.data
                },
                Err(expr) => E::UnaryOp(Box::new(expr), op)
            }
        },
        E::BinaryOp(lhs, rhs, op) => {
            let lhs = optimize(*lhs)?;
            let rhs = optimize(*rhs)?;
            if let (E::Literal(l), E::Literal(r)) = (lhs.clone().data, rhs.clone().data) {
                match Value::binary_op(op, l, r) {
                    Err(err) => return Err(err),
                    Ok(value) => E::Literal(value),
                }
            } else {
                E::BinaryOp(Box::new(lhs), Box::new(rhs), op)
            }
        },
        E::FnCall(ident, values) => {
            E::FnCall(ident, values
                   .into_iter()
                   .map(optimize)
                   .collect::<Result<Vec<Expr>>>()?)
        }
        E::FieldAccess(expr, ident) => {
            let expr = optimize(*expr)?;
            E::FieldAccess(Box::new(expr), ident)
        },
        E::List(list) =>
            E::List(list.into_iter()
                       .map(optimize)
                       .collect::<Result<Vec<Expr>>>()?),
        E::Matrix(mat) =>
            E::Matrix((mat.0, mat.1,
                mat.2.into_iter().map(optimize)
                .collect::<Result<Vec<Expr>>>()?)),
        E::Table(table) =>
            E::Table(table
                .into_iter()
                .map(|(k, v)| {
                    let k = optimize(k)?;
                    let v = optimize(v)?;
                    Ok((k, v))
                }).collect::<Result<Vec<(Expr, Expr)>>>()?),
        E::And(lhs, rhs) => {
            let lhs = optimize(*lhs)?;
            let rhs = optimize(*rhs)?;
            if let (E::Literal(l), r) = (lhs.clone().data, rhs.clone().data) {
                match !!l.clone() {
                    true => {
                        pos = rhs.pos;
                        r
                    },
                    false => {
                        pos = lhs.pos;
                        E::Literal(l)
                    },
                }
            } else {
                E::And(Box::new(lhs), Box::new(rhs))
            }
        },
        E::Or(lhs, rhs) => {
            let lhs = optimize(*lhs)?;
            let rhs = optimize(*rhs)?;
            if let (E::Literal(l), r) = (lhs.clone().data, rhs.clone().data) {
                match !l.clone() {
                    true => {
                        pos = rhs.pos;
                        r
                    },
                    false => {
                        pos = lhs.pos;
                        E::Literal(l)
                    },
                }
            } else {
                E::And(Box::new(lhs), Box::new(rhs))
            }
        },
        E::Block(b) => {
            let len = b.len();
            let b: Vec<Expr> =
                b.into_iter()
                .map(optimize)
                .collect::<Result<Vec<Expr>>>()?
                .into_iter()
                .enumerate()
                .filter(|(i, e)| {
                    if let E::Literal(_) = e.data {
                        return i + 1 == len
                    }
                    E::NoOp != e.data
                })
                .map(|e| e.1)
                .collect();
            if b.is_empty() {
                E::NoOp
            } else {
                E::Block(b)
            }
        },
        E::If(cond, block, else_block) => {
            let cond = optimize(*cond)?;
            let block = optimize(*block)?;
            let else_block = else_block.map(|e| optimize(*e)).transpose()?;
            if let E::Literal(lit) = cond.data {
                if !!lit {
                    pos = block.pos;
                    block.data
                } else if let Some(else_block) = else_block {
                    pos = else_block.pos;
                    else_block.data
                } else {
                    E::NoOp
                }
            } else {
                E::If(Box::new(cond), Box::new(block), else_block.map(|e| Box::new(e)))
            }
        },
        E::While(cond, block) => {
            let cond = optimize(*cond)?;
            let block = optimize(*block)?;
            if let E::Literal(lit) = cond.data {
                if !!lit {
                    E::Loop(Box::new(block))
                } else {
                    E::NoOp
                }
            } else {
                E::While(Box::new(cond), Box::new(block))
            }
        },
        E::For(name, cond, block) => {
            let cond = optimize(*cond)?;
            let block = optimize(*block)?;
            E::For(name, Box::new(cond), Box::new(block))
        }
        E::DoWhile(block, cond) => {
            let cond = optimize(*cond)?;
            let block = optimize(*block)?;
            if let E::Literal(lit) = &cond.data {
                if !!lit.clone() {
                    E::Loop(Box::new(block))
                } else {
                    E::DoWhile(Box::new(block), Box::new(cond))
                }
            } else {
                E::DoWhile(Box::new(block), Box::new(cond))
            }
        }
        E::Loop(block) => {
            let block = Box::new(optimize(*block)?);
            E::Loop(block)
        },
        E::Try(expr, err, catch) => {
            let expr = Box::new(optimize(*expr)?);
            let catch = Box::new(optimize(*catch)?);
            E::Try(expr, err, catch)
        },
        E::Function(ident, params, stmt, varadic) => {
            let stmt = Box::new(optimize(*stmt)?);
            E::Function(ident, params, stmt, varadic)
        }
        E::Lambda(params, stmt, varadic) => {
            let stmt = Box::new(optimize(*stmt)?);
            E::Lambda(params, stmt, varadic)
        },
        E::Let(ident, expr) => {
            E::Let(ident, Box::new(optimize(*expr)?))
        },
        E::Const(ident, expr) => {
            E::Const(ident, Box::new(optimize(*expr)?))
        },
        E::Assign(lhs, rhs) => {
            let lhs = Box::new(optimize(*lhs)?);
            let rhs = Box::new(optimize(*rhs)?);
            E::Assign(lhs, rhs)
        },
        E::Return(expr) => {
            let expr = Box::new(optimize(*expr)?);
            E::Return(expr)
        },
        E::Pipeline(lhs, rhs) => {
            let lhs = Box::new(optimize(*lhs)?);
            let rhs = Box::new(optimize(*rhs)?);
            E::Pipeline(lhs, rhs)
        }
        data => data
    };
    expr.data = data;
    expr.pos = pos;
    Ok(expr)
}

impl From<TokenData> for UnaryOp {
    fn from(value: TokenData) -> Self {
        use TokenData as T;
        match value {
            T::Subtract => Self::Negate,
            T::Not => Self::Not,
            _ => panic!("aaaaa")
        }
    }
}

impl From<TokenData> for BinaryOp {
    fn from(value: TokenData) -> Self {
        use TokenData as T;
        match value {
            T::Equal => Self::Equals,
            T::NotEqual => Self::NotEquals,
            T::GreaterEqual => Self::GreaterEquals,
            T::LessEqual => Self::LessEquals,
            T::GreaterThan => Self::GreaterThan,
            T::LessThan => Self::LessThan,
            T::BitwiseShiftLeft => Self::BitwiseShiftLeft,
            T::BitwiseShiftRight => Self::BitwiseShiftRight,
            T::BitwiseAnd => Self::BitwiseAnd,
            T::BitwiseOr => Self::BitwiseOr,
            T::BitwiseXor => Self::BitwiseXor,
            T::Add => Self::Add,
            T::Subtract => Self::Subtract,
            T::Multiply => Self::Multiply,
            T::Divide => Self::Divide,
            T::Modulo => Self::Modulo,
            T::Power => Self::Power,
            _ => panic!("aaaaa")
        }
    }
}

impl Expr {
    pub fn is_assignable(&self) -> bool {
        use ExprData as E;
        match self.data {
            E::Ident(_) => true,
            E::Index(_, _) => true,
            E::FieldAccess(_, _) => true,
            _ => false,
        }
    }
}