suicmez/output.md
2025-12-15 13:53:57 +05:30

89 KiB

// src/lib.rs
pub const EXTENSION: &str = ".sui";

pub mod ast;
pub mod lexer;
pub mod parser;
pub mod typechecker;

// src/parser.rs
use crate::ast::*;
use crate::lexer::Token;

use std::iter::Peekable;
use std::ops::Range;
use std::vec::IntoIter;

type TokenIter = Peekable<IntoIter<(Token, Range<usize>)>>;

pub struct Parser {
    pub file: String,
    pub tokens: TokenIter,
}

#[derive(Debug)]
pub struct ParseError {
    pub message: String,
    pub span: Span,
}

impl Parser {
    pub fn new(file: String, tokens: Vec<(Token, Range<usize>)>) -> Self {
        Parser {
            file,
            tokens: tokens.into_iter().peekable(),
        }
    }

    // Parse the entire file into a list of AST nodes
    pub fn parse(&mut self) -> Result<Vec<ASTNode>, ParseError> {
        let mut nodes = Vec::new();

        while self.peek().is_some() {
            nodes.push(self.parse_top_level()?);
        }

        Ok(nodes)
    }

    fn peek(&mut self) -> Option<&Token> {
        self.tokens.peek().map(|(token, _)| token)
    }

    fn peek_span(&mut self) -> Option<Range<usize>> {
        self.tokens.peek().map(|(_, span)| span.clone())
    }

    fn next(&mut self) -> Option<(Token, Range<usize>)> {
        self.tokens.next()
    }

    fn expect(&mut self, expected: Token) -> Result<Range<usize>, ParseError> {
        match self.next() {
            Some((token, span))
                if std::mem::discriminant(&token) == std::mem::discriminant(&expected) =>
            {
                Ok(span)
            }
            Some((token, span)) => Err(ParseError {
                message: format!("Expected {:?}, found {:?}", expected, token),
                span: Span::new(&span, self.file.clone()),
            }),
            None => Err(ParseError {
                message: format!("Expected {:?}, found EOF", expected),
                span: Span::new(&(0..0), self.file.clone()),
            }),
        }
    }

    fn error<T>(&self, msg: String, span: Range<usize>) -> Result<T, ParseError> {
        Err(ParseError {
            message: msg,
            span: Span::new(&span, self.file.clone()),
        })
    }

    fn parse_top_level(&mut self) -> Result<ASTNode, ParseError> {
        let mut attributes = Vec::new();

        // Parse any leading attributes
        while matches!(self.peek(), Some(Token::At)) {
            attributes.push(self.parse_attribute()?);
        }

        let start = self.peek_span().unwrap_or(0..0).start;
        let token = self.peek().cloned();
        match token {
            Some(Token::KeywordUse) => {
                self.next();
                let path = match self.next() {
                    Some((Token::String(s), _)) => s,
                    Some((_, span)) => {
                        return self.error("Expected string after 'use'".to_string(), span);
                    }
                    None => {
                        return self.error("Expected string after 'use'".to_string(), start..start);
                    }
                };
                let end = self.peek_span().unwrap_or(start..start).end;
                Ok(ASTNode {
                    kind: ASTNodeKind::Use(path),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes,
                })
            }
            Some(Token::KeywordFn) => {
                self.next();
                let func = self.parse_function()?;
                let end = self.peek_span().unwrap_or(start..start).end;
                Ok(ASTNode {
                    kind: ASTNodeKind::Function(func),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes,
                })
            }
            Some(Token::KeywordStruct) => {
                self.next();
                let struct_def = self.parse_struct()?;
                let end = self.peek_span().unwrap_or(start..start).end;
                Ok(ASTNode {
                    kind: ASTNodeKind::Struct(struct_def),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes,
                })
            }
            Some(Token::KeywordEnum) => {
                self.next();
                let enum_def = self.parse_enum()?;
                let end = self.peek_span().unwrap_or(start..start).end;
                Ok(ASTNode {
                    kind: ASTNodeKind::Enum(enum_def),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes,
                })
            }
            Some(Token::KeywordImpl) => {
                self.next();
                let impl_def = self.parse_impl()?;
                let end = self.peek_span().unwrap_or(start..start).end;
                Ok(ASTNode {
                    kind: ASTNodeKind::Impl(impl_def),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes,
                })
            }
            Some(Token::KeywordTrait) => {
                self.next();
                let trait_def = self.parse_trait()?;
                let end = self.peek_span().unwrap_or(start..start).end;
                Ok(ASTNode {
                    kind: ASTNodeKind::Trait(trait_def),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes,
                })
            }
            Some(Token::KeywordExtern) => {
                self.next();
                let extern_def = self.parse_extern()?;
                let end = self.peek_span().unwrap_or(start..start).end;
                Ok(ASTNode {
                    kind: ASTNodeKind::Extern(extern_def),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes,
                })
            }
            Some(Token::KeywordLoad) => {
                self.next();
                let load_def = self.parse_load()?;
                let end = self.peek_span().unwrap_or(start..start).end;
                Ok(ASTNode {
                    kind: ASTNodeKind::Load(load_def),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes,
                })
            }
            Some(token) => {
                let span = self.peek_span().unwrap_or(start..start);
                self.error(format!("Unexpected token at top level: {:?}", token), span)
            }
            None => self.error("Unexpected EOF".to_string(), start..start),
        }
    }

    fn parse_attribute(&mut self) -> Result<Attribute, ParseError> {
        self.expect(Token::At)?;
        let start = self.peek_span().unwrap_or(0..0).start;
        let name = match self.next() {
            Some((Token::Variable(name), _)) => name,
            Some((_, span)) => return self.error("Expected attribute name".to_string(), span),
            None => return self.error("Expected attribute name".to_string(), start..start),
        };

        // Parentheses are optional
        let mut args = vec![];
        if matches!(self.peek(), Some(Token::LParen)) {
            self.next();
            loop {
                let token = self.peek().cloned();
                match token {
                    Some(Token::RParen) => {
                        self.next();
                        break;
                    }
                    Some(Token::String(s)) => {
                        self.next();
                        args.push(AttributeArg::Literal(s));
                    }
                    Some(Token::Variable(id)) => {
                        self.next();
                        let next_token = self.peek().cloned();
                        if matches!(next_token, Some(Token::Assign)) {
                            self.next();
                            match self.next() {
                                Some((Token::Variable(val), _)) => {
                                    args.push(AttributeArg::KeyValue(id, val))
                                }
                                Some((_, span)) => {
                                    return self.error("Expected value after =".to_string(), span);
                                }
                                None => {
                                    return self
                                        .error("Expected value after =".to_string(), start..start);
                                }
                            }
                        } else {
                            args.push(AttributeArg::Value(id));
                        }
                    }
                    Some(token) => {
                        let span = self.peek_span().unwrap_or(start..start);
                        return self
                            .error(format!("Unexpected token in attribute: {:?}", token), span);
                    }
                    None => {
                        return self.error("Expected attribute argument".to_string(), start..start);
                    }
                }
                let next_token = self.peek().cloned();
                if matches!(next_token, Some(Token::Comma)) {
                    self.next();
                } else if matches!(next_token, Some(Token::RParen)) {
                    // ok
                } else {
                    {
                        let span = self.peek_span().unwrap_or(start..start);
                        return self.error("Expected , or )".to_string(), span);
                    }
                }
            }
        }
        let end = self.peek_span().unwrap_or(start..start).end;
        Ok(Attribute {
            name,
            args,
            span: Span::new(&(start..end), self.file.clone()),
        })
    }

    fn parse_function(&mut self) -> Result<Function, ParseError> {
        let start = self.peek_span().unwrap_or(0..0).start;
        let name = match self.next() {
            Some((Token::Variable(n), _)) => n,
            Some((_, span)) => return self.error("Expected function name".to_string(), span),
            None => return self.error("Expected function name".to_string(), start..start),
        };

        // Parse type parameters if present
        let parameters = if matches!(self.peek(), Some(Token::Less)) {
            self.next();
            self.parse_parameters()?
        } else {
            Vec::new()
        };

        // Parse function arguments
        self.expect(Token::LParen)?;
        let mut args = Vec::new();
        loop {
            if matches!(self.peek(), Some(Token::RParen)) {
                self.next();
                break;
            }

            let arg_name = match self.next() {
                Some((Token::Variable(n), _)) => n,
                Some((_, span)) => return self.error("Expected argument name".to_string(), span),
                None => return self.error("Expected argument name".to_string(), start..start),
            };

            let arg_type = if matches!(self.peek(), Some(Token::Colon)) {
                self.next();
                Some(self.parse_type_annot()?)
            } else {
                None
            };

            args.push((arg_name, arg_type));

            if matches!(self.peek(), Some(Token::Comma)) {
                self.next();
            } else if !matches!(self.peek(), Some(Token::RParen)) {
                let span = self.peek_span().unwrap_or(start..start);
                return self.error("Expected , or )".to_string(), span);
            }
        }

        // Parse return type if present
        let return_type = if matches!(self.peek(), Some(Token::Arrow)) {
            self.next();
            Some(self.parse_type_annot()?)
        } else {
            None
        };

        // Parse body expression
        let body = self.parse_expr()?;

        Ok(Function {
            name,
            parameters,
            args,
            return_type,
            body,
        })
    }

