ast.rs

use std::ops::Range;

#[derive(Debug, Clone)]
pub struct ASTNode {
    pub span: Range<usize>,
    pub file: String,
    pub node: ASTNodeKind,
    pub attributes: Vec<Attribute>,
}

#[derive(Debug, Clone)]
pub enum ASTNodeKind {
    Expr(Expr),
    Function(Box<Function>),
    Struct(Struct),
    Enum(Enum),
    TypeAlias(TypeAlias),
    Impl(Impl),
    Trait(TraitDef),
    MacroDef(MacroDef),
    EffectDef(EffectDef),

    Error,
}

#[derive(Debug, Clone)]
pub struct TypeAnnot {
    pub span: Range<usize>,
    pub file: String,
    pub expr: TypeAnnotKind,
}

#[derive(Debug, Clone)]
pub enum TypeAnnotKind {
    // Primitives
    Bool,
    Int,
    Float,
    String,
    Unit,
    Never,

    // Named types
    Named(String),

    // Generic types: List<T>, Option<T>
    Generic {
        name: String,
        args: Vec<TypeAnnot>,
        kind: Option<KindAnnot>,
    },

    // Function types with effects
    Function {
        params: Vec<TypeAnnot>,
        return_type: Box<TypeAnnot>,
        effects: EffectAnnot,
    },

    // Tuple types
    Tuple(Vec<TypeAnnot>),

    // Union types
    Union(Vec<TypeAnnot>),

    // Trait bounds
    Trait(Vec<String>),

    // HKT: Type constructor with explicit kind
    Constructor {
        name: String,
        kind: KindAnnot,
    },

    // Type variable (in forall/exists)
    Variable {
        name: String,
        kind: Option<KindAnnot>,
    },

    // GADT: Universal quantification
    // forall T. T ~ Int => T -> T
    Forall {
        vars: Vec<(String, KindAnnot)>,
        constraints: Vec<ConstraintAnnot>,
        body: Box<TypeAnnot>,
    },

    // GADT: Existential types
    // exists T. (T, T -> Int)
    Exists {
        vars: Vec<(String, KindAnnot)>,
        constraints: Vec<ConstraintAnnot>,
        body: Box<TypeAnnot>,
    },

    Error,
}

// Kind annotations for HKTs
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum KindAnnot {
    Star,                                  // *
    Arrow(Box<KindAnnot>, Box<KindAnnot>), // * -> *
}

// Effect annotations
#[derive(Debug, Clone)]
pub struct EffectAnnot {
    pub effects: Vec<String>, // IO, State, Exn, etc.
    pub rest: Option<String>, // polymorphic effect variable
}

impl EffectAnnot {
    pub fn pure() -> Self {
        Self {
            effects: vec![],
            rest: None,
        }
    }

    pub fn closed(effects: Vec<String>) -> Self {
        Self {
            effects,
            rest: None,
        }
    }

    pub fn open(effects: Vec<String>, rest: String) -> Self {
        Self {
            effects,
            rest: Some(rest),
        }
    }
}

// Type constraints (for GADTs)
#[derive(Debug, Clone)]
pub enum ConstraintAnnot {
    // Type equality: T ~ Int
    Equal(TypeAnnot, TypeAnnot),

    // Trait bound: T: Show
    Trait(String, String), // (type_var, trait_name)
}

#[derive(Debug, Clone)]
pub struct Expr {
    pub span: Range<usize>,
    pub file: String,
    pub expr: ExprKind,
}

#[derive(Debug, Clone)]
pub enum ExprKind {
    Int(i64),
    Float(f64),
    Bool(bool),
    String(String),

    Variable(String),

    BinOp(Box<Expr>, BinOp, Box<Expr>),
    UnOp(UnOp, Box<Expr>),
    Assign {
        l_val: Box<Expr>,
        r_val: Box<Expr>,
        op: AssignOp,
    },
    Let {
        var: String,
        type_annot: Option<TypeAnnot>,
        value: Box<Expr>,
    },

    Array(Vec<Expr>),
    Tuple(Vec<Expr>),
    Map(Vec<(Expr, Expr)>),
    EnumConstruct {
        name: String,
        variant: String,
        args: Vec<Expr>,
    },

    Perform {
        effect: String,
        args: Vec<Expr>,
    },
    Handle {
        body: Box<Expr>,
        handlers: Vec<EffectHandler>,
    },

    Lambda {
        args: Vec<FnArg>,
        expression: Box<Expr>,
    },
    Cast {
        expr: Box<Expr>,
        target_type: TypeAnnot,
    },

    IfElse {
        condition: Box<Expr>,
        then: Box<Expr>,
        else_: Option<Box<Expr>>,
    },
    Block(Vec<Expr>),
    With {
        context: Box<Expr>,
        var: String,
        body: Box<Expr>,
    },
    Loop {
        label: Option<String>,
        body: Box<Expr>,
    },
    Match(Box<Expr>, Vec<MatchArm>),
    For {
        iterator: Box<Expr>, 
        value: String, 
        expression: Box<Expr>
    },
    While(Box<Expr>, Box<Expr>),
    IfLet {
        pattern: Pattern,
        expr: Box<Expr>,
        then: Box<Expr>,
        else_: Option<Box<Expr>>,
    },
    WhileLet {
        pattern: Pattern,
        expr: Box<Expr>,
        body: Box<Expr>,
    },

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

    Call(Box<Expr>, Vec<Expr>),
    Index(Box<Expr>, Box<Expr>),
    FieldAccess(Box<Expr>, String),
    OptionalChain(Box<Expr>, String),

