tessellate-core 0.1.0

use crate::diagnostic::Diagnostic;
use crate::source::{SourceFile, Span};
use unicode_normalization::UnicodeNormalization;

/// A reserved word in the Tess surface language.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum Keyword {
    Module,
    Import,
    Enum,
    Entity,
    State,
    Derive,
    Decision,
    Action,
    Transition,
    On,
    Set,
    Rule,
    Case,
    Invariant,
    Trace,
    When,
    Then,
    Source,
    Override,
    Let,
    Expect,
    For,
    All,
    Some,
    If,
    Initially,
    Always,
    Terminates,
    Within,
    Steps,
    No,
    Dead,
    Ends,
    Exactly,
    One,
    Zero,
    Or,
    Many,
    Cardinality,
    Range,
    Domain,
    Optional,
    And,
    Not,
    True,
    False,
    Unknown,
    Continue,
}

impl Keyword {
    #[must_use]
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::Module => "module",
            Self::Import => "import",
            Self::Enum => "enum",
            Self::Entity => "entity",
            Self::State => "state",
            Self::Derive => "derive",
            Self::Decision => "decision",
            Self::Action => "action",
            Self::Transition => "transition",
            Self::On => "on",
            Self::Set => "set",
            Self::Rule => "rule",
            Self::Case => "case",
            Self::Invariant => "invariant",
            Self::Trace => "trace",
            Self::When => "when",
            Self::Then => "then",
            Self::Source => "source",
            Self::Override => "override",
            Self::Let => "let",
            Self::Expect => "expect",
            Self::For => "for",
            Self::All => "all",
            Self::Some => "some",
            Self::If => "if",
            Self::Initially => "initially",
            Self::Always => "always",
            Self::Terminates => "terminates",
            Self::Within => "within",
            Self::Steps => "steps",
            Self::No => "no",
            Self::Dead => "dead",
            Self::Ends => "ends",
            Self::Exactly => "exactly",
            Self::One => "one",
            Self::Zero => "zero",
            Self::Or => "or",
            Self::Many => "many",
            Self::Cardinality => "cardinality",
            Self::Range => "range",
            Self::Domain => "domain",
            Self::Optional => "optional",
            Self::And => "and",
            Self::Not => "not",
            Self::True => "true",
            Self::False => "false",
            Self::Unknown => "unknown",
            Self::Continue => "continue",
        }
    }

    fn from_identifier(identifier: &str) -> Option<Self> {
        Some(match identifier {
            "module" => Self::Module,
            "import" => Self::Import,
            "enum" => Self::Enum,
            "entity" => Self::Entity,
            "state" => Self::State,
            "derive" => Self::Derive,
            "decision" => Self::Decision,
            "action" => Self::Action,
            "transition" => Self::Transition,
            "on" => Self::On,
            "set" => Self::Set,
            "rule" => Self::Rule,
            "case" => Self::Case,
            "invariant" => Self::Invariant,
            "trace" => Self::Trace,
            "when" => Self::When,
            "then" => Self::Then,
            "source" => Self::Source,
            "override" => Self::Override,
            "let" => Self::Let,
            "expect" => Self::Expect,
            "for" => Self::For,
            "all" => Self::All,
            "some" => Self::Some,
            "if" => Self::If,
            "initially" => Self::Initially,
            "always" => Self::Always,
            "terminates" => Self::Terminates,
            "within" => Self::Within,
            "steps" => Self::Steps,
            "no" => Self::No,
            "dead" => Self::Dead,
            "ends" => Self::Ends,
            "exactly" => Self::Exactly,
            "one" => Self::One,
            "zero" => Self::Zero,
            "or" => Self::Or,
            "many" => Self::Many,
            "cardinality" => Self::Cardinality,
            "range" => Self::Range,
            "domain" => Self::Domain,
            "optional" => Self::Optional,
            "and" => Self::And,
            "not" => Self::Not,
            "true" => Self::True,
            "false" => Self::False,
            "unknown" => Self::Unknown,
            "continue" => Self::Continue,
            _ => return None,
        })
    }
}