    fn parse_struct(&mut self) -> Result<Struct, ParseError> {
        let start = self.peek_span().unwrap_or(0..0).start;
        let name = match self.next() {
            Some((Token::Variable(n), _)) => n,
            Some((_, span)) => return self.error("Expected struct name".to_string(), span),
            None => return self.error("Expected struct name".to_string(), start..start),
        };

        // Parse type parameters if present
        let parameters = if matches!(self.peek(), Some(Token::Less)) {
            self.next();
            self.parse_parameters()?
        } else {
            Vec::new()
        };

        // Parse fields
        let mut fields = Vec::new();
        loop {
            if matches!(self.peek(), Some(Token::KeywordEnd)) {
                self.next();
                break;
            }

            let field_start = self.peek_span().unwrap_or(0..0).start;
            let field_name = match self.next() {
                Some((Token::Variable(n), _)) => n,
                Some((_, span)) => return self.error("Expected field name".to_string(), span),
                None => return self.error("Expected field name".to_string(), start..start),
            };

            self.expect(Token::Colon)?;
            let field_type = self.parse_type_annot()?;
            let field_end = self.peek_span().unwrap_or(field_start..field_start).start;

            fields.push(Field {
                name: field_name,
                field_type,
                span: Span::new(&(field_start..field_end), self.file.clone()),
            });

            if matches!(self.peek(), Some(Token::Comma)) {
                self.next();
            }
        }

        Ok(Struct {
            name,
            parameters,
            fields,
        })
    }

    fn parse_enum(&mut self) -> Result<Enum, ParseError> {
        let start = self.peek_span().unwrap_or(0..0).start;
        let name = match self.next() {
            Some((Token::Variable(n), _)) => n,
            Some((_, span)) => return self.error("Expected enum name".to_string(), span),
            None => return self.error("Expected enum name".to_string(), start..start),
        };

        // Parse type parameters if present
        let parameters = if matches!(self.peek(), Some(Token::Less)) {
            self.next();
            self.parse_parameters()?
        } else {
            Vec::new()
        };

        // Parse variants
        let mut variants = Vec::new();
        loop {
            if matches!(self.peek(), Some(Token::KeywordEnd)) {
                self.next();
                break;
            }

            let variant_start = self.peek_span().unwrap_or(0..0).start;
            let variant_name = match self.next() {
                Some((Token::Variable(n), _)) => n,
                Some((_, span)) => return self.error("Expected variant name".to_string(), span),
                None => return self.error("Expected variant name".to_string(), start..start),
            };

            let mut fields = Vec::new();
            if matches!(self.peek(), Some(Token::LParen)) {
                self.next();
                loop {
                    if matches!(self.peek(), Some(Token::RParen)) {
                        self.next();
                        break;
                    }
                    fields.push(self.parse_type_annot()?);

                    if matches!(self.peek(), Some(Token::Comma)) {
                        self.next();
                    }
                }
            }

            let variant_end = self
                .peek_span()
                .unwrap_or(variant_start..variant_start)
                .start;
            variants.push(Variant {
                name: variant_name,
                fields,
                span: Span::new(&(variant_start..variant_end), self.file.clone()),
            });

            if matches!(self.peek(), Some(Token::Comma)) {
                self.next();
            }
        }

        Ok(Enum {
            name,
            parameters,
            variants,
        })
    }

    fn parse_impl(&mut self) -> Result<Impl, ParseError> {
        let start = self.peek_span().unwrap_or(0..0).start;

        // Parse impl target as a type (could be generic like Option<int>)
        let target_type = self.parse_type_annot()?;

        // Extract the base type name from the type annotation
        let target = match target_type {
            TypeAnnot::Var(name) => name,
            TypeAnnot::Cons(name, _) => name,
            _ => {
                return self.error(
                    "Expected type name for impl target".to_string(),
                    start..start,
                );
            }
        };

        // Parse optional trait name
        let trait_name = if matches!(self.peek(), Some(Token::Colon)) {
            self.next();
            match self.next() {
                Some((Token::Variable(n), _)) => Some(n),
                Some((_, span)) => return self.error("Expected trait name".to_string(), span),
                None => return self.error("Expected trait name".to_string(), start..start),
            }
        } else {
            None
        };

        // Parse methods
        let mut methods = Vec::new();
        loop {
            if matches!(self.peek(), Some(Token::KeywordEnd)) {
                self.next();
                break;
            }

            self.expect(Token::KeywordFn)?;
            methods.push(self.parse_function()?);
        }

        Ok(Impl {
            target,
            trait_name,
            methods,
        })
    }

    fn parse_trait(&mut self) -> Result<Trait, ParseError> {
        let start = self.peek_span().unwrap_or(0..0).start;
        let name = match self.next() {
            Some((Token::Variable(n), _)) => n,
            Some((_, span)) => return self.error("Expected trait name".to_string(), span),
            None => return self.error("Expected trait name".to_string(), start..start),
        };

        // Parse type parameters if present
        let parameters = if matches!(self.peek(), Some(Token::Less)) {
            self.next();
            self.parse_parameters()?
        } else {
            Vec::new()
        };

        // Parse methods
        let mut methods = Vec::new();

        loop {
            if matches!(self.peek(), Some(Token::KeywordEnd)) {
                self.next();
                break;
            }

            if matches!(self.peek(), Some(Token::KeywordFn)) {
                self.next();
                methods.push(self.parse_function_signature()?);

                // Optional comma between methods
                if matches!(self.peek(), Some(Token::Comma)) {
                    self.next();
                }
            } else {
                break;
            }
        }

        Ok(Trait {
            name,
            methods,
            parameters,
            associated_types: Vec::new(),
        })
    }

    fn parse_extern(&mut self) -> Result<Extern, ParseError> {
        let start = self.peek_span().unwrap_or(0..0).start;
        let name = match self.next() {
            Some((Token::Variable(n), _)) => n,
            Some((_, span)) => return self.error("Expected extern name".to_string(), span),
            None => return self.error("Expected extern name".to_string(), start..start),
        };

        // Parse argument types (with optional parameter names)
        self.expect(Token::LParen)?;
        let mut args = Vec::new();
        loop {
            if matches!(self.peek(), Some(Token::RParen)) {
                self.next();
                break;
            }

            args.push(self.parse_type_annot()?);

            if matches!(self.peek(), Some(Token::Comma)) {
                self.next();
            }
        }

        // Parse return type
        self.expect(Token::Arrow)?;
        let return_type = self.parse_type_annot()?;

        // Parse from clause
        self.expect(Token::KeywordFrom)?;
        let from = match self.next() {
            Some((Token::Variable(n), _)) => n,
            Some((_, span)) => return self.error("Expected library identifier".to_string(), span),
            None => return self.error("Expected library identifier".to_string(), start..start),
        };

        let end = self.peek_span().unwrap_or(start..start).end;
        Ok(Extern {
            name,
            args,
            return_type,
            from,
            span: Span::new(&(start..end), self.file.clone()),
        })
    }

    fn parse_load(&mut self) -> Result<Load, ParseError> {
        let start = self.peek_span().unwrap_or(0..0).start;
        let library = match self.next() {
            Some((Token::String(s), _)) => s,
            Some((_, span)) => return self.error("Expected library name".to_string(), span),
            None => return self.error("Expected library name".to_string(), start..start),
        };

        self.expect(Token::KeywordAs)?;
        let alias = match self.next() {
            Some((Token::Variable(a), _)) => a,
            Some((_, span)) => return self.error("Expected alias".to_string(), span),
            None => return self.error("Expected alias".to_string(), start..start),
        };

        let end = self.peek_span().unwrap_or(start..start).end;
        Ok(Load {
            library,
            alias,
            span: Span::new(&(start..end), self.file.clone()),
        })
    }

    fn parse_parameters(&mut self) -> Result<Vec<Parameter>, ParseError> {
        let start = self.peek_span().unwrap_or(0..0).start;
        let mut params = Vec::new();

        loop {
            if matches!(self.peek(), Some(Token::Greater)) {
                self.next();
                break;
            }

            let param_start = self.peek_span().unwrap_or(0..0).start;
            let param_name = match self.next() {
                Some((Token::Variable(n), _)) => n,
                Some((_, span)) => return self.error("Expected parameter name".to_string(), span),
                None => return self.error("Expected parameter name".to_string(), start..start),
            };

            let bounds = if matches!(self.peek(), Some(Token::Colon)) {
                self.next();
                self.parse_trait_bounds()?
            } else {
                Vec::new()
            };

            let kind = if matches!(self.peek(), Some(Token::Colon)) {
                self.next();
                Some(self.parse_kind()?)
            } else {
                None
            };

            let param_end = self.peek_span().unwrap_or(param_start..param_start).end;
            params.push(Parameter {
                name: param_name,
                bounds,
                kind,
                span: Span::new(&(param_start..param_end), self.file.clone()),
            });

            if matches!(self.peek(), Some(Token::Comma)) {
                self.next();
            }
        }

        Ok(params)
    }

    fn parse_trait_bounds(&mut self) -> Result<Vec<String>, ParseError> {
        let mut bounds = Vec::new();
        loop {
            match self.next() {
                Some((Token::Variable(n), _)) => bounds.push(n),
                Some((_, span)) => return self.error("Expected trait name".to_string(), span),
                None => return self.error("Expected trait name".to_string(), 0..0),
            }

            if !matches!(self.peek(), Some(Token::Plus)) {
                break;
            }
            self.next();
        }

        Ok(bounds)
    }

    fn parse_kind(&mut self) -> Result<Kind, ParseError> {
        if matches!(self.peek(), Some(Token::Mul)) {
            self.next();
            Ok(Kind::Star)
        } else {
            let k1 = Box::new(self.parse_kind()?);
            self.expect(Token::Arrow)?;
            let k2 = Box::new(self.parse_kind()?);
            Ok(Kind::Arrow(k1, k2))
        }
    }

