endorphin-aster-vm/crates/frontend/src/lexer.rs

use crate::{Diagnostic, Span, Token, TokenKind};

/// Lex a reader snapshot and preserve provenance in any diagnostic.
pub fn lex_source(source: &crate::SourceUnit) -> Result<Vec<Token>, crate::SourceDiagnostic> {
    lex(source.text()).map_err(|error| source.annotate(error))
}

/// Lex the basic Endorphin draft syntax without reading files or changing
/// locale. Positions always refer to bytes in this exact UTF-8 source.
pub fn lex(source: &str) -> Result<Vec<Token>, Diagnostic> {
    let mut lexer = Lexer { source, offset: 0 };
    let mut tokens = Vec::new();
    while let Some(character) = lexer.peek() {
        if matches!(character, ' ' | '\t' | '\r' | '\n') {
            let _ = lexer.advance();
            continue;
        }
        if lexer.remaining().starts_with("//") {
            while lexer
                .peek()
                .is_some_and(|character| !matches!(character, '\r' | '\n'))
            {
                let _ = lexer.advance();
            }
            continue;
        }
        let start = lexer.offset;
        let kind = match character {
            character if name_start(character) => {
                let _ = lexer.advance();
                while lexer.peek().is_some_and(name_continue) {
                    let _ = lexer.advance();
                }
                TokenKind::Name
            }
            '0'..='9' => lexer.number()?,
            '"' => lexer.string()?,
            _ => {
                let _ = lexer.advance();
                match character {
                    '(' => TokenKind::LeftParen,
                    ')' => TokenKind::RightParen,
                    '{' => TokenKind::LeftBrace,
                    '}' => TokenKind::RightBrace,
                    ':' => TokenKind::Colon,
                    ',' => TokenKind::Comma,
                    '.' => TokenKind::Dot,
                    '=' => TokenKind::Equal,
                    '+' => TokenKind::Plus,
                    '-' => TokenKind::Minus,
                    '*' => TokenKind::Star,
                    '/' => TokenKind::Slash,
                    _ => {
                        return Err(Diagnostic::lexical(
                            Span::new(start, lexer.offset),
                            "unexpected source character",
                        ));
                    }
                }
            }
        };
        tokens.push(Token {
            kind,
            span: Span::new(start, lexer.offset),
        });
    }
    Ok(tokens)
}

struct Lexer<'source> {
    source: &'source str,
    offset: usize,
}

impl<'source> Lexer<'source> {
    fn remaining(&self) -> &'source str {
        &self.source[self.offset..]
    }

    fn peek(&self) -> Option<char> {
        self.remaining().chars().next()
    }

    fn advance(&mut self) -> Option<char> {
        let character = self.peek()?;
        self.offset += character.len_utf8();
        Some(character)
    }

    fn digits(&mut self) {
        while self
            .peek()
            .is_some_and(|character| character.is_ascii_digit())
        {
            let _ = self.advance();
        }
    }

    fn number(&mut self) -> Result<TokenKind, Diagnostic> {
        let start = self.offset;
        self.digits();
        let mut following = self.remaining().chars();
        if following.next() == Some('.')
            && following
                .next()
                .is_some_and(|character| character.is_ascii_digit())
        {
            let _ = self.advance();
            self.digits();
        }
        if self.peek().is_some_and(name_continue) {
            while self.peek().is_some_and(name_continue) {
                let _ = self.advance();
            }
            return Err(Diagnostic::lexical(
                Span::new(start, self.offset),
                "unsupported numeric suffix or missing token separator",
            ));
        }
        Ok(TokenKind::Number)
    }

    fn string(&mut self) -> Result<TokenKind, Diagnostic> {
        let start = self.offset;
        let _ = self.advance();
        loop {
            let position = self.offset;
            match self.advance() {
                Some('"') => return Ok(TokenKind::String),
                Some('\\') => match self.advance() {
                    Some('n' | 't' | '"' | '\\') => {}
                    Some(_) => {
                        return Err(Diagnostic::lexical(
                            Span::new(position, self.offset),
                            "unsupported string escape",
                        ));
                    }
                    None => {
                        return Err(Diagnostic::lexical(
                            Span::new(start, self.offset),
                            "unterminated string",
                        ));
                    }
                },
                Some('\r' | '\n' | '\0') => {
                    return Err(Diagnostic::lexical(
                        Span::new(position, self.offset),
                        "unescaped newline or NUL in string",
                    ));
                }
                None => {
                    return Err(Diagnostic::lexical(
                        Span::new(start, self.offset),
                        "unterminated string",
                    ));
                }
                _ => {}
            }
        }
    }
}

fn name_start(character: char) -> bool {
    character == '_' || character.is_alphabetic()
}

fn name_continue(character: char) -> bool {
    character == '_' || character.is_alphanumeric()
}

pub(crate) fn decode_string(raw: &str) -> String {
    let mut result = String::new();
    let mut characters = raw.chars();
    while let Some(character) = characters.next() {
        if character == '\\' {
            result.push(match characters.next().expect("lexer validated escape") {
                'n' => '\n',
                't' => '\t',
                '"' => '"',
                '\\' => '\\',
                _ => unreachable!("lexer validated escape"),
            });
        } else {
            result.push(character);
        }
    }
    result
}