/// The semantic kind of a lexical token.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum TokenKind {
    Identifier(String),
    Number(String),
    String(String),
    Keyword(Keyword),
    LeftBrace,
    RightBrace,
    LeftParen,
    RightParen,
    Colon,
    DoubleColon,
    Comma,
    Dot,
    Range,
    Arrow,
    Question,
    Semicolon,
    Equal,
    NotEqual,
    Greater,
    GreaterEqual,
    Less,
    LessEqual,
    Plus,
    Minus,
    Star,
    Slash,
    Newline,
    Eof,
}

impl TokenKind {
    #[must_use]
    pub fn description(&self) -> &'static str {
        match self {
            Self::Identifier(_) => "identifier",
            Self::Number(_) => "number",
            Self::String(_) => "string",
            Self::Keyword(_) => "keyword",
            Self::LeftBrace => "`{`",
            Self::RightBrace => "`}`",
            Self::LeftParen => "`(`",
            Self::RightParen => "`)`",
            Self::Colon => "`:`",
            Self::DoubleColon => "`::`",
            Self::Comma => "`,`",
            Self::Dot => "`.`",
            Self::Range => "`..`",
            Self::Arrow => "`->`",
            Self::Question => "`?`",
            Self::Semicolon => "`;`",
            Self::Equal => "`=`",
            Self::NotEqual => "`!=`",
            Self::Greater => "`>`",
            Self::GreaterEqual => "`>=`",
            Self::Less => "`<`",
            Self::LessEqual => "`<=`",
            Self::Plus => "`+`",
            Self::Minus => "`-`",
            Self::Star => "`*`",
            Self::Slash => "`/`",
            Self::Newline => "newline",
            Self::Eof => "end of file",
        }
    }
}

/// A token paired with its half-open byte range in the original source.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Token {
    pub kind: TokenKind,
    pub span: Span,
}

/// A line comment retained as source trivia.
///
/// Comments are kept out of the parser token stream so they cannot affect the
/// grammar, while formatters and other source-aware tools can still preserve
/// their original spelling and placement.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Comment {
    /// The complete comment text, including its `#` or `//` marker.
    pub text: String,
    pub span: Span,
}

#[derive(Clone, Debug, Default)]
pub struct LexOutput {
    pub tokens: Vec<Token>,
    pub comments: Vec<Comment>,
    pub diagnostics: Vec<Diagnostic>,
}

/// Tokenize a UTF-8 source file. Newlines are significant and are retained.
#[must_use]
pub fn lex(source: &SourceFile) -> LexOutput {
    Lexer::new(&source.text).run()
}

struct Lexer<'a> {
    text: &'a str,
    cursor: usize,
    output: LexOutput,
}

impl<'a> Lexer<'a> {
    fn new(text: &'a str) -> Self {
        Self {
            text,
            cursor: 0,
            output: LexOutput::default(),
        }
    }

    fn run(mut self) -> LexOutput {
        while let Some(character) = self.peek() {
            let start = self.cursor;
            match character {
                ' ' | '\t' | '\u{000C}' => {
                    self.bump();
                }
                '\n' | '\r' => self.lex_newline(start),
                '#' => self.lex_comment(start),
                '/' if self.peek_second() == Some('/') => self.lex_comment(start),
                '"' => self.lex_string(start),
                '0'..='9' => self.lex_number(start),
                '_' => self.lex_identifier(start),
                value if unicode_ident::is_xid_start(value) => self.lex_identifier(start),
                '{' => self.single(TokenKind::LeftBrace),
                '}' => self.single(TokenKind::RightBrace),
                '(' => self.single(TokenKind::LeftParen),
                ')' => self.single(TokenKind::RightParen),
                ':' if self.peek_second() == Some(':') => self.double(TokenKind::DoubleColon),
                ':' => self.single(TokenKind::Colon),
                ',' => self.single(TokenKind::Comma),
                '?' => self.single(TokenKind::Question),
                ';' => self.single(TokenKind::Semicolon),
                '+' => self.single(TokenKind::Plus),
                '*' => self.single(TokenKind::Star),
                '/' => self.single(TokenKind::Slash),
                '.' if self.peek_second() == Some('.') => self.double(TokenKind::Range),
                '.' => self.single(TokenKind::Dot),
                '-' if self.peek_second() == Some('>') => self.double(TokenKind::Arrow),
                '-' => self.single(TokenKind::Minus),
                '=' => self.single(TokenKind::Equal),
                '!' if self.peek_second() == Some('=') => self.double(TokenKind::NotEqual),
                '>' if self.peek_second() == Some('=') => self.double(TokenKind::GreaterEqual),
                '>' => self.single(TokenKind::Greater),
                '<' if self.peek_second() == Some('=') => self.double(TokenKind::LessEqual),
                '<' => self.single(TokenKind::Less),
                _ => {
                    self.bump();
                    self.output.diagnostics.push(
                        Diagnostic::error(
                            "L0001",
                            format!("unexpected character `{character}`"),
                            Span::new(start, self.cursor),
                        )
                        .with_help("remove the character or replace it with a Tess token"),
                    );
                }
            }
        }
        self.output.tokens.push(Token {
            kind: TokenKind::Eof,
            span: Span::new(self.cursor, self.cursor),
        });
        self.output
    }

