AILang/crates/ailang-surface/src/lex.rs

//! Lexer for form (A).
//!
//! The lexical layer has three rules:
//!
//! 1. Whitespace (` `, `\t`, `\n`, `\r`) and parens (`(`, `)`) are the
//!    only delimiters. Any maximal non-whitespace, non-paren run is a
//!    single token.
//! 2. `;` introduces a comment to end-of-line. Comments produce no
//!    tokens.
//! 3. After tokenisation, classify by first character:
//!    - `"` ⇒ string literal (consume until closing `"`; `\"` and
//!      `\n` escapes supported).
//!    - digit, or `-` followed by digit ⇒ integer literal.
//!    - otherwise ⇒ ident.
//!
//! Operators (`+`, `==`, `<=`, `**`), qualified names like
//! `io/print_int`, and dotted cross-module references like
//! `std_list.map` are all single-token idents — there is no special
//! lexical rule for any of them. `(`, `)`, and whitespace are the only
//! reserved tokens. Bool literals (`true`, `false`) and unit
//! (`(lit-unit)`) are disambiguated by the parser via context.

use thiserror::Error;

/// Source byte offset (start, end] into the original input. `end` is
/// one past the last byte. Used for error reporting only.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Span {
    pub start: usize,
    pub end: usize,
}

/// One lexical token.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Tok {
    /// `(`
    LParen,
    /// `)`
    RParen,
    /// Integer literal. Original textual form preserved for error messages;
    /// parsed `i64` value carried for direct use.
    Int(i64),
    /// String literal contents (after escape processing).
    Str(String),
    /// Identifier. Anything that isn't a paren, integer, or string.
    Ident(String),
}

/// A token plus its span.
#[derive(Debug, Clone)]
pub struct Token {
    pub tok: Tok,
    pub span: Span,
}

/// Lexer error.
#[derive(Debug, Error)]
pub enum LexError {
    #[error("unterminated string literal at byte {start}")]
    UnterminatedString { start: usize },
    #[error("invalid integer literal {literal:?} at byte {start}")]
    InvalidInteger { literal: String, start: usize },
    #[error("invalid escape sequence \\{ch} in string literal at byte {pos}")]
    InvalidEscape { ch: char, pos: usize },
}