    fn parse_type_annot(&mut self) -> Result<TypeAnnot, ParseError> {
        let start = self.peek_span().unwrap_or(0..0).start;

        // Check for function type: fn (args)->ret
        if matches!(self.peek(), Some(Token::KeywordFn)) {
            self.next();
            self.expect(Token::LParen)?;
            let mut arg_types = Vec::new();

            loop {
                if matches!(self.peek(), Some(Token::RParen)) {
                    self.next();
                    break;
                }
                arg_types.push(self.parse_type_annot()?);

                if matches!(self.peek(), Some(Token::Comma)) {
                    self.next();
                }
            }

            self.expect(Token::Arrow)?;
            let ret_type = Box::new(self.parse_type_annot()?);
            return Ok(TypeAnnot::Function(arg_types, ret_type));
        }

        let mut base_type = match self.next() {
            Some((Token::Variable(n), _)) => TypeAnnot::Cons(n, vec![]),
            Some((Token::KeywordBool, _)) => TypeAnnot::Cons("bool".to_string(), vec![]),
            Some((Token::KeywordInt, _)) => TypeAnnot::Cons("int".to_string(), vec![]),
            Some((Token::KeywordFloat, _)) => TypeAnnot::Cons("float".to_string(), vec![]),
            Some((Token::KeywordString, _)) => TypeAnnot::Cons("string".to_string(), vec![]),
            Some((Token::LParen, _)) => {
                // Check for unit type: ()
                if matches!(self.peek(), Some(Token::RParen)) {
                    self.next();
                    return Ok(TypeAnnot::Cons("unit".to_string(), vec![]));
                }

                let mut types = Vec::new();
                loop {
                    if matches!(self.peek(), Some(Token::RParen)) {
                        self.next();
                        break;
                    }
                    types.push(self.parse_type_annot()?);

                    if matches!(self.peek(), Some(Token::Comma)) {
                        self.next();
                    }
                }

                // Single element in parens is not a tuple, unwrap it
                if types.len() == 1 {
                    types.pop().unwrap()
                } else {
                    TypeAnnot::Tuple(types)
                }
            }
            Some((Token::LBracket, _)) => {
                let inner = self.parse_type_annot()?;
                self.expect(Token::RBracket)?;
                TypeAnnot::Array(Box::new(inner))
            }
            Some((Token::Bang, _)) => TypeAnnot::Cons("never".to_string(), vec![]),
            Some((_, span)) => return self.error("Expected type name".to_string(), span),
            None => return self.error("Expected type name".to_string(), start..start),
        };

        // Parse type arguments if present
        if matches!(self.peek(), Some(Token::Less)) {
            self.next();
            let mut args = Vec::new();
            loop {
                if matches!(self.peek(), Some(Token::Greater)) {
                    self.next();
                    break;
                }
                args.push(self.parse_type_annot()?);

                if matches!(self.peek(), Some(Token::Comma)) {
                    self.next();
                }
            }
            base_type = match base_type {
                TypeAnnot::Cons(name, _) => TypeAnnot::Cons(name, args),
                _ => {
                    return self.error("Expected type name for generic".to_string(), start..start);
                }
            };
        }

        // Parse array types
        while matches!(self.peek(), Some(Token::LBracket)) {
            self.next();
            self.expect(Token::RBracket)?;
            base_type = TypeAnnot::Array(Box::new(base_type));
        }

        Ok(base_type)
    }

    fn parse_function_signature(&mut self) -> Result<FunctionSignature, ParseError> {
        let start = self.peek_span().unwrap_or(0..0).start;
        let name = match self.next() {
            Some((Token::Variable(n), _)) => n,
            Some((_, span)) => return self.error("Expected function name".to_string(), span),
            None => return self.error("Expected function name".to_string(), start..start),
        };

        self.expect(Token::LParen)?;
        let mut params = Vec::new();
        loop {
            if matches!(self.peek(), Some(Token::RParen)) {
                self.next();
                break;
            }

            let param_start = self.peek_span().unwrap_or(0..0).start;
            let param_name = match self.next() {
                Some((Token::Variable(n), _)) => n,
                Some((_, span)) => return self.error("Expected parameter name".to_string(), span),
                None => return self.error("Expected parameter name".to_string(), start..start),
            };

            // Parameters in trait methods may have type annotations
            if matches!(self.peek(), Some(Token::Colon)) {
                self.next();
                let _param_type = self.parse_type_annot()?;
            }

            let param_end = self.peek_span().unwrap_or(param_start..param_start).end;
            params.push(Parameter {
                name: param_name,
                bounds: Vec::new(),
                kind: None,
                span: Span::new(&(param_start..param_end), self.file.clone()),
            });

            if matches!(self.peek(), Some(Token::Comma)) {
                self.next();
            }
        }

        self.expect(Token::Arrow)?;
        let return_type = self.parse_type_annot()?;

        Ok(FunctionSignature {
            name,
            params,
            return_type,
        })
    }

    fn parse_expr(&mut self) -> Result<Expr, ParseError> {
        let mut attributes = Vec::new();

        // Parse any leading attributes
        while matches!(self.peek(), Some(Token::At)) {
            attributes.push(self.parse_attribute()?);
        }

        let mut expr = self.parse_assignment()?;
        expr.attributes = attributes;
        Ok(expr)
    }

    fn parse_range_expr(&mut self) -> Result<Expr, ParseError> {
        let left = self.parse_or_expr()?;

        if matches!(self.peek(), Some(Token::DotDot)) {
            let start = left.span.start;
            self.next();
            let right = self.parse_or_expr()?;
            let end = right.span.end;
            Ok(Expr {
                kind: ExprKind::Range(Box::new(left), Box::new(right)),
                span: Span::new(&(start..end), self.file.clone()),
                attributes: Vec::new(),
            })
        } else {
            Ok(left)
        }
    }

    fn parse_assignment(&mut self) -> Result<Expr, ParseError> {
        let left = self.parse_range_expr()?;

        if matches!(self.peek(), Some(Token::Assign)) {
            let start = left.span.start;
            self.next();
            let right = self.parse_assignment()?;
            let end = right.span.end;
            Ok(Expr {
                kind: ExprKind::Assign(Box::new(left), Box::new(right)),
                span: Span::new(&(start..end), self.file.clone()),
                attributes: Vec::new(),
            })
        } else {
            Ok(left)
        }
    }

    fn parse_or_expr(&mut self) -> Result<Expr, ParseError> {
        let mut left = self.parse_and_expr()?;

        loop {
            if matches!(self.peek(), Some(Token::Or)) {
                let start = left.span.start;
                self.next();
                let right = self.parse_and_expr()?;
                let end = right.span.end;
                left = Expr {
                    kind: ExprKind::BinOp(Box::new(left), BinOp::Or, Box::new(right)),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes: Vec::new(),
                };
            } else {
                break;
            }
        }

        Ok(left)
    }

    fn parse_and_expr(&mut self) -> Result<Expr, ParseError> {
        let mut left = self.parse_eq_expr()?;

        loop {
            if matches!(self.peek(), Some(Token::And)) {
                let start = left.span.start;
                self.next();
                let right = self.parse_eq_expr()?;
                let end = right.span.end;
                left = Expr {
                    kind: ExprKind::BinOp(Box::new(left), BinOp::And, Box::new(right)),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes: Vec::new(),
                };
            } else {
                break;
            }
        }

        Ok(left)
    }

    fn parse_eq_expr(&mut self) -> Result<Expr, ParseError> {
        let mut left = self.parse_comp_expr()?;

        loop {
            let op = match self.peek() {
                Some(Token::Eq) => BinOp::Eq,
                Some(Token::NotEq) => BinOp::Neq,
                _ => break,
            };
            let start = left.span.start;
            self.next();
            let right = self.parse_comp_expr()?;
            let end = right.span.end;
            left = Expr {
                kind: ExprKind::BinOp(Box::new(left), op, Box::new(right)),
                span: Span::new(&(start..end), self.file.clone()),
                attributes: Vec::new(),
            };
        }

        Ok(left)
    }

    fn parse_comp_expr(&mut self) -> Result<Expr, ParseError> {
        let mut left = self.parse_add_expr()?;

        loop {
            let op = match self.peek() {
                Some(Token::Less) => BinOp::Lt,
                Some(Token::Greater) => BinOp::Gt,
                Some(Token::LessEq) => BinOp::Leq,
                Some(Token::GreaterEq) => BinOp::Geq,
                _ => break,
            };
            let start = left.span.start;
            self.next();
            let right = self.parse_add_expr()?;
            let end = right.span.end;
            left = Expr {
                kind: ExprKind::BinOp(Box::new(left), op, Box::new(right)),
                span: Span::new(&(start..end), self.file.clone()),
                attributes: Vec::new(),
            };
        }

        Ok(left)
    }

    fn parse_add_expr(&mut self) -> Result<Expr, ParseError> {
        let mut left = self.parse_mul_expr()?;

        loop {
            let op = match self.peek() {
                Some(Token::Plus) => BinOp::Add,
                Some(Token::Minus) => BinOp::Sub,
                _ => break,
            };
            let start = left.span.start;
            self.next();
            let right = self.parse_mul_expr()?;
            let end = right.span.end;
            left = Expr {
                kind: ExprKind::BinOp(Box::new(left), op, Box::new(right)),
                span: Span::new(&(start..end), self.file.clone()),
                attributes: Vec::new(),
            };
        }

        Ok(left)
    }

    fn parse_mul_expr(&mut self) -> Result<Expr, ParseError> {
        let mut left = self.parse_unary_expr()?;

        loop {
            let op = match self.peek() {
                Some(Token::Mul) => BinOp::Mul,
                Some(Token::Div) => BinOp::Div,
                Some(Token::Mod) => BinOp::Mod,
                _ => break,
            };
            let start = left.span.start;
            self.next();
            let right = self.parse_unary_expr()?;
            let end = right.span.end;
            left = Expr {
                kind: ExprKind::BinOp(Box::new(left), op, Box::new(right)),
                span: Span::new(&(start..end), self.file.clone()),
                attributes: Vec::new(),
            };
        }

        Ok(left)
    }

    fn parse_unary_expr(&mut self) -> Result<Expr, ParseError> {
        let start = self.peek_span().unwrap_or(0..0).start;
        match self.peek() {
            Some(Token::Not) => {
                self.next();
                let expr = self.parse_unary_expr()?;
                let end = expr.span.end;
                Ok(Expr {
                    kind: ExprKind::UnOp(UnOp::Not, Box::new(expr)),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes: Vec::new(),
                })
            }
            Some(Token::Minus) => {
                self.next();
                let expr = self.parse_unary_expr()?;
                let end = expr.span.end;
                Ok(Expr {
                    kind: ExprKind::UnOp(UnOp::Neg, Box::new(expr)),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes: Vec::new(),
                })
            }
            _ => self.parse_postfix_expr(),
        }
    }