    fn peek(&self) -> Option<char> {
        self.text[self.cursor..].chars().next()
    }

    fn peek_second(&self) -> Option<char> {
        let mut characters = self.text[self.cursor..].chars();
        characters.next()?;
        characters.next()
    }

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

    fn push(&mut self, kind: TokenKind, start: usize) {
        self.output.tokens.push(Token {
            kind,
            span: Span::new(start, self.cursor),
        });
    }

    fn single(&mut self, kind: TokenKind) {
        let start = self.cursor;
        self.bump();
        self.push(kind, start);
    }

    fn double(&mut self, kind: TokenKind) {
        let start = self.cursor;
        self.bump();
        self.bump();
        self.push(kind, start);
    }

    fn lex_newline(&mut self, start: usize) {
        if self.bump() == Some('\r') && self.peek() == Some('\n') {
            self.bump();
        }
        self.push(TokenKind::Newline, start);
    }

    fn lex_comment(&mut self, start: usize) {
        while !matches!(self.peek(), None | Some('\n' | '\r')) {
            self.bump();
        }
        self.output.comments.push(Comment {
            text: self.text[start..self.cursor].to_owned(),
            span: Span::new(start, self.cursor),
        });
    }

    fn lex_identifier(&mut self, start: usize) {
        self.bump();
        while self
            .peek()
            .is_some_and(|value| value == '_' || unicode_ident::is_xid_continue(value))
        {
            self.bump();
        }
        let raw = &self.text[start..self.cursor];
        let normalized: String = raw.nfc().collect();
        if let Some(keyword) = Keyword::from_identifier(&normalized) {
            self.push(TokenKind::Keyword(keyword), start);
        } else {
            self.push(TokenKind::Identifier(normalized), start);
        }
    }

    fn lex_number(&mut self, start: usize) {
        while self
            .peek()
            .is_some_and(|value| value.is_ascii_digit() || value == '_')
        {
            self.bump();
        }
        if self.peek() == Some('.')
            && self
                .peek_second()
                .is_some_and(|value| value.is_ascii_digit())
        {
            self.bump();
            while self
                .peek()
                .is_some_and(|value| value.is_ascii_digit() || value == '_')
            {
                self.bump();
            }
        }
        let text = &self.text[start..self.cursor];
        if text.as_bytes().iter().enumerate().any(|(index, value)| {
            *value == b'_'
                && (!index
                    .checked_sub(1)
                    .and_then(|previous| text.as_bytes().get(previous))
                    .is_some_and(u8::is_ascii_digit)
                    || !text
                        .as_bytes()
                        .get(index + 1)
                        .is_some_and(u8::is_ascii_digit))
        }) {
            self.output.diagnostics.push(
                Diagnostic::error(
                    "L0005",
                    "invalid `_` placement in numeric literal",
                    Span::new(start, self.cursor),
                )
                .with_help("place `_` only between digits, for example `100_000`"),
            );
        }
        self.push(TokenKind::Number(text.to_owned()), start);
    }

