skullbrain-parser 0.1.0

use crate::cst::utils::token_with;

use super::*;

/// The set of all expression-like AST nodes and tokens.
pub enum Expr {
    /// Literal: numbers, strings, etc.
    Literal(LiteralToken),

    /// Identifier.
    Ident(IdentToken),

    /// Parenthesized sub-expression: `( ... )`
    Paren(ParenExpr),

    /// Brace block as an expression: `{ ... }`
    Block(BlockExpr),
    /// Binary operator expression.
    Binary(BinaryExpr),
    ///
    Unary(UnaryExpr),
    ///
    Chain(ChainExpr),

    /// Control-flow expressions/statements:
    Break(Break),
    ///
    Continue(Continue),
    ///
    Return(Return),
    ///
    Throw(Throw),

    /// Loops:
    While(While),
    ///
    For(ForLoop),

    /// Try / Catch:
    TryCatch(TryCatch),

    /// Conditional:
    IfElse(IfElse),

    /// Match:
    Match(Match),

    /// Declarations as expressions (e.g., in REPL contexts):
    Variable(VariableStmt),
    ///
    Import(ImportStmt),
    ///
    TypeAlias(TypeAlias),
    ///
    Function(Function),
    ///
    Trait(Trait),
    ///
    TraitImpl(TraitImpl),
    ///
    Class(Class),
    ///
    Enum(Enum),

    ///
    Fragment(Fragment),
}

impl CstNode for Expr {
    fn can_cast(_kind: SyntaxKind) -> bool {
        true
    }

    fn syntax(&self) -> &SyntaxNode {
        match self {
            Expr::Break(n) => n.syntax(),
            Expr::Continue(n) => n.syntax(),
            Expr::Return(n) => n.syntax(),
            Expr::Throw(n) => n.syntax(),
            Expr::While(n) => n.syntax(),
            Expr::For(n) => n.syntax(),
            Expr::TryCatch(n) => n.syntax(),
            Expr::IfElse(n) => n.syntax(),
            Expr::Match(n) => n.syntax(),
            Expr::Variable(n) => n.syntax(),
            Expr::Import(n) => n.syntax(),
            Expr::TypeAlias(n) => n.syntax(),
            Expr::Function(n) => n.syntax(),
            Expr::Trait(n) => n.syntax(),
            Expr::TraitImpl(n) => n.syntax(),
            Expr::Class(n) => n.syntax(),
            Expr::Enum(n) => n.syntax(),

            Expr::Binary(n) => &n.syntax,
            Expr::Unary(n) => &n.syntax,
            Expr::Chain(n) => &n.syntax,
            Expr::Paren(n) => &n.syntax,
            Expr::Block(n) => &n.syntax,
            Expr::Fragment(n) => match n {
                Fragment::EnumVariant(enum_variant_decl) => &enum_variant_decl.syntax,
                Fragment::Modifiers(modifiers) => &modifiers.syntax,
                Fragment::Generic { syntax } => &syntax,
                Fragment::Unknown(_) => todo!(),
                Fragment::MatchArm(match_arm) => match match_arm {
                    MatchArm::WithCondition(match_arm_with_condition) => {
                        &match_arm_with_condition.syntax
                    }
                    MatchArm::Default(match_arm_default) => &match_arm_default.syntax,
                },
            },
            Expr::Literal(_) | Expr::Ident(_) => todo!(),
        }
    }