    MacroCall(String, Vec<Expr>, Delimiter),
    Import(Import),

    Error,
}

#[derive(Debug, Clone)]
pub struct EffectHandler {
    pub span: Range<usize>,
    pub effect: String,
    pub params: Vec<String>,
    pub resume_param: String, // The continuation parameter
    pub body: Expr,
}

#[derive(Debug, Clone)]
pub struct Import {
    pub span: Range<usize>,
    pub file: String,
    pub path: Vec<String>, // ["abc", "ab"] for "abc/ab.ni"
    pub items: Vec<(String, Option<String>)>, // for selective imports + aliases
    // none implies all being imported
    pub alias: Option<String>,
}

#[derive(Debug, Clone)]
pub struct TraitDef {
    pub span: Range<usize>,
    pub name: String,
    pub type_params: Vec<TypeParam>,
    pub super_traits: Vec<TypeAnnot>,
    pub methods: Vec<Function>,
    pub associated_types: Vec<AssociatedType>,
}

#[derive(Debug, Clone)]
pub struct AssociatedType {
    pub name: String,
    pub bounds: Vec<TypeAnnot>,
}

#[derive(Debug, Clone)]
pub enum UnOp {
    Not,
    Plus,
    Minus,
    Unwrap,
}

#[derive(Debug, Clone)]
pub enum BinOp {
    Add,
    Sub,
    Mul,
    Div,
    Mod,
    Pow,

    Eq,
    Greater,
    Less,
    GreaterEq,
    LessEq,
    NotEq,

    Pipe,

    And,
    Or,
    Nor,
    Xor,
}

#[derive(Debug, Clone)]
pub enum AssignOp {
    Assign,
    AddAssign,
    SubAssign,
    MulAssign,
    DivAssign,
    ModAssign,
}

#[derive(Debug, Clone)]
pub struct MatchArm {
    pub pattern: Pattern,
    pub guard: Option<Expr>,
    pub body: Box<Expr>,
    pub span: Range<usize>,
}

#[derive(Debug, Clone)]
pub struct Pattern {
    pub span: Range<usize>,
    pub file: String,
    pub pat: PatKind,
}

#[derive(Debug, Clone)]
pub enum PatKind {
    Wildcard,
    Bind(String),
    Literal(Literal),
    Array(Vec<Pattern>),
    Or(Vec<Pattern>),
    As {
        name: String,
        pattern: Box<Pattern>,
    },
    Struct {
        name: String,
        fields: Vec<(String, Pattern)>,
    },
    Enum {
        name: String,
        variant: String,
        params: Vec<Pattern>,
    },
    Range(Box<Expr>, Box<Expr>),
    Rest(String),

    Error,
}

#[derive(Debug, Clone)]
pub enum Literal {
    Int(i64),
    Float(f64),
    Bool(bool),
    String(String),
}

#[derive(Debug, Clone)]
pub struct Function {
    pub span: Range<usize>,
    pub file: String,

