lang/src/parser/nodes/expr.rs

use crate::parser::VecDspT;

pub use super::*;

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

pub enum ExprTy {
    Block(Body),
    Group(Box<Expr>),
    Ident(Ident),
    Lit(Lit),
    Negate(Box<Expr>),
    Call { target: Box<Expr>, args: Vec<Expr> },
    Assign { target: Box<Expr>, val: Box<Expr> },
    If { cond: Box<Expr>, body: Box<Expr> },
    Loop { body: Box<Expr> },
    While { cond: Box<Expr>, body: Box<Expr> },
    Fn(Box<Func>),
}

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::OpenParen => {
                    ctx.next();
                    let target = Box::new(res);
                    let args = ctx.list(Token::Comma, Token::CloseParen)?;
                    ExprTy::Call { target, args }
                }
                _ => 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::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(())
            }
        }
    }
}

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)
            }
            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::Fn(f) => f.ends_with_block(),
            _ => false,
        }
    }
}