device-driver-lexer 2.0.0

use std::{borrow::Cow, fmt::Display};

use device_driver_common::{
    span::{SpanExt, Spanned},
    specifiers::{Access, AddressMode, BaseType, ByteOrder, Integer},
};
use logos::Logos;

pub fn lex(source: &str) -> Vec<Spanned<Token<'_>>> {
    Token::lexer(source)
        .spanned()
        .map(|(token, span)| match token {
            Ok(token) => token.with_span(span),
            Err(()) => Token::Error.with_span(span),
        })
        .collect()
}

// Don't forget to update the book when tokens are added, changed or removed!
#[derive(Debug, Clone, Copy, PartialEq, Logos)]
#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
#[logos(skip r"[ \t\r\n]+")] // Skip (common) whitespace
#[logos(skip(r"//[^\n]*", allow_greedy = true))] // Skip comments
pub enum Token<'src> {
    #[regex(r"///[^\n]*", allow_greedy = true, callback = |lex| lex.slice().trim_start_matches("///"))]
    DocCommentLine(&'src str),
    #[regex(r"\p{XID_Start}[\p{XID_Continue}-]*")]
    Ident(&'src str),
    #[token("{")]
    CurlyOpen,
    #[token("}")]
    CurlyClose,
    #[token("[")]
    BracketOpen,
    #[token("]")]
    BracketClose,
    #[token(",")]
    Comma,
    #[token(":")]
    Colon,
    #[token("_")]
    Underscore,
    #[token("->")]
    Arrow,
    #[token("*")]
    Star,
    #[token("try")]
    Try,
    #[token("as")]
    As,
    #[token("allow")]
    Allow,
    #[token("default")]
    Default,
    #[token("catch-all")]
    CatchAll,
    #[token("stride")]
    Stride,
    #[regex(r"-?[0-9][_0-9]*")] // Decimal
    #[regex(r"-?0b[_0-1]+")] // Binary
    #[regex(r"-?0o[_0-7]+")] // Octal
    #[regex(r"-?0x[_0-9a-fA-F]+")] // Hex
    Num(&'src str),
    #[token("RW", |_| Access::RW)]
    #[token("RO", |_| Access::RO)]
    #[token("WO", |_| Access::WO)]
    Access(Access),
    #[token("BE", |_| ByteOrder::BE)]
    #[token("LE", |_| ByteOrder::LE)]
    ByteOrder(ByteOrder),
    /// All the base types except fixed integers. Those are [`Self::Integer`].
    #[token("uint", |_| BaseType::Uint)]
    #[token("int", |_| BaseType::Int)]
    #[token("bool", |_| BaseType::Bool)]
    BaseType(BaseType),
    #[token("u8", |_| Integer::U8)]
    #[token("u16", |_| Integer::U16)]
    #[token("u32", |_| Integer::U32)]
    #[token("u64", |_| Integer::U64)]
    #[token("i8", |_| Integer::I8)]
    #[token("i16", |_| Integer::I16)]
    #[token("i32", |_| Integer::I32)]
    #[token("i64", |_| Integer::I64)]
    Integer(Integer),
    #[token("mapped", |_| AddressMode::Mapped)]
    #[token("indexed", |_| AddressMode::Indexed)]
    AddressMode(AddressMode),
    // Very simple definition without string escaping
    #[regex(r#""[^"]*""#, callback = |lex| lex.slice().strip_prefix('"').unwrap().strip_suffix('"').unwrap())]
    String(&'src str),
    #[regex(r"\S", priority = 0)] // Any non-whitespace character
    Unexpected(&'src str),
    Error, // Catch-all for errors
}

impl Display for Token<'_> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Token::DocCommentLine(_) => write!(f, "doc comment"),
            Token::Ident(_) => write!(f, "identifier"),
            Token::CurlyOpen => write!(f, "{{"),
            Token::CurlyClose => write!(f, "}}"),
            Token::BracketOpen => write!(f, "["),
            Token::BracketClose => write!(f, "]"),
            Token::Colon => write!(f, ":"),
            Token::Underscore => write!(f, "_"),
            Token::Comma => write!(f, ","),
            Token::Arrow => write!(f, "->"),
            Token::Star => write!(f, "*"),
            Token::Try => write!(f, "try"),
            Token::As => write!(f, "as"),
            Token::Allow => write!(f, "allow"),
            Token::Default => write!(f, "default"),
            Token::CatchAll => write!(f, "catch-all"),
            Token::Stride => write!(f, "stride"),
            Token::Num(_) => write!(f, "number"),
            Token::Access(_) => write!(f, "access specifier"),
            Token::ByteOrder(_) => write!(f, "byte order"),
            Token::BaseType(_) => write!(f, "base type"),
            Token::Integer(_) => write!(f, "integer type"),
            Token::AddressMode(_) => write!(f, "address mode"),
            Token::String(_) => write!(f, "string"),
            Token::Unexpected(val) => write!(f, "{}", val.escape_debug()),
            Token::Error => write!(f, "ERROR"),
        }
    }
}

impl<'src> Token<'src> {
    fn get_human_string(&self) -> Cow<'static, str> {
        match self {
            Token::DocCommentLine(line) => format!("///{line}").into(),
            Token::Ident(ident) => format!("#{ident}").into(),
            Token::CurlyOpen => "{".into(),
            Token::CurlyClose => "}".into(),
            Token::BracketOpen => "[".into(),
            Token::BracketClose => "]".into(),
            Token::Colon => ":".into(),
            Token::Underscore => "_".into(),
            Token::Comma => ",".into(),
            Token::Arrow => "->".into(),
            Token::Try => "try".into(),
            Token::Star => "*".into(),
            Token::As => "as".into(),
            Token::Num(n) => n.to_string().into(),
            Token::Access(val) => val.to_string().into(),
            Token::ByteOrder(val) => val.to_string().into(),
            Token::BaseType(val) => val.to_string().into(),
            Token::Integer(val) => val.to_string().into(),
            Token::AddressMode(val) => val.to_string().into(),
            Token::Allow => "allow".into(),
            Token::Default => "default".into(),
            Token::CatchAll => "catch-all".into(),
            Token::Stride => "stride".into(),
            Token::Unexpected(raw) => format!("!{raw}").into(),
            Token::Error => "UNEXPECTED".into(),
            Token::String(val) => format!("\"{val}\"").into(),
        }
    }

    /// New line before, new line after, indent change
    fn get_print_format(&self) -> (bool, bool, i32) {
        match self {
            Token::DocCommentLine(_) => (true, true, 0),
            Token::Comma => (false, true, 0),
            Token::CurlyOpen | Token::BracketOpen => (false, true, 1),
            Token::CurlyClose | Token::BracketClose => (true, false, -1),
            _ => (false, false, 0),
        }
    }

    pub fn formatted_print<'a, I: Iterator<Item = &'a Token<'src>>>(
        stream: &mut impl std::fmt::Write,
        tokens: I,
    ) -> Result<(), std::fmt::Error>
    where
        'src: 'a,
    {
        let mut indent = 0i32;
        for token in tokens {
            let (newline_before, newline_after, indent_change) = token.get_print_format();

            indent += indent_change;
            if newline_before {
                write!(
                    stream,
                    "\n{:width$}",
                    "",
                    width = indent.max(0) as usize * 4
                )?;
            }

            write!(stream, "{} ", token.get_human_string())?;

            if newline_after {
                write!(
                    stream,
                    "\n{:width$}",
                    "",
                    width = indent.max(0) as usize * 4
                )?;
            }
        }

        Ok(())
    }
}