    fn lex_string(&mut self, start: usize) {
        self.bump();
        let mut value = String::new();
        let mut terminated = false;
        while let Some(character) = self.peek() {
            match character {
                '"' => {
                    self.bump();
                    terminated = true;
                    break;
                }
                '\n' | '\r' => break,
                '\\' => {
                    let escape_start = self.cursor;
                    self.bump();
                    self.lex_escape(&mut value, escape_start);
                }
                other => {
                    self.bump();
                    value.push(other);
                }
            }
        }
        if !terminated {
            self.output.diagnostics.push(
                Diagnostic::error(
                    "L0002",
                    "unterminated string literal",
                    Span::new(start, self.cursor),
                )
                .with_help("add a closing double quote before the end of the line"),
            );
        }
        self.push(TokenKind::String(value), start);
    }

    fn lex_escape(&mut self, value: &mut String, start: usize) {
        let Some(escaped) = self.bump() else {
            return;
        };
        match escaped {
            'n' => value.push('\n'),
            'r' => value.push('\r'),
            't' => value.push('\t'),
            '0' => value.push('\0'),
            '\\' => value.push('\\'),
            '"' => value.push('"'),
            'u' => self.lex_unicode_escape(value, start),
            other => {
                value.push(other);
                self.output.diagnostics.push(
                    Diagnostic::error(
                        "L0003",
                        format!("unknown string escape `\\{other}`"),
                        Span::new(start, self.cursor),
                    )
                    .with_help(r#"supported escapes are \n, \r, \t, \0, \\, \", and \u{...}"#),
                );
            }
        }
    }

    fn lex_unicode_escape(&mut self, value: &mut String, start: usize) {
        if self.peek() != Some('{') {
            self.output.diagnostics.push(Diagnostic::error(
                "L0004",
                "Unicode escape must use `\\u{...}` syntax",
                Span::new(start, self.cursor),
            ));
            return;
        }
        self.bump();
        let digits_start = self.cursor;
        while self
            .peek()
            .is_some_and(|character| character.is_ascii_hexdigit())
        {
            self.bump();
        }
        let digits_end = self.cursor;
        if self.peek() == Some('}') {
            self.bump();
        } else {
            self.output.diagnostics.push(Diagnostic::error(
                "L0004",
                "unterminated Unicode escape",
                Span::new(start, self.cursor),
            ));
            return;
        }
        let scalar = u32::from_str_radix(&self.text[digits_start..digits_end], 16)
            .ok()
            .and_then(char::from_u32);
        if let Some(character) = scalar {
            value.push(character);
        } else {
            self.output.diagnostics.push(Diagnostic::error(
                "L0004",
                "invalid Unicode scalar value",
                Span::new(start, self.cursor),
            ));
        }
    }
}

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

    fn kinds(text: &str) -> Vec<TokenKind> {
        lex(&SourceFile::new("test.tes", text))
            .tokens
            .into_iter()
            .map(|token| token.kind)
            .collect()
    }

    #[test]
    fn recognizes_keywords_and_korean_identifier() {
        assert_eq!(
            kinds("module 강의평가\n"),
            vec![
                TokenKind::Keyword(Keyword::Module),
                TokenKind::Identifier("강의평가".into()),
                TokenKind::Newline,
                TokenKind::Eof,
            ]
        );
        assert_eq!(
            kinds("invariant check"),
            vec![
                TokenKind::Keyword(Keyword::Invariant),
                TokenKind::Identifier("check".into()),
                TokenKind::Eof,
            ]
        );
        assert_eq!(
            kinds("for some"),
            vec![
                TokenKind::Keyword(Keyword::For),
                TokenKind::Keyword(Keyword::Some),
                TokenKind::Eof,
            ]
        );
    }

    #[test]
    fn priority_is_an_identifier_not_a_keyword() {
        assert_eq!(
            kinds("priority"),
            vec![TokenKind::Identifier("priority".into()), TokenKind::Eof]
        );
    }

    #[test]
    fn normalizes_identifiers_to_nfc() {
        let output = kinds("e\u{301}");
        assert_eq!(output[0], TokenKind::Identifier("é".into()));
    }

