lang/src/parser/nodes/expr.rs

use crate::parser::VecDspT;

pub use super::*;

pub struct Expr {
    pub span: Span,
    pub ty: ExprTy,
}

pub enum ExprTy {
    Block(Body),
    Group(Box<Expr>),
    Member {
        of: Box<Expr>,
        field: Ident,
    },
    Ident(Ident),
    Lit(Lit),
    Negate(Box<Expr>),
    Call {
        target: Box<Expr>,
        args: Vec<Expr>,
    },
    Assign {
        target: Box<Expr>,
        val: Box<Expr>,
    },
    Define {
        target: Box<Expr>,
        ty: Option<Type>,
        const_: bool,
        val: Box<Expr>,
    },
    If {
        cond: Box<Expr>,
        body: Box<Expr>,
    },
    Loop {
        body: Box<Expr>,
    },
    While {
        cond: Box<Expr>,
        body: Box<Expr>,
    },
    Import(Ident),
    Fn(Box<Func>),
    Break,
    Asm(AsmBlock),
}

impl Node for Expr {
    fn parse(ctx: &mut ParseCtx) -> Result<Self, CompilerMsg> {
        let mut res = Self::unit(ctx)?;
        while let Some(next) = ctx.peek() {
            let ty = match next {
                Token::Equal => {
                    ctx.next();
                    let target = Box::new(res);
                    let val = Box::new(ctx.parse_with(Self::unit)?);
                    ExprTy::Assign { target, val }
                }
                Token::Colon => {
                    ctx.next();
                    let target = Box::new(res);
                    let mut ty = None;
                    let next = ctx.expect_peek()?;
                    if !matches!(next, Token::Equal | Token::Colon) {
                        ty = Some(ctx.parse()?);
                    }
                    let const_ = match ctx.expect_next()? {
                        Token::Equal => false,
                        Token::Colon => true,
                        t => ctx.unexpected(t, "an equals = or colon :")?,
                    };
                    let val = Box::new(ctx.parse_with(Self::unit)?);
                    ExprTy::Define {
                        target,
                        ty,
                        val,
                        const_,
                    }
                }
                Token::OpenParen => {
                    ctx.next();
                    let target = Box::new(res);
                    let args = ctx.list(Token::Comma, Token::CloseParen)?;
                    ExprTy::Call { target, args }
                }
                Token::Dot => {
                    ctx.next();
                    let of = Box::new(res);
                    let field = ctx.parse()?;
                    ExprTy::Member { of, field }
                }
                _ => break,
            };
            res = Self {
                ty,
                span: ctx.span(),
            };
        }
        Ok(res)
    }

    fn fmt(&self, f: &mut std::fmt::Formatter, ctx: DisplayCtx) -> std::fmt::Result {
        self.ty.fmt(f, ctx)
    }
}

impl ExprTy {
    fn fmt(&self, f: &mut std::fmt::Formatter, mut ctx: DisplayCtx) -> std::fmt::Result {
        match self {
            Self::Ident(ident) => ident.fmt(f, ctx),
            Self::Group(expr) => write!(f, "({})", expr.dsp(ctx)),
            Self::Fn(func) => func.fmt(f, ctx),
            Self::Lit(lit) => write!(f, "{}", lit),
            Self::Negate(expr) => {
                write!(f, "-{}", expr.dsp(ctx))
            }
            Self::Call { target, args } => {
                write!(f, "{}({})", target.dsp(ctx), args.dsp(ctx))
            }
            Self::Assign { target, val } => {
                write!(f, "{} = {}", target.dsp(ctx), val.dsp(ctx))
            }
            Self::Define {
                target,
                ty,
                val,
                const_,
            } => {
                write!(f, "{} :", target.dsp(ctx))?;
                if let Some(ty) = ty {
                    write!(f, " {} ", ty.dsp(ctx))?;
                }
                write!(f, "{} {}", if *const_ { ":" } else { "=" }, val.dsp(ctx))
            }
            Self::Member { of, field } => {
                write!(f, "{}.{field}", of.dsp(ctx))
            }
            Self::If { cond, body } => {
                write!(f, "if {} {}", cond.dsp(ctx), body.dsp(ctx))
            }
            Self::While { cond, body } => {
                write!(f, "while {} {}", cond.dsp(ctx), body.dsp(ctx))
            }
            Self::Loop { body } => {
                write!(f, "loop {}", body.dsp(ctx))
            }
            Self::Block(body) => {
                write!(f, "{{")?;
                if !body.items.is_empty() {
                    writeln!(f)?;
                    ctx.indent += 3;
                    body.fmt(f, ctx)?;
                }
                write!(f, "}}")?;
                Ok(())
            }
            Self::Import(ident) => {
                write!(f, "import {ident}")
            }
            Self::Break => {
                write!(f, "break")
            }
            Self::Asm(asm) => asm.fmt(f, ctx),
        }
    }
}

