use core::{ops::Range, str};
use std::{collections::VecDeque, ops::ControlFlow};
use unicode_ident::{is_xid_continue, is_xid_start};
use super::token::{FStringToken, Keyword, Token};
use crate::ast::Identifier;
pub struct Lexer<'a> {
input: &'a str,
original_length: usize,
peeked: VecDeque<(Result<Token<'a>, ()>, Range<usize>)>,
pub almost_keyword:
Option<(Identifier, Range<usize>, Option<&'static str>)>,
}
impl<'a> Lexer<'a> {
pub fn next(&mut self) -> Option<(Result<Token<'a>, ()>, Range<usize>)> {
if let Some(t) = self.peeked.pop_front() {
return Some(t);
}
self.next_inner()
}
pub fn peek(&mut self) -> Option<&(Result<Token<'a>, ()>, Range<usize>)> {
if self.peeked.is_empty()
&& let Some(t) = self.next_inner()
{
self.peeked.push_back(t);
}
self.peeked.front()
}
pub fn peek_many<const N: usize>(&mut self) -> Option<[&Token<'a>; N]> {
for _ in 0..N - self.peeked.len() {
let t = self.next_inner()?;
self.peeked.push_back(t);
}
let mut tokens = [const { &Token::AmpAmp }; N];
for (i, token) in tokens.iter_mut().enumerate() {
let (Ok(t), _) = self.peeked.get(i).unwrap() else {
return None;
};
*token = t;
}
Some(tokens)
}
fn next_inner(
&mut self,
) -> Option<(Result<Token<'a>, ()>, Range<usize>)> {
match self.next_token() {
ControlFlow::Continue(()) => {
if self.input.is_empty() {
None
} else {
let start = self.original_length - self.input.len();
let end = start + 1;
Some((Err(()), start..end))
}
}
ControlFlow::Break((tok, span)) => Some((Ok(tok), span)),
}
}
}
impl<'s> Lexer<'s> {
pub fn new(input: &'s str) -> Self {
Self {
input,
original_length: input.len(),
peeked: VecDeque::new(),
almost_keyword: None,
}
}
fn bump(&mut self, n: usize) -> (&'s str, Range<usize>) {
let start = self.original_length - self.input.len();
let (a, b) = self.input.split_at(n);
self.input = b;
let end = self.original_length - self.input.len();
(a, start..end)
}
fn bump_to(&mut self, tail: &str) -> (&'s str, Range<usize>) {
self.bump(self.input.len() - tail.len())
}
fn is_empty(&self) -> bool {
self.input.is_empty()
}
fn next_token(&mut self) -> ControlFlow<(Token<'s>, Range<usize>)> {
self.skip_whitespace();
if self.is_empty() {
return ControlFlow::Continue(());
}
self.ipv6()?;
self.ipv4()?;
self.two_char_punctuation()?;
self.one_char_punctuation()?;
self.as_number()?;
self.hex_number()?;
self.number()?;
self.f_string()?;
self.string()?;
self.char()?;
self.keyword_or_ident()?;
ControlFlow::Continue(())
}
pub fn skip_shebang(&mut self) {
let mut tail = self.input;
if tail.eat_str("#!")
&& tail.starts_with(|c: char| !c.is_whitespace())
{
tail.eat_until('\n');
self.bump_to(tail);
}
}
fn skip_whitespace(&mut self) {
let mut tail = self.input;
loop {
tail.eat_whitespace();
if tail.eat_str("//") {
tail.eat_until('\n');
} else {
break;
}
}
if tail.len() < self.input.len() {
self.bump_to(tail);
}
}
fn two_char_punctuation(
&mut self,
) -> ControlFlow<(Token<'s>, Range<usize>)> {
let Some(x) = self.input.as_bytes().first_chunk() else {
return ControlFlow::Continue(());
};
let tok = match *x {
[b'=', b'='] => Token::EqEq,
[b'!', b'='] => Token::BangEq,
[b'&', b'&'] => Token::AmpAmp,
[b'|', b'|'] => Token::PipePipe,
[b'>', b'='] => Token::AngleRightEq,
[b'<', b'='] => Token::AngleLeftEq,
[b'-', b'>'] => Token::Arrow,
[b'=', b'>'] => Token::FatArrow,
[b'+', b'='] => Token::PlusEq,
[b'-', b'='] => Token::MinusEq,
[b'*', b'='] => Token::StarEq,
[b'/', b'='] => Token::SlashEq,
[b'%', b'='] => Token::PercentEq,
[b'/', b'*'] => Token::SlashStar,
[b'-', b'-'] => Token::HyphenHyphen,
_ => return ControlFlow::Continue(()),
};
let (_, span) = self.bump(2);
ControlFlow::Break((tok, span))
}
fn one_char_punctuation(
&mut self,
) -> ControlFlow<(Token<'s>, Range<usize>)> {
let Some(x) = self.input.as_bytes().first() else {
return ControlFlow::Continue(());
};
let tok = match x {
b'=' => Token::Eq,
b'|' => Token::Pipe,
b'-' => Token::Hyphen,
b':' => Token::Colon,
b';' => Token::SemiColon,
b',' => Token::Comma,
b'.' => Token::Period,
b'+' => Token::Plus,
b'*' => Token::Star,
b'/' => Token::Slash,
b'!' => Token::Bang,
b'{' => Token::CurlyLeft,
b'}' => Token::CurlyRight,
b'?' => Token::QuestionMark,
b'[' => Token::SquareLeft,
b']' => Token::SquareRight,
b'(' => Token::RoundLeft,
b')' => Token::RoundRight,
b'<' => Token::AngleLeft,
b'>' => Token::AngleRight,
b'%' => Token::Percent,
b'#' => Token::Hash,
_ => return ControlFlow::Continue(()),
};
let (_, span) = self.bump(1);
ControlFlow::Break((tok, span))
}
fn ipv6(&mut self) -> ControlFlow<(Token<'s>, Range<usize>)> {
let mut tail = self.input;
let mut count = 0;
while count < 2 {
tail.eat_while(char::is_ascii_hexdigit);
if !tail.eat_char(':') {
return ControlFlow::Continue(());
}
count += 1;
}
tail.eat_while(|c| char::is_ascii_hexdigit(c) || *c == ':');
let (tok, span) = self.bump_to(tail);
ControlFlow::Break((Token::IpV6(tok), span))
}
fn ipv4(&mut self) -> ControlFlow<(Token<'s>, Range<usize>)> {
let mut tail = self.input;
let mut count = 0;
while count < 3 {
if !tail.eat_while(char::is_ascii_digit) {
return ControlFlow::Continue(());
}
if !tail.eat_char('.') {
return ControlFlow::Continue(());
}
count += 1;
}
tail.eat_while(char::is_ascii_digit);
let (tok, span) = self.bump_to(tail);
ControlFlow::Break((Token::IpV4(tok), span))
}
fn as_number(&mut self) -> ControlFlow<(Token<'s>, Range<usize>)> {
let mut tail = self.input;
if !tail.eat_str("AS") {
return ControlFlow::Continue(());
};
let ate_digits = tail.eat_while(char::is_ascii_digit);
if !ate_digits {
return ControlFlow::Continue(());
}
let (tok, span) = self.bump_to(tail);
ControlFlow::Break((Token::Asn(tok), span))
}
fn hex_number(&mut self) -> ControlFlow<(Token<'s>, Range<usize>)> {
let mut tail = self.input;
if !tail.eat_str("0x") {
return ControlFlow::Continue(());
};
tail.eat_while(char::is_ascii_hexdigit);
let (tok, span) = self.bump_to(tail);
ControlFlow::Break((Token::Hex(tok), span))
}
fn number(&mut self) -> ControlFlow<(Token<'s>, Range<usize>)> {
let mut tail = self.input;
if !tail.starts_with(|c: char| c.is_ascii_digit()) {
return ControlFlow::Continue(());
}
tail.eat_while(is_roto_digit);
let mut is_float = false;
'float: {
if tail.starts_with('.') {
if let Some(c) = tail.chars().nth(1)
&& (is_xid_start(c) || c == '.' || c == '_')
{
break 'float;
}
is_float = true;
tail.eat_char('.');
tail.eat_while(is_roto_digit);
}
if tail.eat_one_of(['e', 'E']) {
is_float = true;
tail.eat_one_of(['+', '-']);
tail.eat_while(is_roto_digit);
}
}
let length = self.input.len() - tail.len();
tail.eat_while(|c| is_xid_continue(*c) || *c == '_');
let (tok, span) = self.bump_to(tail);
let (num, suffix) = tok.split_at(length);
ControlFlow::Break((
if is_float {
Token::Float(num, suffix)
} else {
Token::Integer(num, suffix)
},
span,
))
}
fn f_string(&mut self) -> ControlFlow<(Token<'s>, Range<usize>)> {
let mut tail = self.input;
if !tail.eat_str("f\"") {
return ControlFlow::Continue(());
};
let (_tok, span) = self.bump_to(tail);
ControlFlow::Break((Token::FStringStart, span))
}
pub fn f_string_part(
&mut self,
) -> Option<(FStringToken<'s>, Range<usize>)> {
let mut chars = self.input.chars().enumerate();
'outer: while let Some((i, c)) = chars.next() {
match c {
'\\' => {
let (_, c) = chars.next()?;
if c == 'u' || c == 'U' {
let (_, c) = chars.next()?;
if c != '{' {
return None;
}
for (_, c) in chars.by_ref() {
if c == '}' {
continue 'outer;
}
}
return None;
}
continue 'outer;
}
'{' => {
let (i, c) = chars.next()?;
if c == '{' {
continue 'outer;
} else {
let (tok, span) = self.bump(i - 1);
return Some((
FStringToken::StringIntermediate(tok),
span,
));
}
}
'"' => {
let (tok, span) = self.bump(i);
self.bump(1);
return Some((FStringToken::StringEnd(tok), span));
}
_ => {}
}
}
None
}
fn char(&mut self) -> ControlFlow<(Token<'s>, Range<usize>)> {
let mut tail = self.input;
if !tail.eat_char('\'') {
return ControlFlow::Continue(());
};
let mut last_is_backslash = false;
tail.eat_until_fn(|c| {
if last_is_backslash {
last_is_backslash = false;
return false;
}
match c {
'\'' => true,
'\\' => {
last_is_backslash = true;
false
}
_ => false,
}
});
if tail.is_empty() {
return ControlFlow::Continue(());
};
let (tok, span) = self.bump_to(tail);
ControlFlow::Break((Token::Char(tok), span))
}
fn string(&mut self) -> ControlFlow<(Token<'s>, Range<usize>)> {
let mut tail = self.input;
if !tail.eat_char('\"') {
return ControlFlow::Continue(());
};
let mut last_is_backslash = false;
tail.eat_until_fn(|c| {
if last_is_backslash {
last_is_backslash = false;
return false;
}
match c {
'"' => true,
'\\' => {
last_is_backslash = true;
false
}
_ => false,
}
});
if tail.is_empty() {
return ControlFlow::Continue(());
};
let (tok, span) = self.bump_to(tail);
ControlFlow::Break((Token::String(tok), span))
}
fn keyword_or_ident(&mut self) -> ControlFlow<(Token<'s>, Range<usize>)> {
let mut tail = self.input;
let Some(c) = tail.chars().next() else {
return ControlFlow::Continue(());
};
if !(is_xid_start(c) || c == '_') {
return ControlFlow::Continue(());
}
tail = &tail[1..];
tail.eat_while(|c: &char| is_xid_continue(*c));
let (ident, span) = self.bump_to(tail);
let kw = match ident {
"accept" => Keyword::Accept,
"const" => Keyword::Const,
"dep" => Keyword::Dep,
"else" => Keyword::Else,
"enum" => Keyword::Enum,
"filter" => Keyword::Filter,
"filtermap" => Keyword::FilterMap,
"for" => Keyword::For,
"fn" => Keyword::Fn,
"if" => Keyword::If,
"import" => Keyword::Import,
"in" => Keyword::In,
"let" => Keyword::Let,
"match" => Keyword::Match,
"pkg" => Keyword::Pkg,
"record" => Keyword::Record,
"reject" => Keyword::Reject,
"return" => Keyword::Return,
"std" => Keyword::Std,
"super" => Keyword::Super,
"test" => Keyword::Test,
"while" => Keyword::While,
"true" => return ControlFlow::Break((Token::Bool(true), span)),
"false" => return ControlFlow::Break((Token::Bool(false), span)),
x => {
self.record_almost_keyword(x, span.clone());
return ControlFlow::Break((Token::Ident(x), span));
}
};
ControlFlow::Break((Token::Keyword(kw), span))
}
fn record_almost_keyword(&mut self, x: &str, span: Range<usize>) {
let suggestion = match x {
"loop" => None,
"struct" => Some("record"),
"class" => Some("record"),
"data" => Some("record"),
"use" => Some("import"),
"switch" => Some("match"),
"var" => Some("let"),
"local" => Some("let"),
"function" => Some("fn"),
"fun" => Some("fn"),
"func" => Some("fn"),
"def" => Some("fn"),
_ => return,
};
let ident = Identifier::from(x);
self.almost_keyword = Some((ident, span, suggestion));
}
}
fn is_roto_digit(c: &char) -> bool {
c.is_ascii_digit() || *c == '_'
}
trait StrExt {
fn eat_char(&mut self, c: char) -> bool;
fn eat_str(&mut self, s: &str) -> bool;
fn eat_one_of<const N: usize>(&mut self, options: [char; N]) -> bool;
fn eat_until(&mut self, c: char);
fn eat_until_fn(&mut self, c: impl FnMut(&char) -> bool);
fn eat_while(&mut self, pat: impl FnMut(&char) -> bool) -> bool;
fn eat_whitespace(&mut self);
}
impl StrExt for &str {
fn eat_char(&mut self, s: char) -> bool {
if let Some(new) = self.strip_prefix(s) {
*self = new;
true
} else {
false
}
}
fn eat_str(&mut self, s: &str) -> bool {
if let Some(new) = self.strip_prefix(s) {
*self = new;
true
} else {
false
}
}
fn eat_one_of<const N: usize>(&mut self, options: [char; N]) -> bool {
if let Some(new) = self.strip_prefix(options) {
*self = new;
true
} else {
false
}
}
fn eat_until(&mut self, c: char) {
*self = if let Some((_, tail)) = self.split_once(c) {
tail
} else {
""
};
}
fn eat_until_fn(&mut self, mut f: impl FnMut(&char) -> bool) {
*self = if let Some((_, tail)) = self.split_once(|c: char| f(&c)) {
tail
} else {
""
};
}
fn eat_while(&mut self, mut pat: impl FnMut(&char) -> bool) -> bool {
let new = self.trim_start_matches(|c| pat(&c));
if self.len() != new.len() {
*self = new;
true
} else {
false
}
}
fn eat_whitespace(&mut self) {
*self = self.trim_start();
}
}