    #[test]
    fn distinguishes_decimal_dot_and_range() {
        assert_eq!(
            kinds("0..100 0.70 s.점수"),
            vec![
                TokenKind::Number("0".into()),
                TokenKind::Range,
                TokenKind::Number("100".into()),
                TokenKind::Number("0.70".into()),
                TokenKind::Identifier("s".into()),
                TokenKind::Dot,
                TokenKind::Identifier("점수".into()),
                TokenKind::Eof,
            ]
        );
    }

    #[test]
    fn accepts_digit_separators_in_integer_and_decimal_literals() {
        let output = lex(&SourceFile::new("test.tes", "100_000 1_234.50_01 0..9_999"));
        assert!(output.diagnostics.is_empty(), "{:?}", output.diagnostics);
        assert_eq!(
            output
                .tokens
                .into_iter()
                .map(|token| token.kind)
                .collect::<Vec<_>>(),
            vec![
                TokenKind::Number("100_000".into()),
                TokenKind::Number("1_234.50_01".into()),
                TokenKind::Number("0".into()),
                TokenKind::Range,
                TokenKind::Number("9_999".into()),
                TokenKind::Eof,
            ]
        );
    }

    #[test]
    fn reports_misplaced_digit_separators() {
        for literal in ["1__000", "1_", "1_.0", "1.0_", "1.0__1"] {
            let output = lex(&SourceFile::new("test.tes", literal));
            assert_eq!(output.diagnostics.len(), 1, "{literal}");
            assert_eq!(output.diagnostics[0].code, "L0005", "{literal}");
        }
    }

    #[test]
    fn keeps_one_token_for_crlf() {
        let output = lex(&SourceFile::new("test.tes", "a\r\nb\n"));
        assert_eq!(
            output
                .tokens
                .iter()
                .filter(|token| token.kind == TokenKind::Newline)
                .count(),
            2
        );
        assert_eq!(output.tokens[1].span, Span::new(1, 3));
    }

    #[test]
    fn skips_hash_and_slash_comments_but_keeps_newlines() {
        let source = SourceFile::new("test.tes", "a # first\n// second\nb / 2");
        let output = lex(&source);
        assert_eq!(
            output
                .tokens
                .into_iter()
                .map(|token| token.kind)
                .collect::<Vec<_>>(),
            vec![
                TokenKind::Identifier("a".into()),
                TokenKind::Newline,
                TokenKind::Newline,
                TokenKind::Identifier("b".into()),
                TokenKind::Slash,
                TokenKind::Number("2".into()),
                TokenKind::Eof,
            ]
        );
        assert_eq!(
            output.comments,
            vec![
                Comment {
                    text: "# first".into(),
                    span: Span::new(2, 9),
                },
                Comment {
                    text: "// second".into(),
                    span: Span::new(10, 19),
                },
            ]
        );
    }

    #[test]
    fn comment_markers_inside_strings_are_not_trivia() {
        let output = lex(&SourceFile::new("test.tes", "\"# text // text\" // actual"));
        assert_eq!(output.comments.len(), 1);
        assert_eq!(output.comments[0].text, "// actual");
        assert_eq!(
            output.tokens[0].kind,
            TokenKind::String("# text // text".into())
        );
    }

    #[test]
    fn decodes_string_escapes() {
        assert_eq!(
            kinds(r#""line\n\u{D55C}\"""#),
            vec![TokenKind::String("line\n한\"".into()), TokenKind::Eof]
        );
    }

    #[test]
    fn reports_unterminated_string_and_recovers_at_newline() {
        let output = lex(&SourceFile::new("test.tes", "\"oops\nmodule ok"));
        assert_eq!(output.diagnostics.len(), 1);
        assert_eq!(output.diagnostics[0].code, "L0002");
        assert!(
            output
                .tokens
                .iter()
                .any(|token| { token.kind == TokenKind::Keyword(Keyword::Module) })
        );
    }

    #[test]
    fn reports_invalid_char_and_continues() {
        let output = lex(&SourceFile::new("test.tes", "a @ b"));
        assert_eq!(output.diagnostics.len(), 1);
        assert_eq!(output.tokens.len(), 3);
    }
}