Skip to main content
1 of 5

top-down lexer in rust

the other day, i decided to build a top-down lexer in rust, just for fun.
this is what i have so far:

use std::str::Chars;
use std::env;
use std::fs;

#[derive(Debug, Clone)]
enum Op {
    Plus,
    Minus,
    Multiply,
    Divide,
    Modulo
}

#[derive(Debug, Clone)]
enum TokenKind {
    Opr(Op),
    Ident(String),
    Str(String),
    Num(String),
}

#[derive(Debug, Clone)]
struct Lexer<'a> {
    // string containing entire program
    source: &'a str,
    // the peekable character
    peek: Option<char>,
    // the characters iterator
    chars: Chars<'a>,
    // utf-8 position in the string
    pos: usize,

    // the bottom two are for error reporting/debugging
    row: usize,
    col: usize
}

#[allow(unused)]
impl<'a> Lexer<'a> {
    fn new(source: &'a str) -> Self {
        let mut chars = source.chars();
        Self {
            source,
            peek: chars.next(),
            chars,
            pos: 0,
            row: 1,
            col: 1
        }
    }

    #[inline]
    fn position(&self) -> (usize, usize) {
        (self.row, self.col)
    }

    #[inline]
    fn is_over(&self) -> bool {
        self.peek().is_none()
    }

    #[inline]
    fn bump(&mut self) {
        if let Some(ch) = self.chars.next() {
            self.pos += ch.len_utf8();
            self.col += 1;
            if ch == '\n' {
                self.col = 0;
                self.row += 1;
            }
            self.peek = Some(ch);
        } else {
            self.peek = None;
        }
    }

    #[inline]
    fn peek(&self) -> Option<char> {
        self.peek
    }

    #[inline]
    fn unicode_escape(&mut self, ch: char) -> Option<char> {
        // e.g. \n \t \u{...}
        let res = match ch {
            'n'   => '\n',
            't'   => '\t',
            'r'   => '\r',
            'u'   => todo!("implement unicode character codes"),
            'x'   => todo!("implement hex character codes"),
            other => other
        };
        Some(res)
    }

    fn trim_ident(&mut self) -> TokenKind {
        let start_pos = self.pos;

        while let Some(s) = self.peek() {
            if !(s.is_alphanumeric() || s == '_') {
                break;
            }
            self.next();
        }

        TokenKind::Ident(self.source[start_pos..self.pos].to_owned())
    }

    fn trim_number(&mut self) -> TokenKind {
        let start_pos = self.pos;

        while let Some(s) = self.peek() {
            if !s.is_numeric() {
                break;
            }
            self.next();
        }
        TokenKind::Num(self.source[start_pos..self.pos].to_owned())
    }

    fn trim_string(&mut self) -> TokenKind {
        let start_pos = self.pos;

        while let Some(s) = self.peek() {
            if s == '\\' {
                self.bump();
            }
            if s == '"' {
                break
            }
            self.bump();
        }
        if self.is_over() {
            todo!("Could not lex string");
        }
        TokenKind::Str(self.source[start_pos..self.pos].to_owned())
    }
    
    #[inline]
    fn trim_comment(&mut self) {
        while let Some(s) = self.peek() {
            if s == '\n' {
                self.row = 0;
                self.col += 1;
                break
            }
            self.bump();
        }
    }

    #[inline]
    fn trim_block_comment(&mut self) {
        let mut recursion: usize = 1;
        while let Some(s) = self.next() {
            match s {
                '\n' => {
                    self.row = 0;
                    self.col += 1;
                },
                '/' => {
                    if let Some('*') = self.peek() {
                        self.bump();
                        recursion += 1;
                    }
                },
                '*' => {
                    if let Some('/') = self.peek() {
                        self.bump();
                        recursion -= 1;
                    }
                }
                _ => {},
            }
            if recursion == 0 {
                break
            }
        }
    }

    #[inline]
    fn trim_whitespace(&mut self) {
        while let Some(s) = self.peek() {
            if !s.is_whitespace() {
                break;
            }
            self.bump();
        }
    }
    
    #[inline]
    fn next(&mut self) -> Option<char> {
        let peek = self.peek;
        self.bump();
        peek
    }
}

impl<'a> Iterator for Lexer<'a> {
    type Item = TokenKind;

    fn next(&mut self) -> Option<TokenKind> {
        loop {
            self.trim_whitespace();
            return match self.peek()? {
                x if x.is_alphabetic() => return Some(self.trim_ident()),
                x if x.is_numeric() => return Some(self.trim_number()),
                '"' => {
                    self.bump();
                    let string = self.trim_string();
                    self.bump();
                    Some(string)
                },
                '+' => {
                    self.bump();
                    Some(TokenKind::Opr(Op::Plus))
                },
                '-' => {
                    self.bump();
                    Some(TokenKind::Opr(Op::Minus))
                },
                '*' => {
                    self.bump();
                    Some(TokenKind::Opr(Op::Multiply))
                },
                '/' => {
                    self.bump();
                    if let Some('/') = self.peek() {
                        self.trim_comment();
                        continue;
                    }
                    if let Some('*') = self.peek() {
                        self.trim_block_comment();
                        continue;
                    }
                    Some(TokenKind::Opr(Op::Divide))
                },
                '%' => {
                    self.next();
                    Some(TokenKind::Opr(Op::Modulo))
                },
                _ => None
            };
        }
    }
}

HOWEVER i am sure there are many ways i could improve this code, readability and performance-wise. for instance, the loop {} in Lexer::next should be removed, but that messes up comments. also, as a side note, whenever i see other people implement lexers, their .peek() method clones the chars iterator and then calls .next(). my version stores peek as a separate value. can you tell me which is better?