impl Expr {
    pub fn fmt_body(&self, f: &mut std::fmt::Formatter, ctx: DisplayCtx) -> std::fmt::Result {
        match self.ty {
            ExprTy::Block(_) => self.fmt(f, ctx),
            _ => write!(f, "=> {}", self.dsp(ctx)),
        }
    }

    fn unit(ctx: &mut ParseCtx) -> Result<Self, CompilerMsg> {
        let ty = match ctx.expect_next()? {
            Token::Dash => ExprTy::Negate(ctx.parse_box()?),
            Token::Ident(s) => ExprTy::Ident(ctx.ident(s)),
            Token::Lit(l) => ExprTy::Lit(ctx.lit(l)),
            Token::Fn => ExprTy::Fn(ctx.parse_box()?),
            Token::If => {
                let cond = ctx.parse_box()?;
                let body = Box::new(Self::body(ctx)?);
                ExprTy::If { cond, body }
            }
            Token::While => {
                let cond = ctx.parse_box()?;
                let body = Box::new(Self::body(ctx)?);
                ExprTy::While { cond, body }
            }
            Token::Loop => {
                let body = ctx.parse_box()?;
                ExprTy::Loop { body }
            }
            Token::OpenParen => {
                if ctx.next_if(Token::CloseParen) {
                    ExprTy::Lit(Lit {
                        ty: LitTy::Unit,
                        span: ctx.span(),
                    })
                } else {
                    let inner = ctx.parse_box()?;
                    ctx.expect(Token::CloseParen)?;
                    ExprTy::Group(inner)
                }
            }
            Token::OpenCurly => {
                let body = ctx.parse()?;
                ctx.expect(Token::CloseCurly)?;
                ExprTy::Block(body)
            }
            Token::Break => ExprTy::Break,
            Token::Import => {
                let ident = ctx.parse()?;
                ExprTy::Import(ident)
            }
            Token::Asm => ExprTy::Asm(ctx.parse()?),
            other => return ctx.unexpected(other, "an expression"),
        };
        Ok(Self {
            ty,
            span: ctx.span(),
        })
    }

    pub fn is_group(&self) -> bool {
        matches!(self.ty, ExprTy::Group(_))
    }

    pub fn is_block(&self) -> bool {
        matches!(self.ty, ExprTy::Block(_))
    }

    pub fn block(ctx: &mut ParseCtx) -> Result<Expr, CompilerMsg> {
        ctx.expect(Token::OpenCurly)?;
        let id = ctx.parse()?;
        ctx.expect(Token::CloseCurly)?;
        Ok(Expr {
            ty: ExprTy::Block(id),
            span: ctx.span(),
        })
    }

    pub fn body(ctx: &mut ParseCtx) -> Result<Expr, CompilerMsg> {
        if ctx.next_if(Token::DoubleArrow) {
            ctx.parse()
        } else {
            ctx.parse_with(Expr::block)
        }
    }

    pub fn ends_with_block(&self) -> bool {
        match &self.ty {
            ExprTy::Block(..) => true,
            ExprTy::Loop { body }
            | ExprTy::While { body, .. }
            | ExprTy::If { body, .. }
            | ExprTy::Negate(body)
            | ExprTy::Assign { val: body, .. } => body.ends_with_block(),
            ExprTy::Define { val: body, .. } => body.ends_with_block(),
            ExprTy::Fn(f) => f.ends_with_block(),
            _ => false,
        }
    }

    pub fn needs_semicolon(&self) -> bool {
        !self.ends_with_block()
    }
}