Skip to main content

oak_d2/lexer/
mod.rs

1/// Token type definitions for D2.
2pub mod token_type;
3
4use crate::{D2TokenType, language::D2Language};
5use oak_core::{
6    Lexer, LexerCache, LexerState,
7    lexer::LexOutput,
8    source::{Source, TextEdit},
9};
10
11/// Lexer for D2 diagram language.
12pub struct D2Lexer<'config> {
13    config: &'config D2Language,
14}
15
16impl<'config> D2Lexer<'config> {
17    /// Creates a new D2Lexer with the given language configuration.
18    pub fn new(config: &'config D2Language) -> Self {
19        Self { config }
20    }
21}
22
23impl<'config> Lexer<D2Language> for D2Lexer<'config> {
24    fn lex<'a, S: Source + ?Sized>(&self, source: &S, _edits: &[TextEdit], cache: &'a mut impl LexerCache<D2Language>) -> LexOutput<D2Language> {
25        let mut state = LexerState::new(source);
26
27        while state.not_at_end() {
28            let start = state.get_position();
29
30            match state.peek() {
31                Some(' ') | Some('\t') => {
32                    // Whitespace
33                    while state.peek().map_or(false, |c| c == ' ' || c == '\t') {
34                        state.advance(1);
35                    }
36                    let end = state.get_position();
37                    state.add_token(D2TokenType::Whitespace, start, end);
38                }
39                Some('\n') | Some('\r') => {
40                    // Newline
41                    if state.peek() == Some('\r') {
42                        state.advance(1);
43                    }
44                    if state.peek() == Some('\n') {
45                        state.advance(1);
46                    }
47                    let end = state.get_position();
48                    state.add_token(D2TokenType::Newline, start, end);
49                }
50                Some('#') => {
51                    // Comment
52                    while state.peek().map_or(false, |c| c != '\n' && c != '\r') {
53                        state.advance(1);
54                    }
55                    let end = state.get_position();
56                    state.add_token(D2TokenType::Comment, start, end);
57                }
58                Some(':') => {
59                    // Colon
60                    state.advance(1);
61                    let end = state.get_position();
62                    state.add_token(D2TokenType::Colon, start, end);
63                }
64                Some('-') if state.peek_next_n(1) == Some('>') => {
65                    // Arrow
66                    state.advance(2);
67                    let end = state.get_position();
68                    state.add_token(D2TokenType::Arrow, start, end);
69                }
70                Some('{') => {
71                    // Left brace
72                    state.advance(1);
73                    let end = state.get_position();
74                    state.add_token(D2TokenType::LeftBrace, start, end);
75                }
76                Some('}') => {
77                    // Right brace
78                    state.advance(1);
79                    let end = state.get_position();
80                    state.add_token(D2TokenType::RightBrace, start, end);
81                }
82                Some(c) if c.is_alphabetic() || c == '_' => {
83                    // Identifier
84                    while state.peek().map_or(false, |c| c.is_alphanumeric() || c == '_') {
85                        state.advance(1);
86                    }
87                    let end = state.get_position();
88                    state.add_token(D2TokenType::Id, start, end);
89                }
90                Some('"') => {
91                    // Label
92                    state.advance(1); // Skip opening quote
93                    while state.peek().map_or(false, |c| c != '"') {
94                        state.advance(1);
95                    }
96                    if state.peek() == Some('"') {
97                        state.advance(1); // Skip closing quote
98                    }
99                    let end = state.get_position();
100                    state.add_token(D2TokenType::Label, start, end);
101                }
102                _ => {
103                    // Error or unknown character
104                    state.advance(1);
105                    let end = state.get_position();
106                    state.add_token(D2TokenType::Error, start, end);
107                }
108            }
109        }
110
111        state.add_eof();
112        state.finish_with_cache(Ok(()), cache)
113    }
114}