    fn cast(syntax: SyntaxNode) -> Option<Self> {
        let kind = syntax.kind();
        Some(match kind {
            SyntaxKind::DelimitedRange => disambiguate_delimited(syntax),
            SyntaxKind::Operator => disambiguate_operator(syntax),
            SyntaxKind::LoopBreak => Expr::Break(Break { syntax }),
            SyntaxKind::LoopContinue => Expr::Continue(Continue { syntax }),
            SyntaxKind::Return => Expr::Return(Return { syntax }),
            SyntaxKind::ThrowError => Expr::Throw(Throw { syntax }),

            SyntaxKind::WhileLoop => Expr::While(While { syntax }),
            SyntaxKind::ForLoop => Expr::For(ForLoop { syntax }),
            SyntaxKind::TryCatch => Expr::TryCatch(TryCatch { syntax }),
            SyntaxKind::IfElse => Expr::IfElse(IfElse { syntax }),
            SyntaxKind::Match => Expr::Match(Match { syntax }),
            SyntaxKind::VariableDeclaration => Expr::Variable(VariableStmt { syntax }),
            SyntaxKind::Import => Expr::Import(ImportStmt { syntax }),
            SyntaxKind::TypeAlias => Expr::TypeAlias(TypeAlias { syntax }),
            SyntaxKind::Function => Expr::Function(Function { syntax }),
            SyntaxKind::Trait => Expr::Trait(Trait { syntax }),
            SyntaxKind::TraitImpl => Expr::TraitImpl(TraitImpl { syntax }),
            SyntaxKind::Class => Expr::Class(Class { syntax }),
            SyntaxKind::Enum => Expr::Enum(Enum { syntax }),
            SyntaxKind::EnumVariant => {
                Expr::Fragment(Fragment::EnumVariant(EnumVariant { syntax }))
            }
            SyntaxKind::Modifiers => Expr::Fragment(Fragment::Modifiers(Modifiers { syntax })),
            SyntaxKind::MatchCondition => Expr::Fragment(Fragment::MatchArm(
                MatchArm::WithCondition(MatchArmWithCondition { syntax }),
            )),
            SyntaxKind::MatchConditionDefault => {
                Expr::Fragment(Fragment::MatchArm(MatchArm::Default(MatchArmDefault {
                    syntax,
                })))
            }
            _ => Expr::Fragment(Fragment::Generic { syntax }),
        })
    }
}

fn disambiguate_delimited(syntax: SyntaxNode) -> Expr {
    let para = ParenRange { syntax };

    if para.open_token().is_some() || para.close_token().is_some() {
        Expr::Paren(ParenExpr {
            syntax: para.syntax,
        })
    } else {
        Expr::Block(BlockExpr {
            syntax: para.syntax,
        })
    }
}

fn disambiguate_operator(syntax: SyntaxNode) -> Expr {
    const fn is(kind: SyntaxKind) -> bool {
        matches!(
            kind,
            SyntaxKind::OpMethodChaining
                | OpMinus
                | OpNot
                | OpEq
                | OpNotEq
                | OpLt
                | OpGt
                | OpLtEq
                | OpGtEq
                | OpAnd
                | OpOr
                | OpPlus
                | OpMul
                | OpDivide
                | OpMod
        )
    }

    let Some(token) = token_with(&syntax, is) else {
        unreachable!()
    };

    use SyntaxKind::*;
    match token.kind() {
        OpMethodChaining => Expr::Chain(ChainExpr { syntax }),
        OpEq | OpNotEq | OpLt | OpGt | OpLtEq | OpGtEq | OpAnd | OpOr | OpPlus | OpMul
        | OpDivide | OpMod => Expr::Binary(BinaryExpr { syntax }),
        OpNot => Expr::Unary(UnaryExpr { syntax }),
        OpMinus => {
            let binary = BinaryExpr { syntax };

            // Since we dont create a lhs node when its unary-minus
            match binary.lhs() {
                None => Expr::Unary(UnaryExpr {
                    syntax: binary.syntax,
                }),
                _ => Expr::Binary(binary),
            }
        }
        _ => unreachable!(),
    }
}

impl Expr {
    pub(crate) fn cast_element(element: SyntaxElementRef) -> Expr {
        match element {
            cstree::util::NodeOrToken::Node(node) => {
                let Some(expr) = Expr::cast(node.clone()) else {
                    unreachable!()
                };

                expr
            }
            cstree::util::NodeOrToken::Token(token) => {
                let token = token.clone();
                match token.kind() {
                    SyntaxKind::Unknown => {
                        Expr::Fragment(Fragment::Unknown(UnknownToken { token }))
                    }
                    SyntaxKind::Literal => Expr::Literal(LiteralToken { token }),
                    SyntaxKind::Ident => Expr::Ident(IdentToken { token }),

                    _ => unreachable!(),
                }
            }
        }
    }
}