/// Tokenise an input string into a flat token stream.
///
/// Whitespace and comments are dropped. Returns a [`LexError`] on
/// unterminated string literals or numeric tokens that fail to parse.
pub fn tokenize(input: &str) -> Result<Vec<Token>, LexError> {
    let bytes = input.as_bytes();
    let mut out = Vec::new();
    let mut i = 0usize;
    while i < bytes.len() {
        let b = bytes[i];
        // Whitespace.
        if b == b' ' || b == b'\t' || b == b'\n' || b == b'\r' {
            i += 1;
            continue;
        }
        // Comment to end-of-line.
        if b == b';' {
            while i < bytes.len() && bytes[i] != b'\n' {
                i += 1;
            }
            continue;
        }
        // Parens.
        if b == b'(' {
            out.push(Token { tok: Tok::LParen, span: Span { start: i, end: i + 1 } });
            i += 1;
            continue;
        }
        if b == b')' {
            out.push(Token { tok: Tok::RParen, span: Span { start: i, end: i + 1 } });
            i += 1;
            continue;
        }
        // String literal. Iterate over UTF-8 chars (not raw bytes) so
        // multi-byte glyphs in the source survive round-trip.
        if b == b'"' {
            let start = i;
            i += 1;
            let mut value = String::new();
            // Walk via char_indices on the unread tail to keep i a byte offset.
            loop {
                if i >= bytes.len() {
                    return Err(LexError::UnterminatedString { start });
                }
                // Quick byte-level checks for ASCII control bytes that
                // can never start a multi-byte UTF-8 sequence.
                let c0 = bytes[i];
                if c0 == b'"' {
                    i += 1;
                    break;
                }
                if c0 == b'\\' {
                    if i + 1 >= bytes.len() {
                        return Err(LexError::UnterminatedString { start });
                    }
                    let esc = bytes[i + 1];
                    match esc {
                        b'"' => value.push('"'),
                        b'\\' => value.push('\\'),
                        b'n' => value.push('\n'),
                        b't' => value.push('\t'),
                        b'r' => value.push('\r'),
                        other => {
                            return Err(LexError::InvalidEscape {
                                ch: other as char,
                                pos: i,
                            });
                        }
                    }
                    i += 2;
                    continue;
                }
                // Decode one full UTF-8 character starting at i. The
                // input is guaranteed to be valid UTF-8 by Rust's
                // `&str` invariant, so a single `chars()` step on the
                // tail gives us the next scalar.
                if let Some(ch) = input[i..].chars().next() {
                    value.push(ch);
                    i += ch.len_utf8();
                } else {
                    return Err(LexError::UnterminatedString { start });
                }
            }
            out.push(Token {
                tok: Tok::Str(value),
                span: Span { start, end: i },
            });
            continue;
        }
        // Otherwise: maximal non-paren / non-whitespace / non-comment run.
        let start = i;
        while i < bytes.len() {
            let c = bytes[i];
            if c == b' '
                || c == b'\t'
                || c == b'\n'
                || c == b'\r'
                || c == b'('
                || c == b')'
                || c == b';'
            {
                break;
            }
            i += 1;
        }
        let raw = &input[start..i];
        // Classify.
        let first = raw.as_bytes()[0];
        let is_int = first.is_ascii_digit()
            || (first == b'-' && raw.len() > 1 && raw.as_bytes()[1].is_ascii_digit());
        if is_int {
            // Numeric parse must succeed; else parse error.
            match raw.parse::<i64>() {
                Ok(v) => out.push(Token {
                    tok: Tok::Int(v),
                    span: Span { start, end: i },
                }),
                Err(_) => {
                    return Err(LexError::InvalidInteger {
                        literal: raw.to_string(),
                        start,
                    });
                }
            }
        } else {
            out.push(Token {
                tok: Tok::Ident(raw.to_string()),
                span: Span { start, end: i },
            });
        }
    }
    Ok(out)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parens_and_idents() {
        let toks = tokenize("(foo bar)").unwrap();
        assert_eq!(toks.len(), 4);
        assert!(matches!(toks[0].tok, Tok::LParen));
        assert!(matches!(&toks[1].tok, Tok::Ident(s) if s == "foo"));
        assert!(matches!(&toks[2].tok, Tok::Ident(s) if s == "bar"));
        assert!(matches!(toks[3].tok, Tok::RParen));
    }

    #[test]
    fn comments_are_dropped() {
        let toks = tokenize("(a ; this is a comment\n b)").unwrap();
        assert_eq!(toks.len(), 4);
    }

    #[test]
    fn integers() {
        let toks = tokenize("42 -7 0").unwrap();
        assert_eq!(toks.len(), 3);
        assert!(matches!(toks[0].tok, Tok::Int(42)));
        assert!(matches!(toks[1].tok, Tok::Int(-7)));
        assert!(matches!(toks[2].tok, Tok::Int(0)));
    }

    #[test]
    fn operators_are_idents() {
        let toks = tokenize("+ == <= io/print_int std_list.map").unwrap();
        let names: Vec<_> = toks
            .iter()
            .map(|t| match &t.tok {
                Tok::Ident(s) => s.clone(),
                _ => panic!("not ident"),
            })
            .collect();
        assert_eq!(names, vec!["+", "==", "<=", "io/print_int", "std_list.map"]);
    }

    #[test]
    fn string_with_escapes() {
        let toks = tokenize(r#""hello\nworld""#).unwrap();
        assert_eq!(toks.len(), 1);
        assert!(matches!(&toks[0].tok, Tok::Str(s) if s == "hello\nworld"));
    }

    #[test]
    fn negative_lone_minus_is_ident() {
        // `-` alone (no digit follows) is an ident, not a (failed) integer.
        let toks = tokenize("-").unwrap();
        assert_eq!(toks.len(), 1);
        assert!(matches!(&toks[0].tok, Tok::Ident(s) if s == "-"));
    }

    #[test]
    fn string_preserves_multibyte_utf8() {
        let toks = tokenize(r#""em — dash and naïve""#).unwrap();
        assert_eq!(toks.len(), 1);
        assert!(matches!(&toks[0].tok, Tok::Str(s) if s == "em — dash and naïve"));
    }
}