    pub vis: Visibility,
    pub name: String,
    pub type_params: Vec<TypeParam>,
    pub args: Vec<FnArg>,
    pub return_type: Option<TypeAnnot>,
    pub where_constraints: Vec<TypeConstraint>,
    pub effects: EffectAnnot,
    pub body: Option<Expr>, // no body implies declaration
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Visibility {
    Public,
    Private,
}

#[derive(Debug, Clone)]
pub struct TypeParam {
    pub name: String,
    pub kind: Option<KindAnnot>,
    pub bounds: Vec<TypeAnnot>,
}

#[derive(Debug, Clone)]
pub struct TypeConstraint {
    pub span: Range<usize>,
    pub left: TypeAnnot,
    pub right: TypeAnnot,
}

#[derive(Debug, Clone)]
pub struct FnArg {
    pub span: Range<usize>,
    pub file: String,
    pub name: String,
    pub type_: Option<TypeAnnot>,
    // pub default: Option<Expr>,
}

#[derive(Debug, Clone)]
pub struct Struct {
    pub span: Range<usize>,
    pub file: String,
    pub name: String,
    pub type_params: Vec<TypeParam>,
    pub fields: Vec<(FnArg, Visibility)>,
    pub methods: Vec<Function>, // (desugar to impl)
    pub vis: Visibility,
}

#[derive(Debug, Clone)]
pub struct Enum {
    pub span: Range<usize>,
    pub file: String,
    pub name: String,
    pub type_params: Vec<TypeParam>,
    pub variants: Vec<EnumVariant>,
    pub methods: Vec<Function>, // (desugar to impl)
    pub vis: Visibility,
}

#[derive(Debug, Clone)]
pub struct EnumVariant {
    pub span: Range<usize>,
    pub file: String,
    pub name: String,
    pub types: Vec<TypeAnnot>,
    pub constraints: Vec<TypeConstraint>,
}

#[derive(Debug, Clone)]
pub struct TypeAlias {
    pub span: Range<usize>,
    pub file: String,
    pub name: String,
    pub type_params: Vec<TypeParam>,
    pub target_type: TypeAnnot,
    pub where_constraints: Vec<TypeConstraint>,
}

#[derive(Debug, Clone)]
pub struct Impl {
    pub span: Range<usize>,
    pub file: String,
    pub type_name: String,
    pub type_params: Vec<TypeParam>,
    pub methods: Vec<Function>,
    pub trait_: Option<String>,
    pub where_constraints: Vec<TypeConstraint>,
}

#[derive(Debug, Clone)]
pub struct Attribute {
    pub span: Range<usize>,
    pub name: String,
    pub args: Vec<AttributeArg>,
}

#[derive(Debug, Clone)]
pub enum AttributeArg {
    Ident(String),
    Literal(Literal),
    KeyValue(String, Box<Expr>),
}

#[derive(Debug, Clone)]
pub struct MacroDef {
    pub span: Range<usize>,
    pub file: String,
    pub name: String,
    pub rules: Vec<MacroRule>,
    pub hygiene: MacroHygiene,
}

#[derive(Debug, Clone)]
pub enum MacroHygiene {
    Hygienic,
    Unhygienic,
}

#[derive(Debug, Clone)]
pub struct MacroRule {
    pub span: Range<usize>,
    pub pattern: Vec<MacroToken>,
    pub body: Vec<MacroToken>,
}

#[derive(Debug, Clone)]
pub enum MacroToken {
    Ident(String),
    Literal(Literal),
    Punct(String),
    Group {
        delimiter: Delimiter,
        tokens: Vec<MacroToken>,
    },
    MetaVar {
        name: String,
        kind: String,
    }, // $x:expr, $t:ty
    Repeat {
        tokens: Vec<MacroToken>,
        separator: Option<String>,
        kind: RepeatKind,
    },
}

#[derive(Debug, Clone)]
pub enum Delimiter {
    Paren,   // ()
    Bracket, // []
    Brace,   // {}
}

#[derive(Debug, Clone)]
pub enum RepeatKind {
    ZeroOrMore, // *
    OneOrMore,  // +
    ZeroOrOne,  // ?
}

#[derive(Debug, Clone)]
pub struct EffectDef {
    pub span: Range<usize>,
    pub file: String,
    pub vis: Visibility,
    pub name: String,
    pub type_params: Vec<TypeParam>,
    pub operations: Vec<EffectOperation>,
    pub where_constraints: Vec<TypeConstraint>,
}

#[derive(Debug, Clone)]
pub struct EffectOperation {
    pub span: Range<usize>,
    pub name: String,
    pub params: Vec<TypeAnnot>,
    pub return_type: TypeAnnot,
}

#[derive(Debug, PartialEq, Eq, Hash, Clone)]
pub enum Kind {
    Star,
    Arrow(Box<Kind>, Box<Kind>),
}

grammar

// ? at the end implies "optional"
// * implies 0 or more
// + implies 1 or more
// | implies "or"
// kindof like regex

ASTNodeKind := 
    Expr
    | Function
    | Struct
    | Enum
    | Impl_block
    | Type_alias
    | Macro_def
    | Trait_def
    | Effect_def


Expr := 
	Int = Token::Int
	| Float = Token::Float
	| Bool = Token::Bool
	| String = Token::String
	| Variable = Token::Variable
	| BinOp = Expr BinOp Expr
	| UnOp = UnOp Expr // special case is unwrap. instead of going in FRONT it goes behind. so its Expr UnOp syntactically
	| Let = Token::KeywordLet name: Token::Variable type_annot: (Token::Colon TypeAnnot)? Token::Assign value: Expr
	| Array = Token::LBracket (Expr Token::Comma)* Token::LBracket
	| Tuple = Expr Token::Comma Expr*
	| Token::LParen Expr Token::RParen
	| Block = Token::LBrace (Expr Token::Semicolon?)+ Token::RBrace
	| Map = Token::LBrace (Expr Token::Colon Expr Token::Comma)+ Token::RBrace
	| EnumConstruct = Token::Variable Token::Access Token::Variable Token::LParen (Expr Token::Comma)* Token::RParen
	| Perform = Token::Perform effect: Token::Variable args: (Token::LParen (Expr Token::Comma)* Token::RParen)
	| Handle = Token::Handle body: Expr Token::With Token::LBrace handlers: EffectHandler+ Token::RBrace
	| Lambda = Token::Fn Token::LParen args: (Token::Variable (Token::Colon TypeAnnot)? Token::Comma)* Token::RParen expression: Expr
	| Cast = expr: Expr Token::KeywordAs target_type: TypeAnnot
	| IfElse = Token::KeywordIf condition: Expr then: Expr (Token::KeywordElse else_: Expr)?
	| With = Token::KeywordWith context: Expr Token::KeywordAs var: Token::Variable body: Expr
	// ill add `loop` later
	| Match = Token::KeywordMatch Expr Token::LBrace MatchArm* Token::RBrace
	| For = Token::KeywordFor value: Token::Variable Token::KeywordIn iterator: Expr expression: Expr
	| While = Token::KeywordWhile Expr Expr
	| IfLet = Token::KeywordIf Token::KeywordLet pattern: Pattern Token::Assign expr: Expr then: Expr (Token::KeywordElse else: Expr)?
	| WhileLet = Token::KeywordWhile Token::KeywordLet pattern: Pattern Token::Assign expr: Expr body: Expr
	| Return = Token::KeywordReturn Expr?
	| Break = Token::KeywordBreak Expr?
	| Continue = Token::KeywordContinue
	| Call = Expr Token::LParen (Expr Token::Comma)* Token::RParen
	| Index = Expr Token::LBracket Expr Token::RBracket
	| FieldAccess = Expr Token::Dot Token::Variable
	| OptionalChain = Expr Token::OptionalChain Token::Variable
	| MacroCall = Token::MacroInvocation MacroDelimiterStart (Expr Token::Comma)* MacroDelimiterEnd
	| Import = Token::Import path: Token::String Token::RBrace items: (Token::Variable (Token::KeywordAs Token::Variable)?)* Token::KeywordAs Token::Variable
	| Error // just placeholder if needed

Function := 
	| Token::KeywordDef vis: Token::KeywordPub? name: Token::Variable type_params: (Token::Less TypeParam+ Token::Greater)? args: Token::LParen (Token::Variable (Token::Colon TypeAnnot)? Token::Comma)* Token::RParen (Token::Arrow return_type: TypeAnnot)? (Token::Where where_constrains: TypeConstraint*)? (Token::Effects effects: EffectAnnot*)? Expr 

Struct := 
	| Token::KeywordStruct vis: Token::KeywordPub? name: Token::Variable type_params: (Token::Less TypeParam+ Token::Greater)? Token::LBrace fields|methods: ((Token::Variable Token::KeywordPub? Token::Colon TypeAnnot)|Function)* Token::RBrace 

Enum := 
	| Token::KeywordEnum vis: Token::KeywordPub? name: Token::Variable type_params: (Token::Less TypeParam+ Token::Greater)? fields|methods: (EnumVariant|Function)* Token::RBrace 

TypeAlias = 
	| Token::KeywordType name: Token::Variable type_params: (Token::Less TypeParam+ Token::Greater)? Token::KeywordWhere where_constrains: TypeConstraint* Token::Assign target_type: TypeAnnot

Impl = 
	| Token::KeywordImpl name: Token::Variable (Token::Colon trait_: Token::Variable)? type_params: (Token::Less TypeParam+ Token::Greater)? Token::KeywordWhere where_constrains: TypeConstraint* Token::LBrace methods: Function* Token::RBrace

Trait = 
	| Token::KeywordTrait name: Token::Variable (Token::Colon trait_: Token::Variable)? associated_types|methods: (Token::KeywordType TypeAnnot|Function)*
	// todo: add supertraits support. Show: Debug + Clone

// todo: MacroDef 

EffectDef = 
	| Token::KeywordEffect vis: Token::KeywordPub? name: Token::Variable type_params: (Token::Less TypeParam+ Token::Greater)? where_constrains: TypeConstraint* Token::LBrace EffectOperation* Token::RBrace

EffectAnnot =
	| Token::Variable

EffectHandler = 
	| Token::LParen args: Token::Variable* Token::RParen

TypeAnnot = 
	| Token::KeywordBool
	| Token::KeywordInt
	| Token::KeywordFloat
	| Token::KeywordString
	| Named = Token::Variable
	| Unit = Token::LParen Token::RParen
	| Never = Token::Bang
	| Generic = idk
	| Token::Fn Token::LParen args: (TypeAnnot Token::Comma)* Token::RParen Token::Effects EffectAnnot*
	| Tuple = Token::LParen TypeAnnot+ Token::RParen
	| Union = TypeAnnot | (Token::Union TypeAnnot) +
	| Trait = Token::KeywordTrait Token::Variable+
	| Constructor = idk
	| Forall = Token::KeywordForall idk
	| idk the rest

parser.rs

// only the relevant function
fn expr<'src>() -> impl Parser<'src, &'src [Token], Expr, extra::Err<Rich<'src, Token, Span>>> + Clone {
    recursive(|expr| {
        let literal = choice((
            select! { Token::Int(n) => ExprKind::Int(n) },
            select! { Token::Float(f) => ExprKind::Float(f) },
            select! { Token::String(s) => ExprKind::String(s) },
            select! { Token::Bool(b) => ExprKind::Bool(b) },
        ));
        
        let variable = ident().map(ExprKind::Variable);
        
        // Lambda
        let lambda = just(Token::KeywordFn)
            .ignore_then(
                ident()
                    .then(just(Token::Colon).ignore_then(type_annot()).or_not())
                    .map_with(|(name, ty), e| FnArg {
                        span: to_range(e.span()),
                        file: String::new(),
                        name,
                        type_: ty,
                    })
                    .separated_by(just(Token::Comma))
                    .allow_trailing()
                    .collect::<Vec<_>>()
                    .delimited_by(just(Token::LParen), just(Token::RParen))
            )
            .then(expr.clone())
            .map(|(args, body)| ExprKind::Lambda {
                args,
                expression: Box::new(body),
            });
        
        // Block
        let block = expr.clone()
            .then_ignore(just(Token::Semicolon).or_not())
            .repeated()
            .at_least(1)
            .collect::<Vec<_>>()
            .delimited_by(just(Token::LBrace), just(Token::RBrace))
            .recover_with(via_parser(nested_delimiters(
                Token::LBrace, Token::RBrace,
                [(Token::LParen, Token::RParen), (Token::LBracket, Token::RBracket)],
                |_| vec![],
            )))
            .map(ExprKind::Block);
        
        // Let
        let let_expr = just(Token::KeywordLet)
            .ignore_then(ident())
            .then(just(Token::Colon).ignore_then(type_annot()).or_not())
            .then_ignore(just(Token::Assign))
            .then(expr.clone())
            .map(|((var, ty), value)| ExprKind::Let {
                var,
                type_annot: ty,
                value: Box::new(value),
            });
        
        // If-let
        let if_let = just(Token::KeywordIf)
            .ignore_then(just(Token::KeywordLet))
            .ignore_then(pattern())
            .then_ignore(just(Token::Assign))
            .then(expr.clone())
            .then(expr.clone())
            .then(just(Token::KeywordElse).ignore_then(expr.clone()).or_not())
            .map(|(((pattern, expr_val), then), else_)| ExprKind::IfLet {
                pattern,
                expr: Box::new(expr_val),
                then: Box::new(then),
                else_: else_.map(Box::new),
            });
        
        // If-else
        let if_expr = just(Token::KeywordIf)
            .ignore_then(expr.clone())
            .then(expr.clone())
            .then(just(Token::KeywordElse).ignore_then(expr.clone()).or_not())
            .map(|((cond, then_block), else_block)| ExprKind::IfElse {
                condition: Box::new(cond),
                then: Box::new(then_block),
                else_: else_block.map(Box::new),
            });
        
        // While-let
        let while_let = just(Token::KeywordWhile)
            .ignore_then(just(Token::KeywordLet))
            .ignore_then(pattern())
            .then_ignore(just(Token::Assign))
            .then(expr.clone())
            .then(expr.clone())
            .map(|((pattern, expr_val), body)| ExprKind::WhileLet {
                pattern,
                expr: Box::new(expr_val),
                body: Box::new(body),
            });
        
        // While
        let while_expr = just(Token::KeywordWhile)
            .ignore_then(expr.clone())
            .then(expr.clone())
            .map(|(cond, body)| ExprKind::While(Box::new(cond), Box::new(body)));
        
        // Match - CRITICAL: Don't inline match_arm
        let match_arm = pattern()
            .then(just(Token::KeywordIf).ignore_then(expr.clone()).or_not())
            .then_ignore(just(Token::FatArrow))
            .then(expr.clone())
            .map_with(|((pattern, guard), body), e| MatchArm {
                pattern,
                guard,
                body: Box::new(body),
                span: to_range(e.span()),
            });
        
        let match_expr = just(Token::KeywordMatch)
            .ignore_then(expr.clone())
            .then(
                match_arm
                    .separated_by(just(Token::Comma))
                    .allow_trailing()
                    .collect::<Vec<_>>()
                    .delimited_by(just(Token::LBrace), just(Token::RBrace))
            )
            .map(|(scrutinee, arms)| ExprKind::Match(Box::new(scrutinee), arms));
        
        // For
        let for_expr = just(Token::KeywordFor)
            .ignore_then(ident())
            .then_ignore(just(Token::KeywordIn))
            .then(expr.clone())
            .then(expr.clone())
            .map(|((value, iterator), body)| ExprKind::For {
                value,
                iterator: Box::new(iterator),
                expression: Box::new(body),
            });
        
        // Loop
        let loop_expr = just(Token::KeywordLoop)
            .ignore_then(expr.clone())
            .map(|body| ExprKind::Loop {
                label: None,
                body: Box::new(body),
            });
        
        // With
        let with_expr = just(Token::KeywordWith)
            .ignore_then(expr.clone())
            .then_ignore(just(Token::KeywordAs))
            .then(ident())
            .then(expr.clone())
            .map(|((context, var), body)| ExprKind::With {
                context: Box::new(context),
                var,
                body: Box::new(body),
            });
        
        // Return, break, continue
        let return_expr = just(Token::KeywordReturn)
            .then(expr.clone().or_not())
            .map(|(_, e)| ExprKind::Return(e.map(Box::new)));
        
        let break_expr = just(Token::KeywordBreak)
            .then(expr.clone().or_not())
            .map(|(_, e)| ExprKind::Break(e.map(Box::new)));
        
        let continue_expr = just(Token::KeywordContinue).to(ExprKind::Continue);
        
        // Perform
        let perform_expr = just(Token::KeywordPerform)
            .ignore_then(ident())
            .then(
                expr.clone()
                    .separated_by(just(Token::Comma))
                    .allow_trailing()
                    .collect::<Vec<_>>()
                    .delimited_by(just(Token::LParen), just(Token::RParen))
            )
            .map(|(effect, args)| ExprKind::Perform { effect, args });
        
        // Handle
        let effect_handler = ident()
            .then(
                ident()
                    .separated_by(just(Token::Comma))
                    .allow_trailing()
                    .collect::<Vec<_>>()
                    .delimited_by(just(Token::LParen), just(Token::RParen))
            )
            .then_ignore(just(Token::FatArrow))
            .then(expr.clone())
            .map_with(|((effect, mut params), body), e| {
                let resume_param = params.pop().unwrap_or_else(|| "k".to_string());
                EffectHandler {
                    span: to_range(e.span()),
                    effect,
                    params,
                    resume_param,
                    body,
                }
            });
        
        let handle_expr = just(Token::KeywordHandle)
            .ignore_then(expr.clone())
            .then_ignore(just(Token::KeywordWith))
            .then(
                effect_handler
                    .separated_by(just(Token::Comma))
                    .allow_trailing()
                    .collect::<Vec<_>>()
                    .delimited_by(just(Token::LBrace), just(Token::RBrace))
            )
            .map(|(body, handlers)| ExprKind::Handle {
                body: Box::new(body),
                handlers,
            });
        
        // Enum construction
        let enum_construct = ident()
            .then_ignore(just(Token::Access))
            .then(ident())
            .then(
                expr.clone()
                    .separated_by(just(Token::Comma))
                    .allow_trailing()
                    .collect::<Vec<_>>()
                    .delimited_by(just(Token::LParen), just(Token::RParen))
            )
            .map(|((name, variant), args)| ExprKind::EnumConstruct {
                name,
                variant,
                args,
            });
        
        // Import
        let import = just(Token::KeywordImport)
            .ignore_then(select! { Token::String(s) => s })
            .then(
                choice((
                    ident()
                        .then(just(Token::KeywordAs).ignore_then(ident()).or_not())
                        .separated_by(just(Token::Comma))
                        .allow_trailing()
                        .collect::<Vec<_>>()
                        .delimited_by(just(Token::LBrace), just(Token::RBrace))
                        .map(|items| (items, None)),
                    just(Token::KeywordAs)
                        .ignore_then(ident())
                        .map(|alias| (vec![], Some(alias))),
                ))
                .or_not()
            )
            .map_with(|(path, items), e| {
                let (items, alias) = items.unwrap_or((vec![], None));
                let path_parts: Vec<String> = path.split('/').map(|s| s.to_string()).collect();
                ExprKind::Import(Import {
                    span: to_range(e.span()),
                    file: String::new(),
                    path: path_parts,
                    items,
                    alias,
                })
            });
        
        // Macro call
        let macro_call = select! { Token::MacroInvocation(name) => name }
            .then(
                expr.clone()
                    .separated_by(just(Token::Comma))
                    .allow_trailing()
                    .collect::<Vec<_>>()
                    .delimited_by(just(Token::LParen), just(Token::RParen))
            )
            .map(|(name, args)| ExprKind::MacroCall(name, args, Delimiter::Paren));
        
        // Parenthesized expression
        let paren_expr = expr.clone()
            .delimited_by(just(Token::LParen), just(Token::RParen))
            .map(|e| e.expr);
        
        // ATOM - NO COLLECTIONS HERE!
        // Collections will be parsed AFTER as part of the delimiter check
        let atom = choice((
            literal,
            
            // All keywords first
            let_expr,
            if_let,
            if_expr,
            while_let,
            while_expr,
            match_expr,
            for_expr,
            loop_expr,
            with_expr,
            return_expr,
            break_expr,
            continue_expr,
            perform_expr,
            handle_expr,
            
            // Other constructs
            import,
            macro_call,
            enum_construct,
            lambda,
            
            // Block
            block,
            
            // Variable AFTER keywords
            variable,
            
            // Parenthesized last
            paren_expr,
        ))
        .map_with(|kind, e| Expr {
            span: to_range(e.span()),
            file: String::new(),
            expr: kind,
        });
        
        // NOW handle collections by checking delimiters FIRST
        // This breaks the left recursion
        let array_or_expr = just(Token::LBracket)
            .ignore_then(
                expr.clone()
                    .separated_by(just(Token::Comma))
                    .allow_trailing()
                    .collect::<Vec<_>>()
            )
            .then_ignore(just(Token::RBracket))
            .map_with(|exprs, e| Expr {
                span: to_range(e.span()),
                file: String::new(),
                expr: ExprKind::Array(exprs),
            });
        
        let map_or_block = just(Token::LBrace)
            .then(
                choice((
                    // Try map first: expr : expr
                    expr.clone()
                        .then_ignore(just(Token::Colon))
                        .then(expr.clone())
                        .separated_by(just(Token::Comma))
                        .at_least(1)
                        .allow_trailing()
                        .collect::<Vec<_>>()
                        .map(ExprKind::Map),
                    // Otherwise it's a block
                    expr.clone()
                        .then_ignore(just(Token::Semicolon).or_not())
                        .repeated()
                        .at_least(1)
                        .collect::<Vec<_>>()
                        .map(ExprKind::Block),
                ))
            )
            .then_ignore(just(Token::RBrace))
            .map_with(|(_, kind), e| Expr {
                span: to_range(e.span()),
                file: String::new(),
                expr: kind,
            })
            .recover_with(via_parser(nested_delimiters(
                Token::LBrace, Token::RBrace,
                [(Token::LParen, Token::RParen), (Token::LBracket, Token::RBracket)],
                |_| Expr {
                    span: 0..0,
                    file: String::new(),
                    expr: ExprKind::Error,
                },
            )));
        
        let tuple_or_paren = just(Token::LParen)
            .ignore_then(
                expr.clone()
                    .separated_by(just(Token::Comma))
                    .at_least(1)
                    .allow_trailing()
                    .collect::<Vec<_>>()
            )
            .then_ignore(just(Token::RParen))
            .map_with(|exprs, e| {
                if exprs.len() == 1 {
                    exprs.into_iter().next().unwrap()
                } else {
                    Expr {
                        span: to_range(e.span()),
                        file: String::new(),
                        expr: ExprKind::Tuple(exprs),
                    }
                }
            });
        
        // Try delimiter-based parsers first, then atom
        let base = choice((
            array_or_expr,
            map_or_block,
            tuple_or_paren,
            atom,
        )).boxed();
        
        // Postfix operations
        let call_args = expr.clone()
            .separated_by(just(Token::Comma))
            .allow_trailing()
            .collect::<Vec<_>>()
            .delimited_by(just(Token::LParen), just(Token::RParen));
        
        let index = expr.clone()
            .delimited_by(just(Token::LBracket), just(Token::RBracket));
        
        let field = just(Token::Dot).ignore_then(ident());
        let optional_chain = just(Token::OptionalChain).ignore_then(ident());
        let cast = just(Token::KeywordAs).ignore_then(type_annot());
        
        let postfix_op = choice((
            call_args.map(PostOp::Call),
            index.map(|idx| PostOp::Index(Box::new(idx))),
            optional_chain.map(PostOp::OptionalChain),
            field.map(PostOp::Field),
            cast.map(PostOp::Cast),
        ));
        
        let postfix = base.foldl(postfix_op.repeated(), |lhs, op| {
            let span = lhs.span.clone();
            Expr {
                span,
                file: String::new(),
                expr: match op {
                    PostOp::Call(args) => ExprKind::Call(Box::new(lhs), args),
                    PostOp::Index(idx) => ExprKind::Index(Box::new(lhs), idx),
                    PostOp::Field(f) => ExprKind::FieldAccess(Box::new(lhs), f),
                    PostOp::OptionalChain(f) => ExprKind::OptionalChain(Box::new(lhs), f),
                    PostOp::Cast(ty) => ExprKind::Cast { expr: Box::new(lhs), target_type: ty },
                },
            }
        });
        
        // Binary/unary operators
        postfix.pratt((
            infix(right(1), just(Token::Pipe), |l: Expr, _op, r: Expr, _extra| {
                let span = l.span.start..r.span.end;
                Expr { span, file: String::new(), expr: ExprKind::BinOp(Box::new(l), BinOp::Pipe, Box::new(r)) }
            }),
            infix(right(2), just(Token::Assign), |l: Expr, _op, r: Expr, _extra| {
                let span = l.span.start..r.span.end;
                Expr { span, file: String::new(), expr: ExprKind::Assign {
                    l_val: Box::new(l),
                    r_val: Box::new(r),
                    op: AssignOp::Assign,
                }}
            }),
            infix(right(2), just(Token::AddAssign), |l: Expr, _op, r: Expr, _extra| {
                let span = l.span.start..r.span.end;
                Expr { span, file: String::new(), expr: ExprKind::Assign {
                    l_val: Box::new(l),
                    r_val: Box::new(r),
                    op: AssignOp::AddAssign,
                }}
            }),
            infix(right(2), just(Token::SubAssign), |l: Expr, _op, r: Expr, _extra| {
                let span = l.span.start..r.span.end;
                Expr { span, file: String::new(), expr: ExprKind::Assign {
                    l_val: Box::new(l),
                    r_val: Box::new(r),
                    op: AssignOp::SubAssign,
                }}
            }),
            infix(right(2), just(Token::MulAssign), |l: Expr, _op, r: Expr, _extra| {
                let span = l.span.start..r.span.end;
                Expr { span, file: String::new(), expr: ExprKind::Assign {
                    l_val: Box::new(l),
                    r_val: Box::new(r),
                    op: AssignOp::MulAssign,
                }}
            }),
            infix(right(2), just(Token::DivAssign), |l: Expr, _op, r: Expr, _extra| {
                let span = l.span.start..r.span.end;
                Expr { span, file: String::new(), expr: ExprKind::Assign {
                    l_val: Box::new(l),
                    r_val: Box::new(r),
                    op: AssignOp::DivAssign,
                }}
            }),
            infix(right(2), just(Token::ModAssign), |l: Expr, _op, r: Expr, _extra| {
                let span = l.span.start..r.span.end;
                Expr { span, file: String::new(), expr: ExprKind::Assign {
                    l_val: Box::new(l),
                    r_val: Box::new(r),
                    op: AssignOp::ModAssign,
                }}
            }),
            infix(left(3), just(Token::Or), |l: Expr, _op, r: Expr, _extra| {
                let span = l.span.start..r.span.end;
                Expr { span, file: String::new(), expr: ExprKind::BinOp(Box::new(l), BinOp::Or, Box::new(r)) }
            }),
            infix(left(3), just(Token::Nor), |l: Expr, _op, r: Expr, _extra| {
                let span = l.span.start..r.span.end;
                Expr { span, file: String::new(), expr: ExprKind::BinOp(Box::new(l), BinOp::Nor, Box::new(r)) }
            }),
            infix(left(5), just(Token::Xor), |l: Expr, _op, r: Expr, _extra| {
                let span = l.span.start..r.span.end;
                Expr { span, file: String::new(), expr: ExprKind::BinOp(Box::new(l), BinOp::Xor, Box::new(r)) }
            }),
            infix(left(7), just(Token::And), |l: Expr, _op, r: Expr, _extra| {
                let span = l.span.start..r.span.end;
                Expr { span, file: String::new(), expr: ExprKind::BinOp(Box::new(l), BinOp::And, Box::new(r)) }
            }),
            infix(left(9), just(Token::Eq), |l: Expr, _op, r: Expr, _extra| {
                let span = l.span.start..r.span.end;
                Expr { span, file: String::new(), expr: ExprKind::BinOp(Box::new(l), BinOp::Eq, Box::new(r)) }
            }),
            infix(left(9), just(Token::NotEq), |l: Expr, _op, r: Expr, _extra| {
                let span = l.span.start..r.span.end;
                Expr { span, file: String::new(), expr: ExprKind::BinOp(Box::new(l), BinOp::NotEq, Box::new(r)) }
            }),
            infix(left(9), just(Token::Less), |l: Expr, _op, r: Expr, _extra| {
                let span = l.span.start..r.span.end;
                Expr { span, file: String::new(), expr: ExprKind::BinOp(Box::new(l), BinOp::Less, Box::new(r)) }
            }),
            infix(left(9), just(Token::Greater), |l: Expr, _op, r: Expr, _extra| {
                let span = l.span.start..r.span.end;
                Expr { span, file: String::new(), expr: ExprKind::BinOp(Box::new(l), BinOp::Greater, Box::new(r)) }
            }),
            infix(left(9), just(Token::LessEq), |l: Expr, _op, r: Expr, _extra| {
                let span = l.span.start..r.span.end;
                Expr { span, file: String::new(), expr: ExprKind::BinOp(Box::new(l), BinOp::LessEq, Box::new(r)) }
            }),
            infix(left(9), just(Token::GreaterEq), |l: Expr, _op, r: Expr, _extra| {
                let span = l.span.start..r.span.end;
                Expr { span, file: String::new(), expr: ExprKind::BinOp(Box::new(l), BinOp::GreaterEq, Box::new(r)) }
            }),
            infix(left(11), just(Token::Plus), |l: Expr, _op, r: Expr, _extra| {
                let span = l.span.start..r.span.end;
                Expr { span, file: String::new(), expr: ExprKind::BinOp(Box::new(l), BinOp::Add, Box::new(r)) }
            }),
            infix(left(11), just(Token::Minus), |l: Expr, _op, r: Expr, _extra| {
                let span = l.span.start..r.span.end;
                Expr { span, file: String::new(), expr: ExprKind::BinOp(Box::new(l), BinOp::Sub, Box::new(r)) }
            }),
            infix(left(13), just(Token::Mul), |l: Expr, _op, r: Expr, _extra| {
                let span = l.span.start..r.span.end;
                Expr { span, file: String::new(), expr: ExprKind::BinOp(Box::new(l), BinOp::Mul, Box::new(r)) }
            }),
            infix(left(13), just(Token::Div), |l: Expr, _op, r: Expr, _extra| {
                let span = l.span.start..r.span.end;
                Expr { span, file: String::new(), expr: ExprKind::BinOp(Box::new(l), BinOp::Div, Box::new(r)) }
            }),
            infix(left(13), just(Token::Mod), |l: Expr, _op, r: Expr, _extra| {
                let span = l.span.start..r.span.end;
                Expr { span, file: String::new(), expr: ExprKind::BinOp(Box::new(l), BinOp::Mod, Box::new(r)) }
            }),
            infix(right(15), just(Token::Power), |l: Expr, _op, r: Expr, _extra| {
                let span = l.span.start..r.span.end;
                Expr { span, file: String::new(), expr: ExprKind::BinOp(Box::new(l), BinOp::Pow, Box::new(r)) }
            }),
            prefix(17, just(Token::Minus), |_op, e: Expr, _extra| {
                Expr { span: e.span.clone(), file: String::new(), expr: ExprKind::UnOp(UnOp::Minus, Box::new(e)) }
            }),
            prefix(17, just(Token::Not), |_op, e: Expr, _extra| {
                Expr { span: e.span.clone(), file: String::new(), expr: ExprKind::UnOp(UnOp::Not, Box::new(e)) }
            }),
            prefix(17, just(Token::Plus), |_op, e: Expr, _extra| {
                Expr { span: e.span.clone(), file: String::new(), expr: ExprKind::UnOp(UnOp::Plus, Box::new(e)) }
            }),
        ))
        .boxed()
    })
}

fn type_annot<'src>() -> impl Parser<'src, &'src [Token], TypeAnnot, extra::Err<Rich<'src, Token, Span>>> + Clone {
    recursive(|ty| {
        // Primitive types (lowercase)
        let primitive = choice((
            just(Token::KeywordInt).to(TypeAnnotKind::Int),    
            just(Token::KeywordFloat).to(TypeAnnotKind::Float),
            just(Token::KeywordBool).to(TypeAnnotKind::Bool),  
            just(Token::KeywordString).to(TypeAnnotKind::String),
        ));
        
        // Unit type: ()
        let unit = just(Token::LParen)
            .then(just(Token::RParen))
            .to(TypeAnnotKind::Unit);
        
        // Never type: !
        let never = just(Token::Bang).to(TypeAnnotKind::Never);
        
        // Named type
        let named = ident().map(TypeAnnotKind::Named);
        
        // Generic type: Name<T, U>
        let generic = ident()
            .then(
                ty.clone()
                    .separated_by(just(Token::Comma))
                    .at_least(1)
                    .collect::<Vec<_>>()
                    .delimited_by(just(Token::Less), just(Token::Greater))
            )
            .map(|(name, args)| TypeAnnotKind::Generic {
                name,
                args,
                kind: None,
            });
        
        // Tuple type: (T, U, V)
        let tuple = ty.clone()
            .separated_by(just(Token::Comma))
            .at_least(2)
            .collect::<Vec<_>>()
            .delimited_by(just(Token::LParen), just(Token::RParen))
            .map(TypeAnnotKind::Tuple);
        
        // Function type: fn(T, U) -> V effects E
        let function = just(Token::KeywordFn)
            .ignore_then(
                ty.clone()
                    .separated_by(just(Token::Comma))
                    .allow_trailing()
                    .collect::<Vec<_>>()
                    .delimited_by(just(Token::LParen), just(Token::RParen))
            )
            .then_ignore(just(Token::Arrow))
            .then(ty.clone())
            .then(
                select! { Token::Variable(s) if s == "effects" => () }
                    .ignore_then(
                        ident()
                            .separated_by(just(Token::Comma))
                            .collect::<Vec<_>>()
                    )
                    .map(EffectAnnot::closed)
                    .or_not()
            )
            .map(|((params, ret), effects)| TypeAnnotKind::Function {
                params,
                return_type: Box::new(ret),
                effects: effects.unwrap_or_else(EffectAnnot::pure),
            });
        
        // Union type: T | U | V
        let union = ty.clone()
            .separated_by(just(Token::Union))
            .at_least(2)
            .collect::<Vec<_>>()
            .map(TypeAnnotKind::Union);
        
        choice((function, union, tuple, generic, primitive, unit, never, named))
            .map_with(|kind, e| TypeAnnot {
                span: to_range(e.span()),
                file: String::new(),
                expr: kind,
            })
            .labelled("type")
    })
}