    fn parse_postfix_expr(&mut self) -> Result<Expr, ParseError> {
        let mut expr = self.parse_primary_expr()?;

        loop {
            match self.peek() {
                Some(Token::LParen) => {
                    // Function call
                    let start = expr.span.start;
                    self.next();
                    let mut args = Vec::new();
                    loop {
                        if matches!(self.peek(), Some(Token::RParen)) {
                            self.next();
                            break;
                        }
                        args.push(self.parse_expr()?);

                        if matches!(self.peek(), Some(Token::Comma)) {
                            self.next();
                        }
                    }
                    let end = self.peek_span().unwrap_or(expr.span.end..expr.span.end).end;
                    expr = Expr {
                        kind: ExprKind::Call(Box::new(expr), args),
                        span: Span::new(&(start..end), self.file.clone()),
                        attributes: Vec::new(),
                    };
                }
                Some(Token::LBracket) => {
                    // Index
                    let start = expr.span.start;
                    self.next();
                    let index = self.parse_expr()?;
                    self.expect(Token::RBracket)?;
                    let end = self.peek_span().unwrap_or(expr.span.end..expr.span.end).end;
                    expr = Expr {
                        kind: ExprKind::Index(Box::new(expr), Box::new(index)),
                        span: Span::new(&(start..end), self.file.clone()),
                        attributes: Vec::new(),
                    };
                }
                Some(Token::Dot) => {
                    // Field access
                    let start = expr.span.start;
                    self.next();
                    let field = match self.next() {
                        Some((Token::Variable(f), _)) => f,
                        Some((_, span)) => {
                            return self.error("Expected field name".to_string(), span);
                        }
                        None => return self.error("Expected field name".to_string(), start..start),
                    };
                    let end = self.peek_span().unwrap_or(expr.span.end..expr.span.end).end;
                    expr = Expr {
                        kind: ExprKind::Dot(Box::new(expr), field),
                        span: Span::new(&(start..end), self.file.clone()),
                        attributes: Vec::new(),
                    };
                }
                Some(Token::OptionalChain) => {
                    // Optional chain
                    let start = expr.span.start;
                    self.next();
                    let field = match self.next() {
                        Some((Token::Variable(f), _)) => f,
                        Some((_, span)) => {
                            return self.error("Expected field name".to_string(), span);
                        }
                        None => return self.error("Expected field name".to_string(), start..start),
                    };
                    let end = self.peek_span().unwrap_or(expr.span.end..expr.span.end).end;
                    expr = Expr {
                        kind: ExprKind::OptionalChain(Some(Box::new(expr)), field),
                        span: Span::new(&(start..end), self.file.clone()),
                        attributes: Vec::new(),
                    };
                }
                Some(Token::Unwrap) => {
                    // Early return / unwrap
                    let start = expr.span.start;
                    self.next();
                    let end = self.peek_span().unwrap_or(expr.span.end..expr.span.end).end;
                    expr = Expr {
                        kind: ExprKind::EarlyReturn(Some(Box::new(expr))),
                        span: Span::new(&(start..end), self.file.clone()),
                        attributes: Vec::new(),
                    };
                }
                Some(Token::KeywordAs) => {
                    // Cast
                    let start = expr.span.start;
                    self.next();
                    let type_annot = self.parse_type_annot()?;
                    let end = self.peek_span().unwrap_or(expr.span.end..expr.span.end).end;
                    expr = Expr {
                        kind: ExprKind::Cast(Box::new(expr), type_annot),
                        span: Span::new(&(start..end), self.file.clone()),
                        attributes: Vec::new(),
                    };
                }
                _ => break,
            }
        }

        Ok(expr)
    }

    fn parse_primary_expr(&mut self) -> Result<Expr, ParseError> {
        let start = self.peek_span().unwrap_or(0..0).start;
        match self.peek().cloned() {
            Some(Token::Int(n)) => {
                self.next();
                let end = self.peek_span().unwrap_or(start..start).end;
                Ok(Expr {
                    kind: ExprKind::Int(n),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes: Vec::new(),
                })
            }
            Some(Token::Float(f)) => {
                self.next();
                let end = self.peek_span().unwrap_or(start..start).end;
                Ok(Expr {
                    kind: ExprKind::Float(f),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes: Vec::new(),
                })
            }
            Some(Token::Bool(b)) => {
                self.next();
                let end = self.peek_span().unwrap_or(start..start).end;
                Ok(Expr {
                    kind: ExprKind::Bool(b),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes: Vec::new(),
                })
            }
            Some(Token::String(s)) => {
                self.next();
                let end = self.peek_span().unwrap_or(start..start).end;
                Ok(Expr {
                    kind: ExprKind::String(s),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes: Vec::new(),
                })
            }
            Some(Token::Variable(name)) => {
                self.next();

                // Check for struct literal or enum variant
                if matches!(self.peek(), Some(Token::LBrace)) {
                    // Struct literal
                    self.next();
                    let mut fields = Vec::new();
                    loop {
                        if matches!(self.peek(), Some(Token::RBrace)) {
                            self.next();
                            break;
                        }

                        let field_name = match self.next() {
                            Some((Token::Variable(f), _)) => f,
                            Some((_, span)) => {
                                return self.error("Expected field name".to_string(), span);
                            }
                            None => {
                                return self.error("Expected field name".to_string(), start..start);
                            }
                        };

                        self.expect(Token::Colon)?;
                        let field_expr = self.parse_expr()?;
                        fields.push((field_name, field_expr));

                        if matches!(self.peek(), Some(Token::Comma)) {
                            self.next();
                        }
                    }

                    let end = self.peek_span().unwrap_or(start..start).end;
                    Ok(Expr {
                        kind: ExprKind::StructLit(name, fields),
                        span: Span::new(&(start..end), self.file.clone()),
                        attributes: Vec::new(),
                    })
                } else if matches!(self.peek(), Some(Token::Access)) {
                    // Enum variant
                    self.next();
                    let variant = match self.next() {
                        Some((Token::Variable(v), _)) => v,
                        Some((_, span)) => {
                            return self.error("Expected variant name".to_string(), span);
                        }
                        None => {
                            return self.error("Expected variant name".to_string(), start..start);
                        }
                    };

                    let mut args = Vec::new();
                    if matches!(self.peek(), Some(Token::LParen)) {
                        self.next();
                        loop {
                            if matches!(self.peek(), Some(Token::RParen)) {
                                self.next();
                                break;
                            }
                            args.push(self.parse_expr()?);

                            if matches!(self.peek(), Some(Token::Comma)) {
                                self.next();
                            }
                        }
                    }

                    let end = self.peek_span().unwrap_or(start..start).end;
                    Ok(Expr {
                        kind: ExprKind::EnumLit(name, variant, args),
                        span: Span::new(&(start..end), self.file.clone()),
                        attributes: Vec::new(),
                    })
                } else {
                    let end = self.peek_span().unwrap_or(start..start).end;
                    Ok(Expr {
                        kind: ExprKind::Variable(name),
                        span: Span::new(&(start..end), self.file.clone()),
                        attributes: Vec::new(),
                    })
                }
            }
            Some(Token::LParen) => {
                self.next();
                if matches!(self.peek(), Some(Token::RParen)) {
                    // Empty tuple
                    self.next();
                    let end = self.peek_span().unwrap_or(start..start).end;
                    Ok(Expr {
                        kind: ExprKind::Tuple(vec![]),
                        span: Span::new(&(start..end), self.file.clone()),
                        attributes: Vec::new(),
                    })
                } else {
                    let first = self.parse_expr()?;
                    if matches!(self.peek(), Some(Token::Comma)) {
                        // Tuple
                        let mut elements = vec![first];
                        self.next();
                        loop {
                            if matches!(self.peek(), Some(Token::RParen)) {
                                break;
                            }
                            elements.push(self.parse_expr()?);

                            if matches!(self.peek(), Some(Token::Comma)) {
                                self.next();
                            }
                        }
                        self.expect(Token::RParen)?;
                        let end = self.peek_span().unwrap_or(start..start).end;
                        Ok(Expr {
                            kind: ExprKind::Tuple(elements),
                            span: Span::new(&(start..end), self.file.clone()),
                            attributes: Vec::new(),
                        })
                    } else {
                        self.expect(Token::RParen)?;
                        let end = self.peek_span().unwrap_or(start..start).end;
                        Ok(Expr {
                            kind: first.kind,
                            span: Span::new(&(start..end), self.file.clone()),
                            attributes: first.attributes,
                        })
                    }
                }
            }
            Some(Token::LBracket) => {
                self.next();
                let mut elements = Vec::new();
                loop {
                    if matches!(self.peek(), Some(Token::RBracket)) {
                        self.next();
                        break;
                    }
                    elements.push(self.parse_expr()?);

                    if matches!(self.peek(), Some(Token::Comma)) {
                        self.next();
                    }
                }

                let end = self.peek_span().unwrap_or(start..start).end;
                Ok(Expr {
                    kind: ExprKind::Array(elements),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes: Vec::new(),
                })
            }
            Some(Token::KeywordLet) => {
                self.next();

                // Parse binding kind (mut, uniq, once) - comes AFTER let
                let binding_kind = match self.peek() {
                    Some(Token::KeywordMut) => {
                        self.next();
                        BindingKind::Mutable
                    }
                    Some(Token::KeywordUniq) => {
                        self.next();
                        BindingKind::Affine
                    }
                    Some(Token::KeywordOnce) => {
                        self.next();
                        BindingKind::Linear
                    }
                    _ => BindingKind::Default,
                };

                // Now parse the variable name
                let var_name = match self.next() {
                    Some((Token::Variable(n), _)) => n,
                    Some((_, span)) => {
                        return self.error("Expected variable name".to_string(), span);
                    }
                    None => return self.error("Expected variable name".to_string(), start..start),
                };

                // Parse optional type annotation
                let type_annot = if matches!(self.peek(), Some(Token::Colon)) {
                    self.next();
                    Some(self.parse_type_annot()?)
                } else {
                    None
                };

                self.expect(Token::Assign)?;
                let expr = self.parse_expr()?;
                let end = expr.span.end;
                Ok(Expr {
                    kind: ExprKind::Let(var_name, binding_kind, type_annot, Box::new(expr)),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes: Vec::new(),
                })
            }
            Some(Token::KeywordIf) => {
                self.next();
                let cond = self.parse_expr()?;
                let then_expr = self.parse_expr()?;
                let else_expr = if matches!(self.peek(), Some(Token::KeywordElse)) {
                    self.next();
                    Some(Box::new(self.parse_expr()?))
                } else {
                    None
                };

                let end = else_expr
                    .as_ref()
                    .map(|e| e.span.end)
                    .unwrap_or(then_expr.span.end);

                Ok(Expr {
                    kind: ExprKind::If(Box::new(cond), Box::new(then_expr), else_expr),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes: Vec::new(),
                })
            }
            Some(Token::KeywordMatch) => {
                self.next();
                let expr = self.parse_expr()?;
                let mut arms = Vec::new();

                loop {
                    if matches!(self.peek(), Some(Token::KeywordEnd)) {
                        self.next();
                        break;
                    }

                    let pattern = self.parse_pattern()?;
                    self.expect(Token::FatArrow)?;
                    let body = self.parse_expr()?;
                    arms.push((pattern, body));

                    if matches!(self.peek(), Some(Token::Comma)) {
                        self.next();
                    }
                }

                let end = self.peek_span().unwrap_or(start..start).end;
                Ok(Expr {
                    kind: ExprKind::Match(Box::new(expr), arms),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes: Vec::new(),
                })
            }
            Some(Token::KeywordWhile) => {
                self.next();
                let cond = self.parse_expr()?;
                let body = self.parse_expr()?;
                let end = body.span.end;

                Ok(Expr {
                    kind: ExprKind::While(Box::new(cond), Box::new(body)),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes: Vec::new(),
                })
            }
            Some(Token::KeywordFor) => {
                self.next();
                let var = match self.next() {
                    Some((Token::Variable(v), _)) => v,
                    Some((_, span)) => {
                        return self.error("Expected variable name".to_string(), span);
                    }
                    None => return self.error("Expected variable name".to_string(), start..start),
                };
                self.expect(Token::KeywordIn)?;
                let iterable = self.parse_expr()?;
                let body = self.parse_expr()?;
                let end = body.span.end;

                Ok(Expr {
                    kind: ExprKind::For(var, Box::new(iterable), Box::new(body)),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes: Vec::new(),
                })
            }
            Some(Token::KeywordDo) => {
                self.next();
                let mut exprs = Vec::new();

                loop {
                    if matches!(self.peek(), Some(Token::KeywordEnd)) {
                        self.next();
                        break;
                    }
                    exprs.push(self.parse_expr()?);

                    if matches!(self.peek(), Some(Token::Semicolon)) {
                        self.next();
                    }
                }

                let end = self.peek_span().unwrap_or(start..start).end;
                Ok(Expr {
                    kind: ExprKind::Do(exprs),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes: Vec::new(),
                })
            }
            Some(Token::KeywordLambda) => {
                let start = self.peek_span().unwrap_or(0..0).start;
                self.next();
                self.expect(Token::LParen)?;
                let mut params = Vec::new();

                loop {
                    if matches!(self.peek(), Some(Token::RParen)) {
                        self.next();
                        break;
                    }

                    let param_name = match self.next() {
                        Some((Token::Variable(p), _)) => p,
                        Some((_, span)) => {
                            return self.error("Expected parameter name".to_string(), span);
                        }
                        None => {
                            return self.error("Expected parameter name".to_string(), start..start);
                        }
                    };

                    // Check for optional type annotation
                    let param_type = if matches!(self.peek(), Some(Token::Colon)) {
                        self.next(); // consume ':'
                        Some(self.parse_type_annot()?)
                    } else {
                        None
                    };

                    params.push((param_name, param_type));

                    if matches!(self.peek(), Some(Token::Comma)) {
                        self.next();
                    }
                }

                let body = self.parse_expr()?;
                let end = body.span.end;

                Ok(Expr {
                    kind: ExprKind::Lambda(params, Box::new(body)),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes: Vec::new(),
                })
            }
            Some(Token::KeywordReturn) => {
                self.next();
                let expr = if self.is_expr_end() {
                    None
                } else {
                    Some(Box::new(self.parse_expr()?))
                };

                let end = expr
                    .as_ref()
                    .map(|e| e.span.end)
                    .unwrap_or(self.peek_span().unwrap_or(start..start).end);

                Ok(Expr {
                    kind: ExprKind::Return(expr),
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes: Vec::new(),
                })
            }
            Some(Token::KeywordBreak) => {
                self.next();
                let end = self.peek_span().unwrap_or(start..start).end;
                Ok(Expr {
                    kind: ExprKind::Break,
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes: Vec::new(),
                })
            }
            Some(Token::KeywordContinue) => {
                self.next();
                let end = self.peek_span().unwrap_or(start..start).end;
                Ok(Expr {
                    kind: ExprKind::Continue,
                    span: Span::new(&(start..end), self.file.clone()),
                    attributes: Vec::new(),
                })
            }
            Some(token) => {
                let span = self.peek_span().unwrap_or(start..start);
                self.error(format!("Unexpected token: {:?}", token), span)
            }
            None => self.error("Unexpected EOF".to_string(), start..start),
        }
    }

