rusty-minic/src/lexer/lexer.rs

use std::{str::FromStr, sync::LazyLock};

use thiserror::Error;

use crate::{diagnostic::{Diagnositics, span::Span}, lexer::{err::LexingError, types::{TokenValue, TypeIdent}}};

use super::types::Token;

pub struct Lexer {
    pub tokens: Vec<Token>,
    pub diagnostics: Diagnositics,
    old_char_count: usize,
    block_comment_span: Option<Span>,
    in_skip_line: bool
}

const WHITESPACE_CHARS: &[char] = &[' ', '\t', '\n', '\r'];
const DELIMITER_CHARS: &[char] = &[
    '+', '-', '*', '/', '%', '=', '!', '<', '>', '(', ')', '[', ']', ',', ';', '{', '}', '|', '&'
];
struct Cursor {
    chars: Vec<char>,
    pos: usize,
}
enum LexParseError {
    NotMatched,
    InvalidInMatch(LexingError)
}
impl Cursor {
    pub fn new(s: &str) -> Self {
        Self { chars: s.chars().collect(), pos: 0 }
    }
    fn peek(&self) -> Option<char> {
        self.chars.get(self.pos).copied()
    }
    fn peek_multiple(&self, n: usize) -> Option<&[char]> {
        if self.pos + n <= self.chars.len() {
            Some(&self.chars[self.pos..self.pos + n])
        } else {
            None
        }
    }
    fn advance(&mut self, n: usize) {
        self.pos += n;
    }
    fn next(&mut self) -> Option<char> {
        let c = self.chars.get(self.pos).copied();
        if c.is_some() {
            self.advance(1);
        }
        c
    }
    fn back(&mut self, n: usize) {
        self.pos = self.pos.saturating_sub(n);
    }
    fn pos(&self) -> usize {
        self.pos
    }
}
/// try parse using the giving function, return whether should continue
fn try_parse_as(
    f: fn(&mut Cursor) -> Result<TokenValue, LexParseError>,
    tokens: &mut Vec<Token>,
    str_iter: &mut Cursor,
    diagnostics: &mut Diagnositics,
    last_char_count: usize,
) -> bool {
    let last_pos = str_iter.pos() + last_char_count;
    match f(str_iter) {
        Ok(token_value) => {
            let span = Span { start: last_pos, end: str_iter.pos() + last_char_count };
            tokens.push(Token { value: token_value, span });
            return true;
        }
        Err(LexParseError::NotMatched) => false,
        Err(LexParseError::InvalidInMatch(err)) => {
            // try recover from delimiter char or whitespace char
            while let Some(c) = str_iter.peek() {
                if DELIMITER_CHARS.contains(&c) || WHITESPACE_CHARS.contains(&c) {
                    break;
                }
                str_iter.advance(1);
            }
            let span = Span { start: last_pos, end: str_iter.pos() + last_char_count };
            diagnostics.add_from_frontend_error(err, span);
            return true;
        }
    }
}
macro_rules! if_true_then_continue {
    ($e: expr) => {
        if $e {
            continue;
        }
    };
}
#[derive(Debug, Error)]
pub enum LexerError {
    #[error("Io error: {0}")]
    Io(#[from] std::io::Error),
    #[error("Too much errors, stop lexing")]
    TooManyErrors,
}
impl Lexer {
    pub fn new() -> Self {
        Self { tokens: vec![], diagnostics: Diagnositics::new(), old_char_count: 0, block_comment_span: None, in_skip_line: false }
    }
    pub fn finish(mut self) -> (Vec<Token>, Diagnositics) {
        if let Some(span) = self.block_comment_span.take() {
            self.diagnostics.add_from_frontend_error(LexingError::UnterminatedComment, span);
        }
        (self.tokens, self.diagnostics)
    }
    /// call `parse_str` will continue to parse the input from current state
    /// please also pass the whitespace to ensure the correct char position in diagnostics
    pub fn parse_next_str(&mut self, s: &str) {
        let mut cursor = Cursor::new(s);
        loop {
            if let Some(c) = cursor.peek() {
                if self.in_skip_line && c != '\n' {
                    cursor.advance(1);
                    continue;
                }
                // check white space first, if it's white space, skip it and continue to the next character
                if WHITESPACE_CHARS.contains(&c) {
                    if c == '\n' {
                        self.in_skip_line = false;
                    }
                    cursor.advance(1);
                    continue;
                }
                // check comment
                match cursor.peek_multiple(2) {
                    Some(['/', '/']) => {
                        // skip the rest of the line
                        self.in_skip_line = true;
                        cursor.advance(2);
                        continue;
                    }
                    Some(['/', '*']) => {
                        let start = cursor.pos() + self.old_char_count;
                        self.block_comment_span = Some(Span { start, end: start + 2 });
                        cursor.advance(2);
                        continue;
                    }
                    Some(['*', '/']) => {
                        self.block_comment_span = None;
                        cursor.advance(2);
                        continue;
                    }
                    _ => {}
                }
                if self.block_comment_span.is_some() {
                    cursor.advance(1);
                    continue;
                }
            } else {
                break;
            }
            if_true_then_continue!(try_parse_as(parse_litint, &mut self.tokens, &mut cursor, &mut self.diagnostics, self.old_char_count));
            if_true_then_continue!(try_parse_as(parse_delimiter, &mut self.tokens, &mut cursor, &mut self.diagnostics, self.old_char_count));
            if_true_then_continue!(try_parse_as(parse_puncuation, &mut self.tokens, &mut cursor, &mut self.diagnostics, self.old_char_count));
            if_true_then_continue!(try_parse_as(parse_ident, &mut self.tokens, &mut cursor, &mut self.diagnostics, self.old_char_count));
            // unrecognized token
            let last_pos = cursor.pos() + self.old_char_count;
            let mut unrecognized = Vec::new();
            while let Some(c) = cursor.peek() {
                if DELIMITER_CHARS.contains(&c) || WHITESPACE_CHARS.contains(&c) {
                    break;
                }
                unrecognized.push(c);
                cursor.advance(1);
            }
            let span = Span { start: last_pos, end: cursor.pos() + self.old_char_count };
            let unrecognized = unrecognized.into_iter().collect::<String>();
            self.diagnostics.add_from_frontend_error(LexingError::UnrecognizedToken(unrecognized), span);
            self.tokens.push(Token { value: TokenValue::Unrecognized, span });
        }
        self.old_char_count += s.len();
    }
}
fn parse_litint(
    str_iter: &mut Cursor,
) -> Result<TokenValue, LexParseError> {
    let mut c1 = str_iter.peek().ok_or(LexParseError::NotMatched)?;
    // c1 is the peek value from here
    let mut base: i64 = 10;
    if !(c1.is_ascii_digit()) {
        return Err(LexParseError::NotMatched);
    }
    let mut number = 0i64;
    let mut has_digits = false;
    if c1 == '0' {
        str_iter.advance(1);
        match str_iter.peek() {
            Some('x') | Some('X') => {
                base = 16;
                str_iter.advance(1);
            }
            Some(c) if c.is_ascii_digit() => {
                base = 8;
            }
            _ => {
                has_digits = true;
                // only zero
            }
        }
    }
    // from here, the cursor points to:
    // 0x1234 -> cursor at '1'
    // 0123 -> cursor at '1'
    // 0 -> cursor at end
    // 1234 -> cursor at '1'
    loop {
        let c = match str_iter.peek() {
            Some(c) => c,
            None => break,
        };
        let digit = match c {
            '0'..='9' if (c as u8 - b'0') < base as u8 => c as i64 - '0' as i64 ,
            'a'..='f' if base == 16 => c as i64 - 'a' as i64 + 10 ,
            'A'..='F' if base == 16 => c as i64 - 'A' as i64 + 10,
            c => if WHITESPACE_CHARS.contains(&c) || DELIMITER_CHARS.contains(&c) {
                break;
            } else {
                // unrecognized character in number literal
                return Err(LexParseError::InvalidInMatch(LexingError::InvalidIntLiteral));
            }
        };
        has_digits = true;
        number = number * base + digit;
        str_iter.advance(1);
    }
    if !has_digits {
        // No valid digits found, add a diagnostic
        return Err(LexParseError::InvalidInMatch(LexingError::InvalidIntLiteral));
    }
    Ok(TokenValue::IntLit(number))
}

fn parse_delimiter(
    str_iter: &mut Cursor,
) -> Result<TokenValue, LexParseError> {
    let c = str_iter.peek().ok_or(LexParseError::NotMatched)?;
    let token_value = match c {
        '(' => TokenValue::LParen,
        ')' => TokenValue::RParen,
        '{' => TokenValue::LBrace,
        '}' => TokenValue::RBrace,
        '[' => TokenValue::LBracket,
        ']' => TokenValue::RBracket,
        _ => return Err(LexParseError::NotMatched),
    };
    str_iter.advance(1);
    Ok(token_value)
}
fn parse_puncuation(
    str_iter: &mut Cursor,
) -> Result<TokenValue, LexParseError> {
    let get_value_by_next_char = 
        |str_iter: &mut Cursor, not_equal_value: TokenValue, equal_value: TokenValue| {
            str_iter.advance(1);
            if let Some('=') = str_iter.peek() {
                equal_value
            } else {
                str_iter.back(1);
                not_equal_value
            }
    };
    let c = str_iter.peek().ok_or(LexParseError::NotMatched)?;
    let token_value = match c {
        '+' => {
            str_iter.advance(1);
            if let Some('+') = str_iter.peek() {
                TokenValue::PlusPlus
            } else if let Some('=') = str_iter.peek() {
                TokenValue::PlusEqual
            } else {                
                str_iter.back(1);
                TokenValue::Plus
            }
        },
        '-' => {
            str_iter.advance(1);
            if let Some('-') = str_iter.peek() {
                TokenValue::MinusMinus
            } else if let Some('=') = str_iter.peek() {
                TokenValue::MinusEqual
            } else {                
                str_iter.back(1);
                TokenValue::Minus
            }
        },
        '*' => {
            str_iter.advance(1);
            if let Some('=') = str_iter.peek() {
                TokenValue::StarEqual
            } else {
                str_iter.back(1);
                TokenValue::Star
            }
        },
        '/' => {
            str_iter.advance(1);
            if let Some('=') = str_iter.peek() {
                TokenValue::SlashEqual
            } else {
                str_iter.back(1);
                TokenValue::Slash
            }
        },
        '%' => {
            str_iter.advance(1);
            if let Some('=') = str_iter.peek() {
                TokenValue::PercentEqual
            } else {
                str_iter.back(1);
                TokenValue::Percent
            }
        },

        '=' => get_value_by_next_char(str_iter, TokenValue::Equal, TokenValue::DoubleEqual),
        '!' => {
            str_iter.advance(1);
            if let Some('=') = str_iter.peek() {
                TokenValue::NotEqual
            } else {
                str_iter.back(1);
                TokenValue::Not
            }
        },
        '<' => get_value_by_next_char(str_iter, TokenValue::Less, TokenValue::LessEqual),
        '>' => get_value_by_next_char(str_iter, TokenValue::Greater, TokenValue::GreaterEqual),

        ',' => TokenValue::Comma,
        ';' => TokenValue::Semicolon,
        '|' => {
            str_iter.advance(1);
            if let Some('|') = str_iter.peek() {
                TokenValue::Or
            } else {
                // unrecognized token starting with '|'
                return Err(LexParseError::InvalidInMatch(LexingError::UnrecognizedToken("|".to_string())));
            }
        },
        '&' => {
            str_iter.advance(1);
            if let Some('&') = str_iter.peek() {
                TokenValue::And
            } else {
                // unrecognized token starting with '&'
                return Err(LexParseError::InvalidInMatch(LexingError::UnrecognizedToken("&".to_string())));
            }
        },
        _ => return Err(LexParseError::NotMatched),
    };
    str_iter.advance(1);
    Ok(token_value)
}

// static KEYWORDS: LazyLock<HashMap<&'static str, TokenValue>> = LazyLock::new(|| {
//     let mut m = HashMap::new();
//     m.insert("if", TokenValue::If);
//     m.insert("else", TokenValue::Else);
//     m.insert("while", TokenValue::While);
//     m.insert("for", TokenValue::For);
//     m.insert("return", TokenValue::Return);
//     m.insert("break", TokenValue::Break);
//     m.insert("continue", TokenValue::Continue);
//     m
// });
fn parse_ident(
    str_iter: &mut Cursor,
) -> Result<TokenValue, LexParseError> {
    let c = str_iter.peek().ok_or(LexParseError::NotMatched)?;
    if !c.is_ascii_alphabetic() && c != '_' {
        return Err(LexParseError::NotMatched);
    }
    let mut name = Vec::new();
    while let Some(c) = str_iter.peek() {
        if c.is_ascii_alphanumeric() || c == '_' {
            name.push(c);
            str_iter.advance(1);
        } else if DELIMITER_CHARS.contains(&c) || WHITESPACE_CHARS.contains(&c) {
            break;
        } else {
            return Err(LexParseError::InvalidInMatch(LexingError::InvalidIdent));
        }
    }
    let name = name.into_iter().collect::<String>();
    if name.eq("return") {
        return Ok(TokenValue::Return);
    }
    if name.eq("if") {
        return Ok(TokenValue::If);
    }
    if name.eq("else") {
        return Ok(TokenValue::Else);
    }
    if name.eq("while") {
        return Ok(TokenValue::While);
    }
    if name.eq("for") {
        return Ok(TokenValue::For);
    }
    if name.eq("break") {
        return Ok(TokenValue::Break);
    }
    if name.eq("continue") {
        return Ok(TokenValue::Continue);
    }
    if name.eq("static") {
        return Ok(TokenValue::Static);
    }
    if let Some(type_ident) = TypeIdent::from_str(&name).ok() {
        return Ok(TokenValue::TypeIdent(type_ident));
    }
    Ok(TokenValue::Ident(name))
}
#[cfg(test)]
mod tests {
    use std::io::BufRead;
    use std::path::Path;
    use std::fs::File;
    use crate::utils::case_list::CaseList;
    use crate::utils::num_sequence::NumberSequence;

    pub use super::*;
    fn test_case(case_str: &str) {
        let case_sequence = NumberSequence::from_str(case_str).unwrap();
        let case_list = CaseList::from_dir(&Path::new("./testcases")).unwrap();
        let mut error_case_cnt = 0;
        for case_no in case_sequence {
            let case_path = case_list.get_case_path(case_no).unwrap();
            println!("{}", case_path.display());
            let file = File::open(&case_path).unwrap();
            let mut buf_reader = std::io::BufReader::new(file);
            let mut lexer = Lexer::new();
            let mut full_text = String::new();
            loop {
                let mut line = String::new();
                let bytes_read = buf_reader.read_line(&mut line).unwrap();
                if bytes_read == 0 {
                    break;
                }
                full_text.push_str(&line);
                lexer.parse_next_str(&line);
            }
            let (_tokens, diagnostics) = lexer.finish();
            if !diagnostics.is_empty() {
                diagnostics.print(&format!("{}", case_path.display()), &full_text);
                error_case_cnt += 1;
            }
        }
        if error_case_cnt > 0 {
            panic!("Found {} cases with errors", error_case_cnt);
        }

    }
    #[test]
    fn test_expr() {
        test_case("0-3,14-25");
    }
}