ltk_ritobin 0.12.0

use std::fmt::Display;

use crate::span::Span;

#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
#[rustfmt::skip]
pub enum TokenKind {
  Unknown, Eof, Newline,

  LParen, RParen, LCurly, RCurly,
  LBrack, RBrack,
  Eq, Comma, Colon, SemiColon,
  Star, Slash, At,
  Quote,

  String, UnterminatedString,

  Comment,

  // keywords
  True, False, Null,

  Name, Number, HexLit,
}

impl TokenKind {
    /// Whether we are a string/unterminated string
    pub fn is_string(&self) -> bool {
        matches!(self, Self::String | Self::UnterminatedString)
    }

    /// Maps a punctuation byte -> its token kind, inverse of [`Self::print_value`] for single-byte punctuation
    pub fn from_punct_byte(b: u8) -> Option<Self> {
        Some(match b {
            b'(' => Self::LParen,
            b')' => Self::RParen,
            b'{' => Self::LCurly,
            b'}' => Self::RCurly,
            b'[' => Self::LBrack,
            b']' => Self::RBrack,
            b'=' => Self::Eq,
            b',' => Self::Comma,
            b':' => Self::Colon,
            b';' => Self::SemiColon,
            b'*' => Self::Star,
            b'/' => Self::Slash,
            b'@' => Self::At,
            _ => return None,
        })
    }

    pub fn print_value(&self) -> Option<&'static str> {
        match self {
            TokenKind::Unknown => None,
            TokenKind::Eof => None,
            TokenKind::Newline => None,
            TokenKind::LParen => Some("("),
            TokenKind::RParen => Some(")"),
            TokenKind::LCurly => Some("{"),
            TokenKind::RCurly => Some("}"),
            TokenKind::LBrack => Some("["),
            TokenKind::RBrack => Some("]"),
            TokenKind::Eq => Some("="),
            TokenKind::Comma => Some(","),
            TokenKind::Colon => Some(":"),
            TokenKind::SemiColon => Some(";"),
            TokenKind::Star => Some("*"),
            TokenKind::Slash => Some("/"),
            TokenKind::At => Some("@"),
            TokenKind::Quote => Some("\""),
            TokenKind::String => None,
            TokenKind::UnterminatedString => None,
            TokenKind::Comment => None,
            TokenKind::True => Some("true"),
            TokenKind::False => Some("false"),
            TokenKind::Null => Some("null"),
            TokenKind::Name => None,
            TokenKind::Number => None,
            TokenKind::HexLit => None,
        }
    }
}

impl Display for TokenKind {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            TokenKind::Unknown => "unknown text",
            TokenKind::Eof => "end of file",
            TokenKind::Newline => "new line",
            TokenKind::LParen => "'('",
            TokenKind::RParen => "')'",
            TokenKind::LCurly => "'{'",
            TokenKind::RCurly => "'}'",
            TokenKind::LBrack => "'['",
            TokenKind::RBrack => "']'",
            TokenKind::Eq => "'='",
            TokenKind::Comma => "','",
            TokenKind::Colon => "':'",
            TokenKind::SemiColon => "';'",
            TokenKind::Star => "'*'",
            TokenKind::Slash => "'/'",
            TokenKind::At => "'@'",
            TokenKind::Quote => "'\"'",
            TokenKind::String => "string literal",
            TokenKind::UnterminatedString => "unterminated string literal",
            TokenKind::True => "'true'",
            TokenKind::False => "'false'",
            TokenKind::Null => "'null'",
            TokenKind::Name => "keyword",
            TokenKind::Number => "number",
            TokenKind::HexLit => "hexadecimal literal",
            TokenKind::Comment => "comment",
        })
    }
}

#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Token {
    pub kind: TokenKind,
    pub span: Span,
}
pub fn lex(source: &str) -> Vec<Token> {
    use TokenKind::*;

    // real bins run ~6.82 (hashes) to ~7.25 (resolved) bytes/token; /6 stays under both so big files don't realloc
    let mut result = Vec::with_capacity(source.len() / 6);
    let mut cur = Cursor::new(source);

    while !cur.at_end() {
        let start = cur.pos;

        // whitespace: eat the run, emitting a synthetic Newline between value-ending tokens
        // when it spans a line break
        if cur.consume_while(|b| b.is_ascii_whitespace()) > 0 {
            if ends_line(result.last(), cur.consumed(start)) {
                result.push(Token {
                    kind: Newline,
                    span: cur.span(start),
                });
            }

            continue;
        }

        let mut kind = cur.consume_token();
        assert!(cur.pos > start);

        if kind == Name {
            kind = keyword_or_name(&source[start..cur.pos]);
        }

        result.push(Token {
            kind,
            span: cur.span(start),
        });
    }

    result
}

/// A byte cursor over the source. `pos` is always on a UTF-8 char boundary: every token
/// ends on an ASCII byte (or swallows whole multi-byte sequences), so slicing at `pos`
/// never splits a char.
struct Cursor<'a> {
    bytes: &'a [u8],
    pos: usize,
}

impl<'a> Cursor<'a> {
    fn new(source: &'a str) -> Self {
        Cursor {
            bytes: source.as_bytes(),
            pos: 0,
        }
    }

    fn at_end(&self) -> bool {
        self.pos >= self.bytes.len()
    }