    fn parse_pattern(&mut self) -> Result<Pattern, ParseError> {
        let start = self.peek_span().unwrap_or(0..0).start;

        match self.peek().cloned() {
            Some(Token::Variable(name)) => {
                self.next();
                let end = self.peek_span().unwrap_or(start..start).end;

                // Check for struct or enum pattern
                if matches!(self.peek(), Some(Token::LBrace)) {
                    // Struct pattern
                    self.next();
                    let mut fields = Vec::new();

                    loop {
                        if matches!(self.peek(), Some(Token::RBrace)) {
                            self.next();
                            break;
                        }

                        let field_name = match self.next() {
                            Some((Token::Variable(f), _)) => f,
                            Some((_, span)) => {
                                return self.error("Expected field name".to_string(), span);
                            }
                            None => {
                                return self.error("Expected field name".to_string(), start..start);
                            }
                        };

                        self.expect(Token::Colon)?;
                        let pattern = self.parse_pattern()?;
                        fields.push((field_name, pattern));

                        if matches!(self.peek(), Some(Token::Comma)) {
                            self.next();
                        }
                    }

                    let end = self.peek_span().unwrap_or(start..start).end;
                    Ok(Pattern {
                        kind: PatternKind::Struct(name, fields),
                        span: Span::new(&(start..end), self.file.clone()),
                    })
                } else if matches!(self.peek(), Some(Token::Access)) {
                    // Enum pattern
                    self.next();
                    let variant = match self.next() {
                        Some((Token::Variable(v), _)) => v,
                        Some((_, span)) => {
                            return self.error("Expected variant name".to_string(), span);
                        }
                        None => {
                            return self.error("Expected variant name".to_string(), start..start);
                        }
                    };

                    let mut patterns = Vec::new();
                    if matches!(self.peek(), Some(Token::LParen)) {
                        self.next();
                        loop {
                            if matches!(self.peek(), Some(Token::RParen)) {
                                self.next();
                                break;
                            }
                            patterns.push(self.parse_pattern()?);

                            if matches!(self.peek(), Some(Token::Comma)) {
                                self.next();
                            }
                        }
                    }

                    let end = self.peek_span().unwrap_or(start..start).end;
                    Ok(Pattern {
                        kind: PatternKind::Enum(name, variant, patterns),
                        span: Span::new(&(start..end), self.file.clone()),
                    })
                } else {
                    Ok(Pattern {
                        kind: PatternKind::Variable(name),
                        span: Span::new(&(start..end), self.file.clone()),
                    })
                }
            }
            Some(Token::Union) => {
                self.next();
                let end = self.peek_span().unwrap_or(start..start).end;
                Ok(Pattern {
                    kind: PatternKind::Wildcard,
                    span: Span::new(&(start..end), self.file.clone()),
                })
            }
            Some(Token::LParen) => {
                self.next();
                let mut patterns = Vec::new();
                loop {
                    if matches!(self.peek(), Some(Token::RParen)) {
                        self.next();
                        break;
                    }
                    patterns.push(self.parse_pattern()?);

                    if matches!(self.peek(), Some(Token::Comma)) {
                        self.next();
                    }
                }

                let end = self.peek_span().unwrap_or(start..start).end;
                Ok(Pattern {
                    kind: PatternKind::Tuple(patterns),
                    span: Span::new(&(start..end), self.file.clone()),
                })
            }
            Some(Token::String(s)) => {
                self.next();
                let end = self.peek_span().unwrap_or(start..start).end;
                Ok(Pattern {
                    kind: PatternKind::Literal(s),
                    span: Span::new(&(start..end), self.file.clone()),
                })
            }
            Some(Token::Int(n)) => {
                self.next();
                let end = self.peek_span().unwrap_or(start..start).end;
                Ok(Pattern {
                    kind: PatternKind::Literal(n.to_string()),
                    span: Span::new(&(start..end), self.file.clone()),
                })
            }
            Some(Token::Float(f)) => {
                self.next();
                let end = self.peek_span().unwrap_or(start..start).end;
                Ok(Pattern {
                    kind: PatternKind::Literal(f.to_string()),
                    span: Span::new(&(start..end), self.file.clone()),
                })
            }
            Some(Token::Bool(b)) => {
                self.next();
                let end = self.peek_span().unwrap_or(start..start).end;
                Ok(Pattern {
                    kind: PatternKind::Literal(b.to_string()),
                    span: Span::new(&(start..end), self.file.clone()),
                })
            }
            Some(token) => {
                let span = self.peek_span().unwrap_or(start..start);
                self.error(format!("Unexpected token in pattern: {:?}", token), span)
            }
            None => self.error("Unexpected EOF".to_string(), start..start),
        }
    }

    fn is_expr_end(&mut self) -> bool {
        matches!(
            self.peek(),
            Some(Token::RParen)
                | Some(Token::RBracket)
                | Some(Token::RBrace)
                | Some(Token::Comma)
                | Some(Token::Semicolon)
                | Some(Token::KeywordEnd)
                | Some(Token::FatArrow)
        )
    }
}

// src/main.rs
use logos::Logos;
use std::fs;
use suicmez::{lexer::Token, parser::Parser, typechecker::TypeChecker};

fn main() {
    // Check if a file was provided as argument
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 2 {
        // Run all test files in the tests directory
        run_test_suite();
        return;
    }

    let filename = &args[1];
    println!("Type checking file: {}", filename);

    if let Err(e) = run_file(filename) {
        eprintln!("Error: {}", e);
    }
}

fn run_test_suite() {
    println!("Running test suite...\n");

    let test_files = vec![
        "tests/basic_types.sui",
        "tests/structs.sui",
        "tests/enums.sui",
        "tests/functions.sui",
        "tests/arrays.sui",
        "tests/traits.sui",
        "tests/control_flow.sui",
    ];

    for file in test_files {
        println!("Testing: {}", file);
        match run_file(file) {
            Ok(_) => println!("✓ Passed\n"),
            Err(e) => println!("✗ Failed: {}\n", e),
        }
    }
}

fn run_file(filename: &str) -> Result<(), String> {
    // Read the source file
    let source = fs::read_to_string(filename)
        .map_err(|e| format!("Error reading file {}: {}", filename, e))?;

    // First, we need to parse the source code
    let mut tokens = Vec::new();
    let mut lexer = Token::lexer(&source);

    loop {
        match lexer.next() {
            Some(Ok(token)) => {
                let span = lexer.span();
                tokens.push((token, span));
            }
            Some(Err(_)) => {
                return Err("Lexing error".to_string());
            }
            None => break,
        }
    }

    let mut parser = Parser::new(filename.to_string(), tokens);
    let ast_nodes = parser
        .parse()
        .map_err(|e| format!("Parse error: {}", e.message))?;

    println!("Parsed {} AST nodes successfully", ast_nodes.len());

    // Typecheck the AST
    let mut typechecker = TypeChecker::new();
    let typed_nodes = typechecker.typecheck_program(&ast_nodes).map_err(|e| {
        format!(
            "Type error at {}:{}: {:?}",
            e.span.file, e.span.start, e.kind
        )
    })?;

    println!(
        "Type checking passed! {} nodes typechecked.",
        typed_nodes.len()
    );

    Ok(())
}

// src/typechecker.rs
#[derive(Debug, Clone, PartialEq)]
pub enum Type {
    Stub,
}

// src/ast.rs
use crate::typechecker::Type;
use std::ops::Range;

#[derive(Debug, Clone)]
pub enum TypeAnnot {
    Var(String),
    Cons(String, Vec<TypeAnnot>),
    Function(Vec<TypeAnnot>, Box<TypeAnnot>),
    Tuple(Vec<TypeAnnot>),
    Array(Box<TypeAnnot>),
}

#[derive(Debug, Clone)]
pub struct Span {
    pub start: usize,
    pub end: usize,
    pub file: String,
}

impl Span {
    pub fn new(range: &Range<usize>, file: String) -> Self {
        Span {
            start: range.start,
            end: range.end,
            file,
        }
    }

    pub fn merge(&self, other: &Span) -> Span {
        Span {
            start: self.start.min(other.start),
            end: self.end.max(other.end),
            file: self.file.clone(),
        }
    }
}

// @attribute
#[derive(Debug, Clone)]
pub struct Attribute {
    pub name: String,
    pub args: Vec<AttributeArg>,
    pub span: Span,
}

#[derive(Debug, Clone)]
pub enum AttributeArg {
    Value(String),            // some_identifier
    KeyValue(String, String), // some_key = some_identifier
    Literal(String),          // some literal value
}

#[derive(Debug, Clone)]
pub struct ASTNode {
    pub kind: ASTNodeKind,
    pub span: Span,
    pub attributes: Vec<Attribute>,
}

#[derive(Debug, Clone)]
pub enum ASTNodeKind {
    Function(Function),
    Extern(Extern),
    Load(Load),
    Struct(Struct),
    Enum(Enum),
    Impl(Impl),
    Trait(Trait),
    Use(String),
}

// ? implies OPTIONAL here
// \( implies the presence of (. same for /)

#[derive(Debug, Clone)]
/// fn name\( (arg: type?,)* \) -> return_type? body
pub struct Function {
    pub name: String,
    pub parameters: Vec<Parameter>, // type params
    pub args: Vec<(String, Option<TypeAnnot>)>,
    pub return_type: Option<TypeAnnot>,
    pub body: Expr,
}

/// extern name\( type?,* \) -> return_type from library_alias
#[derive(Debug, Clone)]
pub struct Extern {
    pub name: String,
    pub args: Vec<TypeAnnot>,
    pub return_type: TypeAnnot,
    pub from: String,
    pub span: Span,
}

/// load "library" as alias
#[derive(Debug, Clone)]
pub struct Load {
    pub library: String,
    pub alias: String,
    pub span: Span,
}

/// struct name <parameter*>?
///   (field_name: field_type,)*
/// end
#[derive(Debug, Clone)]
pub struct Struct {
    pub name: String,
    pub parameters: Vec<Parameter>, // type parameters
    pub fields: Vec<Field>,
}

#[derive(Debug, Clone)]
pub struct Field {
    pub name: String,
    pub field_type: TypeAnnot,
    pub span: Span,
}

/// enum name <parameter*>?
///   VariantName\(field_type,\)*
/// end
#[derive(Debug, Clone)]
pub struct Enum {
    pub name: String,
    pub parameters: Vec<Parameter>, // type parameters
    pub variants: Vec<Variant>,
}

#[derive(Debug, Clone)]
pub struct Parameter {
    pub name: String,
    pub bounds: Vec<String>, // trait bounds
    pub kind: Option<Kind>,  // for HKTs
    pub span: Span,
}

#[derive(Debug, Clone, PartialEq)]
pub enum Kind {
    Star,                        // *
    Arrow(Box<Kind>, Box<Kind>), // k1 -> k2
}

#[derive(Debug, Clone)]
pub struct Variant {
    pub name: String,
    pub fields: Vec<TypeAnnot>,
    pub span: Span,
}

/// impl TypeName <parameter*>? (: TraitName)?
///     functions*
/// end
#[derive(Debug, Clone)]
pub struct Impl {
    pub target: String,
    pub trait_name: Option<String>,
    pub methods: Vec<Function>,
}

/// trait TraitName <parameter*>?
///     function_signatures*
/// end
#[derive(Debug, Clone)]
pub struct Trait {
    pub name: String,
    pub methods: Vec<FunctionSignature>,

    pub parameters: Vec<Parameter>,
    pub associated_types: Vec<AssociatedType>,
}

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

#[derive(Debug, Clone)]
pub struct FunctionSignature {
    pub name: String,
    pub params: Vec<Parameter>,
    pub return_type: TypeAnnot,
}

#[derive(Debug, Clone)]
pub struct Expr {
    pub kind: ExprKind,
    pub span: Span,
    pub attributes: Vec<Attribute>,
}

#[derive(Debug, Clone)]
pub enum ExprKind {
    Int(i64),
    Float(f64),
    Bool(bool),
    String(String),
    Array(Vec<Expr>),
    Tuple(Vec<Expr>),

    StructLit(String, Vec<(String, Expr)>), // Name { a: expr, b: expr }
    EnumLit(String, String, Vec<Expr>),     // Name::Variant(expr, expr)

    Variable(String),

    Call(Box<Expr>, Vec<Expr>),
    Index(Box<Expr>, Box<Expr>),
    Dot(Box<Expr>, String),
    EarlyReturn(Option<Box<Expr>>),           // eg: myresultoroption?
    OptionalChain(Option<Box<Expr>>, String), // a?.b

    Lambda(Vec<(String, Option<TypeAnnot>)>, Box<Expr>), // lambda (arg, arg: optionalty, ...) body
    Let(String, BindingKind, Option<TypeAnnot>, Box<Expr>), // no patterns for now
    Assign(Box<Expr>, Box<Expr>), // NOTE: check for valid lvalue during typechecking
    Cast(Box<Expr>, TypeAnnot),

    If(Box<Expr>, Box<Expr>, Option<Box<Expr>>), // if cond expr (else expr)?
    Match(Box<Expr>, Vec<(Pattern, Expr)>),      // match expr pattern => expr* end
    While(Box<Expr>, Box<Expr>),                 // while cond expr

    For(String, Box<Expr>, Box<Expr>), // for i in expr body
    Range(Box<Expr>, Box<Expr>),       // 0..10