    fn rest(&self) -> &'a [u8] {
        &self.bytes[self.pos..]
    }

    fn consumed(&self, start: usize) -> &'a [u8] {
        &self.bytes[start..self.pos]
    }

    fn span(&self, start: usize) -> Span {
        Span::new(start as u32, self.pos as u32)
    }

    fn starts_with(&self, prefix: &[u8]) -> bool {
        self.rest().starts_with(prefix)
    }

    #[inline]
    fn consume_while(&mut self, pred: impl Fn(&u8) -> bool) -> usize {
        let n = scan_run(self.rest(), pred);
        self.pos += n;
        n
    }

    #[inline]
    fn consume_token(&mut self) -> TokenKind {
        use TokenKind::*;
        let b = self.bytes[self.pos];

        if let Some(k) = TokenKind::from_punct_byte(b) {
            self.pos += 1;
            return k;
        }

        if b == b'#' {
            self.pos += 1;
            self.consume_while(|b| !matches!(b, b'\n' | b'\r'));
            return Comment;
        }

        if self.starts_with(b"0x") {
            self.pos += 2;
            self.consume_while(|b| b.is_ascii_hexdigit());
            return HexLit;
        }

        if let Some(len) = scan_number(self.rest()) {
            self.pos += len;
            return Number;
        }

        if matches!(b, b'\'' | b'"') {
            return self.consume_string();
        }

        if self.consume_while(is_name_char) > 0 {
            return Name;
        }

        // unrecognized: consume up to the next whitespace
        self.consume_while(|b| !b.is_ascii_whitespace());
        Unknown
    }

    #[inline]
    fn consume_string(&mut self) -> TokenKind {
        use TokenKind::*;
        let quote = self.bytes[self.pos];
        self.pos += 1;

        let mut escaped = false;
        while let Some(&b) = self.bytes.get(self.pos) {
            self.pos += 1;
            match b {
                b'\n' | b'\r' => return UnterminatedString,

                _ if escaped => escaped = false,
                b'\\' => escaped = true,
                _ if b == quote => return String,
                _ => {}
            }
        }

        UnterminatedString
    }
}

fn is_name_char(b: &u8) -> bool {
    matches!(b, b'_' | b'a'..=b'z' | b'A'..=b'Z' | b'0'..=b'9')
}

fn keyword_or_name(word: &str) -> TokenKind {
    use TokenKind::*;
    match word {
        "true" => True,
        "false" => False,
        "null" => Null,
        _ => Name,
    }
}

fn ends_value(kind: TokenKind) -> bool {
    use TokenKind::*;
    matches!(
        kind,
        Name | HexLit | True | False | Number | Null | RCurly | String | Eq | Comment | At
    )
}

fn ends_line(last: Option<&Token>, ws: &[u8]) -> bool {
    last.is_some_and(|t| ends_value(t.kind) || matches!(t.kind, TokenKind::LParen))
        && ws.iter().any(|&b| matches!(b, b'\n' | b'\r'))
}

#[inline]
fn scan_run(bytes: &[u8], pred: impl Fn(&u8) -> bool) -> usize {
    bytes.iter().position(|b| !pred(b)).unwrap_or(bytes.len())
}

#[inline]
fn scan_number_segment(bytes: &[u8]) -> Option<usize> {
    let len = scan_run(bytes, |b| matches!(b, b'0'..=b'9' | b'_'));
    bytes[..len]
        .split(|&b| b == b'_')
        .all(|group| !group.is_empty())
        .then_some(len)
}

#[inline]
fn scan_exponent(bytes: &[u8]) -> Option<usize> {
    if !matches!(bytes.first(), Some(b'e' | b'E')) {
        return Some(0);
    }
    let sign = 1 + matches!(bytes.get(1), Some(b'+' | b'-')) as usize;

    Some(match scan_run(&bytes[sign..], u8::is_ascii_digit) {
        // must be at least 1 digit
        0 => return None,
        digits => sign + digits,
    })
}

#[inline]
fn scan_number(bytes: &[u8]) -> Option<usize> {
    let starts_with_minus = bytes.first() == Some(&b'-');
    let mut i = starts_with_minus as usize;

    // check for inf/nan, only if there's no minus, or a minus as first char
    if matches!(i, 0 | 1) {
        let end = i + 3;
        if bytes.len() >= end {
            // 0 1 2 3 4
            // - i n f
            for word in [b"inf", b"nan"] {
                if bytes[i..end].eq_ignore_ascii_case(word)
                    && bytes.get(end).is_none_or(|c| !is_name_char(c))
                {
                    return Some(end);
                }
            }
        }
    }

    // ".5"
    if bytes.get(i) == Some(&b'.') {
        i += 1;
        i += scan_number_segment(&bytes[i..])?;
        i += scan_exponent(&bytes[i..])?;
        return Some(i);
    }

    // "123" / "123.45"
    i += scan_number_segment(&bytes[i..])?;
    if bytes.get(i) == Some(&b'.') {
        i += 1;
        i += scan_number_segment(&bytes[i..])?;
    }

    Some(i + scan_exponent(&bytes[i..])?)
}

#[cfg(test)]
mod tests {
    use crate::parse::tokenizer::scan_number;

    #[test]
    fn number_scan() {
        assert_eq!(scan_number(b"1234"), Some(4));
        assert_eq!(scan_number(b"-1234"), Some(5));
        assert_eq!(scan_number(b"-1234.212"), Some(9));

        assert_eq!(scan_number(b"-1234.212e"), None);
        assert_eq!(scan_number(b"-1234e120"), Some(9));
        assert_eq!(scan_number(b"-1234.212E+120"), Some(14));
        assert_eq!(scan_number(b"-1234.212e-120"), Some(14));

        assert_eq!(scan_number(b".5e10"), Some(5));

        assert_eq!(scan_number(b"NaN"), Some(3));
        assert_eq!(scan_number(b"-nan"), Some(4));

        assert_eq!(scan_number(b"Inf"), Some(3));
        assert_eq!(scan_number(b"-inf"), Some(4));
    }
}