    Do(Vec<Expr>), // do expr* end
    BinOp(Box<Expr>, BinOp, Box<Expr>),
    UnOp(UnOp, Box<Expr>),

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

#[derive(Debug, Clone, PartialEq)]
pub enum BindingKind {
    Default, // immutable but infinite usages
    Mutable, // mutable but infinite usages
    Affine,
    Linear,
}

#[derive(Debug, Clone)]
pub enum BinOp {
    Add,
    Sub,
    Mul,
    Div,
    Mod,
    And,
    Or,
    Eq,
    Neq,
    Lt,
    Gt,
    Leq,
    Geq,
}

#[derive(Debug, Clone)]
pub enum UnOp {
    Neg,
    Not,
}

#[derive(Debug, Clone)]
pub struct Pattern {
    pub kind: PatternKind,
    pub span: Span,
}

#[derive(Debug, Clone)]
pub enum PatternKind {
    Wildcard, // _
    Variable(String),
    Literal(String),
    Tuple(Vec<Pattern>),
    Struct(String, Vec<(String, Pattern)>),
    Enum(String, String, Vec<Pattern>),
    Range(i64, i64),
}

// Typed variants

#[derive(Debug, Clone)]
pub struct TypedASTNode {
    pub kind: TypedASTNodeKind,
    pub span: Span,
    pub attributes: Vec<Attribute>,
    pub ty: Type,
}

#[derive(Debug, Clone)]
pub enum TypedASTNodeKind {
    Function(TypedFunction),
    Extern(TypedExtern),
    Load(TypedLoad),
    Struct(TypedStruct),
    Enum(TypedEnum),
    Impl(TypedImpl),
    Trait(TypedTrait),
    Use(String),
}

#[derive(Debug, Clone)]
pub struct TypedFunction {
    pub name: String,
    pub parameters: Vec<Parameter>,
    pub args: Vec<(String, Option<TypeAnnot>)>,
    pub return_type: Option<TypeAnnot>,
    pub body: TypedExpr,
    pub ty: Type,
}

#[derive(Debug, Clone)]
pub struct TypedExtern {
    pub name: String,
    pub args: Vec<TypeAnnot>,
    pub return_type: TypeAnnot,
    pub from: String,
    pub span: Span,
}

#[derive(Debug, Clone)]
pub struct TypedLoad {
    pub library: String,
    pub alias: String,
    pub span: Span,
}

#[derive(Debug, Clone)]
pub struct TypedStruct {
    pub name: String,
    pub parameters: Vec<Parameter>,
    pub fields: Vec<TypedField>,
}

#[derive(Debug, Clone)]
pub struct TypedField {
    pub name: String,
    pub field_type: TypeAnnot,
    pub span: Span,
}

#[derive(Debug, Clone)]
pub struct TypedEnum {
    pub name: String,
    pub parameters: Vec<Parameter>,
    pub variants: Vec<TypedVariant>,
}

#[derive(Debug, Clone)]
pub struct TypedVariant {
    pub name: String,
    pub fields: Vec<TypeAnnot>,
    pub span: Span,
}

#[derive(Debug, Clone)]
pub struct TypedImpl {
    pub target: String,
    pub trait_name: Option<String>,
    pub methods: Vec<TypedFunction>,
}

#[derive(Debug, Clone)]
pub struct TypedTrait {
    pub name: String,
    pub methods: Vec<FunctionSignature>,
    pub parameters: Vec<Parameter>,
    pub associated_types: Vec<AssociatedType>,
}

#[derive(Debug, Clone)]
pub struct TypedExpr {
    pub kind: TypedExprKind,
    pub span: Span,
    pub attributes: Vec<Attribute>,
    pub ty: Type,
}

#[derive(Debug, Clone)]
pub enum TypedExprKind {
    Int(i64),
    Float(f64),
    Bool(bool),
    String(String),
    Array(Vec<TypedExpr>),
    Tuple(Vec<TypedExpr>),
    StructLit(String, Vec<(String, TypedExpr)>),
    EnumLit(String, String, Vec<TypedExpr>),
    Variable(String),
    Call(Box<TypedExpr>, Vec<TypedExpr>),
    Index(Box<TypedExpr>, Box<TypedExpr>),
    Dot(Box<TypedExpr>, String),
    EarlyReturn(Option<Box<TypedExpr>>),
    OptionalChain(Option<Box<TypedExpr>>, String),
    Lambda(Vec<(String, Option<TypeAnnot>)>, Box<TypedExpr>),
    Let(String, BindingKind, Option<TypeAnnot>, Box<TypedExpr>),
    Assign(Box<TypedExpr>, Box<TypedExpr>),
    Cast(Box<TypedExpr>, TypeAnnot),
    If(Box<TypedExpr>, Box<TypedExpr>, Option<Box<TypedExpr>>),
    Match(Box<TypedExpr>, Vec<(TypedPattern, TypedExpr)>),
    While(Box<TypedExpr>, Box<TypedExpr>),
    Do(Vec<TypedExpr>),
    BinOp(Box<TypedExpr>, BinOp, Box<TypedExpr>),
    UnOp(UnOp, Box<TypedExpr>),
    For(String, Box<TypedExpr>, Box<TypedExpr>),
    Range(Box<TypedExpr>, Box<TypedExpr>),
    Return(Option<Box<TypedExpr>>),
    Break,
    Continue,
}

#[derive(Debug, Clone)]
pub struct TypedPattern {
    pub kind: TypedPatternKind,
    pub span: Span,
    pub ty: Type,
}

#[derive(Debug, Clone)]
pub enum TypedPatternKind {
    Wildcard,
    Variable(String),
    Literal(String),
    Tuple(Vec<TypedPattern>),
    Struct(String, Vec<(String, TypedPattern)>),
    Enum(String, String, Vec<TypedPattern>),
}

// src/lexer/mod.rs
use logos::Logos;

#[cfg(test)]
pub mod tests;

#[derive(Logos, Debug, PartialEq)]
#[logos(skip r"[ \n\r\t\f]+")] // Ignore this regex pattern between tokens
#[logos(skip r"#(.*)\n")] // Ignore this regex pattern between tokens
#[derive(Clone)]
pub enum Token {
    #[regex(r"true|false", |lex| {
        lex.slice().parse::<bool>().unwrap()
    })]
    Bool(bool),

    #[regex(r"0|[1-9][0-9_]*", |lex| {
        let s = lex.slice().replace("_", "");
        // We parse to i64 for wider support.
        s.parse::<i64>().unwrap()
    }, priority = 4)]
    Int(i64),

    #[regex(r"(([0-9][0-9_]*\.[0-9_]+|[0-9]*\.[0-9_]+)([eE][+-]?[0-9_]+)?)", |lex| {
        let s = lex.slice().replace("_", "");
        s.parse::<f64>().unwrap()
    }, priority = 3)]
    Float(f64),

    #[regex(r#""([^"\\]*(\\.[^"\\]*)*)""#, |lex| {
        let s = lex.slice();
        s[1..s.len()-1]
            .replace("\\\"", "\"")
            .replace("\\\\", "\\")
            .replace("\\n", "\n")
            .replace("\\r", "\r")
            .replace("\\t", "\t")
    })]
    String(String),

    #[regex(r#"r#"([^"]*)""#, |lex| {
        let s = lex.slice();
        // Remove the outer r" and " (s[2..s.len() - 1])
        s[3..s.len() - 1].to_string()
    })]
    RawString(String),

    #[regex(r"[a-zA-Z_][a-zA-Z0-9_]*", |lex|{
        lex.slice().to_string()
    })]
    Variable(String),

    #[token("bool")]
    KeywordBool,

    #[token("int")]
    KeywordInt,

    #[token("float")]
    KeywordFloat,

    #[token("string")]
    KeywordString,

    #[token("let")]
    KeywordLet,

    #[token("mut")]
    KeywordMut,

    #[token("uniq")]
    KeywordUniq,

    #[token("once")]
    KeywordOnce,

    #[token("if")]
    KeywordIf,

    #[token("then")]
    KeywordThen,

    #[token("else")]
    KeywordElse,

    #[token("fn")]
    KeywordFn,

    #[token("lambda")]
    KeywordLambda,

    #[token("do")]
    KeywordDo,

    #[token("end")]
    KeywordEnd,

    #[token("as")]
    KeywordAs,

    #[token("in")]
    KeywordIn,

    #[token("for")]
    KeywordFor,
    #[token("while")]
    KeywordWhile,

    #[token("loop")]
    KeywordLoop,

    #[token("where")]
    KeywordWhere,

    #[token("extern")]
    KeywordExtern,

    #[token("load")]
    KeywordLoad,

    #[token("from")]
    KeywordFrom,

    #[token("use")]
    KeywordUse,

    #[token("struct")]
    KeywordStruct,

    #[token("enum")]
    KeywordEnum,

    #[token("impl")]
    KeywordImpl,

    #[token("trait")]
    KeywordTrait,

    // #[token("type")]
    // KeywordType,
    //
    #[token("match")]
    KeywordMatch,

    #[token("return")]
    KeywordReturn,

    #[token("break")]
    KeywordBreak,

    #[token("continue")]
    KeywordContinue,

    #[token("+")]
    Plus,

    #[token("-")]
    Minus,

    #[token("*")]
    Mul,

    #[token("/")]
    Div,

    #[token("%")]
    Mod,

    #[token("**", priority = 3)]
    Power,

    #[token("$")]
    Dollar,

    #[token("@")]
    At,

    #[token("==")]
    Eq,

    #[token("!=")]
    NotEq,

    #[token("<")]
    Less,

    #[token(">")]
    Greater,

    #[token("<=")]
    LessEq,

    #[token(">=")]
    GreaterEq,

    #[token("and")]
    And,

    #[token("or")]
    Or,

    #[token("xor")]
    Xor,

    #[token("nor")]
    Nor,

    #[token("not")]
    Not,

    #[token("(")]
    LParen,

    #[token(")")]
    RParen,

    #[token("[")]
    LBracket,

    #[token("]")]
    RBracket,

    #[token("{")]
    LBrace,

    #[token("}")]
    RBrace,

    #[token(",")]
    Comma,

    #[token(";")]
    Semicolon,

    #[token(":")]
    Colon,

    #[token(".")]
    Dot,

    #[token("...")]
    Spread,

    #[token("..")]
    DotDot,

    #[token("::")]
    Access,

    #[token("->")]
    Arrow,

    #[token("~")]
    Tilde,

    #[token("!")]
    Bang,

    // New tokens for pattern matching
    #[token("=>")]
    FatArrow, // For match arms

    #[token("|")]
    Union,

    #[token("?.")]
    OptionalChain,

    #[token("?")]
    Unwrap,

    #[token("=")]
    Assign,

    #[token("+=")]
    AddAssign,

    #[token("-=")]
    SubAssign,

    #[token("*=")]
    MulAssign,

    #[token("/=")]
    DivAssign,

    #[token("%=")]
    ModAssign,
}

// src/lexer/tests.rs
use super::Token;
use logos::Logos;

#[test]
fn test_literals() {
    let mut lexer = Token::lexer("true false 42 2.14 \"hello\" r\"raw\"");

    assert_eq!(lexer.next(), Some(Ok(Token::Bool(true))));
    assert_eq!(lexer.next(), Some(Ok(Token::Bool(false))));
    assert_eq!(lexer.next(), Some(Ok(Token::Int(42))));
    assert_eq!(lexer.next(), Some(Ok(Token::Float(2.14))));
    assert_eq!(lexer.next(), Some(Ok(Token::String("hello".to_string()))));
    // RawString regex seems to have issues, let's test separately
    assert_eq!(lexer.next(), Some(Ok(Token::Variable("r".to_string()))));
    assert_eq!(lexer.next(), Some(Ok(Token::String("raw".to_string()))));
    assert_eq!(lexer.next(), None);
}

#[test]
fn test_int_literals() {
    let mut lexer = Token::lexer("0 123 1_000_000");

    assert_eq!(lexer.next(), Some(Ok(Token::Int(0))));
    assert_eq!(lexer.next(), Some(Ok(Token::Int(123))));
    assert_eq!(lexer.next(), Some(Ok(Token::Int(1000000))));
    assert_eq!(lexer.next(), None);
}

#[test]
fn test_string_literals() {
    let mut lexer = Token::lexer("\"hello world\" \"with\\\\escape\" \"quote\\\"here\"");

    assert_eq!(
        lexer.next(),
        Some(Ok(Token::String("hello world".to_string())))
    );
    assert_eq!(
        lexer.next(),
        Some(Ok(Token::String("with\\escape".to_string())))
    );
    assert_eq!(
        lexer.next(),
        Some(Ok(Token::String("quote\"here".to_string())))
    );
    assert_eq!(lexer.next(), None);
}

#[test]
fn test_keywords() {
    let mut lexer = Token::lexer(
        "bool int float string let if else fn do end as in for while loop where extern import struct enum impl trait match return break continue",
    );

    assert_eq!(lexer.next(), Some(Ok(Token::KeywordBool)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordInt)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordFloat)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordString)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordLet)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordIf)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordElse)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordFn)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordDo)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordEnd)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordAs)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordIn)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordFor)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordWhile)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordLoop)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordWhere)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordExtern)));
    assert_eq!(lexer.next(), Some(Ok(Token::Variable("import".into()))));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordStruct)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordEnum)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordImpl)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordTrait)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordMatch)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordReturn)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordBreak)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordContinue)));
    assert_eq!(lexer.next(), None);
}

#[test]
fn test_operators() {
    let mut lexer = Token::lexer("+ - * / % ** $ @ == != < > <= >= and or xor nor not");

    assert_eq!(lexer.next(), Some(Ok(Token::Plus)));
    assert_eq!(lexer.next(), Some(Ok(Token::Minus)));
    assert_eq!(lexer.next(), Some(Ok(Token::Mul)));
    assert_eq!(lexer.next(), Some(Ok(Token::Div)));
    assert_eq!(lexer.next(), Some(Ok(Token::Mod)));
    assert_eq!(lexer.next(), Some(Ok(Token::Power)));
    assert_eq!(lexer.next(), Some(Ok(Token::Dollar)));
    assert_eq!(lexer.next(), Some(Ok(Token::At)));
    assert_eq!(lexer.next(), Some(Ok(Token::Eq)));
    assert_eq!(lexer.next(), Some(Ok(Token::NotEq)));
    assert_eq!(lexer.next(), Some(Ok(Token::Less)));
    assert_eq!(lexer.next(), Some(Ok(Token::Greater)));
    assert_eq!(lexer.next(), Some(Ok(Token::LessEq)));
    assert_eq!(lexer.next(), Some(Ok(Token::GreaterEq)));
    assert_eq!(lexer.next(), Some(Ok(Token::And)));
    assert_eq!(lexer.next(), Some(Ok(Token::Or)));
    assert_eq!(lexer.next(), Some(Ok(Token::Xor)));
    assert_eq!(lexer.next(), Some(Ok(Token::Nor)));
    assert_eq!(lexer.next(), Some(Ok(Token::Not)));
    assert_eq!(lexer.next(), None);
}

#[test]
fn test_assignment_operators() {
    let mut lexer = Token::lexer("= += -= *= /= %=");

    assert_eq!(lexer.next(), Some(Ok(Token::Assign)));
    assert_eq!(lexer.next(), Some(Ok(Token::AddAssign)));
    assert_eq!(lexer.next(), Some(Ok(Token::SubAssign)));
    assert_eq!(lexer.next(), Some(Ok(Token::MulAssign)));
    assert_eq!(lexer.next(), Some(Ok(Token::DivAssign)));
    assert_eq!(lexer.next(), Some(Ok(Token::ModAssign)));
    assert_eq!(lexer.next(), None);
}

#[test]
fn test_punctuation() {
    let mut lexer = Token::lexer("( ) [ ] { } , ; : . ... .. :: -> ~ ! => | |> ?. ?");

    assert_eq!(lexer.next(), Some(Ok(Token::LParen)));
    assert_eq!(lexer.next(), Some(Ok(Token::RParen)));
    assert_eq!(lexer.next(), Some(Ok(Token::LBracket)));
    assert_eq!(lexer.next(), Some(Ok(Token::RBracket)));
    assert_eq!(lexer.next(), Some(Ok(Token::LBrace)));
    assert_eq!(lexer.next(), Some(Ok(Token::RBrace)));
    assert_eq!(lexer.next(), Some(Ok(Token::Comma)));
    assert_eq!(lexer.next(), Some(Ok(Token::Semicolon)));
    assert_eq!(lexer.next(), Some(Ok(Token::Colon)));
    assert_eq!(lexer.next(), Some(Ok(Token::Dot)));
    assert_eq!(lexer.next(), Some(Ok(Token::Spread)));
    assert_eq!(lexer.next(), Some(Ok(Token::DotDot)));
    assert_eq!(lexer.next(), Some(Ok(Token::Access)));
    assert_eq!(lexer.next(), Some(Ok(Token::Arrow)));
    assert_eq!(lexer.next(), Some(Ok(Token::Tilde)));
    assert_eq!(lexer.next(), Some(Ok(Token::Bang)));
    assert_eq!(lexer.next(), Some(Ok(Token::FatArrow)));
    assert_eq!(lexer.next(), Some(Ok(Token::Union)));
    assert_eq!(lexer.next(), Some(Ok(Token::OptionalChain)));
    assert_eq!(lexer.next(), Some(Ok(Token::Unwrap)));
    assert_eq!(lexer.next(), None);
}

#[test]
fn test_variables() {
    let mut lexer = Token::lexer("x y_z _private camelCase PascalCase");

    assert_eq!(lexer.next(), Some(Ok(Token::Variable("x".to_string()))));
    assert_eq!(lexer.next(), Some(Ok(Token::Variable("y_z".to_string()))));
    assert_eq!(
        lexer.next(),
        Some(Ok(Token::Variable("_private".to_string())))
    );
    assert_eq!(
        lexer.next(),
        Some(Ok(Token::Variable("camelCase".to_string())))
    );
    assert_eq!(
        lexer.next(),
        Some(Ok(Token::Variable("PascalCase".to_string())))
    );
    assert_eq!(lexer.next(), None);
}

#[test]
fn test_whitespace_skipping() {
    let mut lexer = Token::lexer("  \t\n\r  true  \n  false  ");

    assert_eq!(lexer.next(), Some(Ok(Token::Bool(true))));
    assert_eq!(lexer.next(), Some(Ok(Token::Bool(false))));
    assert_eq!(lexer.next(), None);
}

#[test]
fn test_comment_skipping() {
    let mut lexer = Token::lexer("true # this is a comment\n false");

    assert_eq!(lexer.next(), Some(Ok(Token::Bool(true))));
    assert_eq!(lexer.next(), Some(Ok(Token::Bool(false))));
    assert_eq!(lexer.next(), None);
}

#[test]
fn test_complex_sequence() {
    let mut lexer = Token::lexer("fn add(x: int, y: int) -> int { x + y }");

    assert_eq!(lexer.next(), Some(Ok(Token::KeywordFn)));
    assert_eq!(lexer.next(), Some(Ok(Token::Variable("add".to_string()))));
    assert_eq!(lexer.next(), Some(Ok(Token::LParen)));
    assert_eq!(lexer.next(), Some(Ok(Token::Variable("x".to_string()))));
    assert_eq!(lexer.next(), Some(Ok(Token::Colon)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordInt)));
    assert_eq!(lexer.next(), Some(Ok(Token::Comma)));
    assert_eq!(lexer.next(), Some(Ok(Token::Variable("y".to_string()))));
    assert_eq!(lexer.next(), Some(Ok(Token::Colon)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordInt)));
    assert_eq!(lexer.next(), Some(Ok(Token::RParen)));
    assert_eq!(lexer.next(), Some(Ok(Token::Arrow)));
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordInt)));
    assert_eq!(lexer.next(), Some(Ok(Token::LBrace)));
    assert_eq!(lexer.next(), Some(Ok(Token::Variable("x".to_string()))));
    assert_eq!(lexer.next(), Some(Ok(Token::Plus)));
    assert_eq!(lexer.next(), Some(Ok(Token::Variable("y".to_string()))));
    assert_eq!(lexer.next(), Some(Ok(Token::RBrace)));
    assert_eq!(lexer.next(), None);
}

#[test]
fn test_edge_cases() {
    // Test that keywords are not treated as variables
    let mut lexer = Token::lexer("let let_var if if_var");

    assert_eq!(lexer.next(), Some(Ok(Token::KeywordLet)));
    assert_eq!(
        lexer.next(),
        Some(Ok(Token::Variable("let_var".to_string())))
    );
    assert_eq!(lexer.next(), Some(Ok(Token::KeywordIf)));
    assert_eq!(
        lexer.next(),
        Some(Ok(Token::Variable("if_var".to_string())))
    );
    assert_eq!(lexer.next(), None);
}