lexer.py

import re
from dataclasses import dataclass
from typing import List, Optional

@dataclass
class Token:
    type: str
    lex: str
    line: int
    col: int

class Lexer:
    def __init__(self, source: str):
        self.source = source
        self.tokens: List[Token] = []
        self.start = 0
        self.current = 0
        self.line = 1
        self.line_start = 0

    def scan_tokens(self) -> List[Token]:
        while not self.is_at_end():
            self.start = self.current
            self.scan_token()
        
        self.tokens.append(Token("EOF", "", self.line, self.current - self.line_start))
        return self.tokens

    def is_at_end(self) -> bool:
        return self.current >= len(self.source)

    def advance(self) -> str:
        self.current += 1
        return self.source[self.current - 1]

    def peek(self) -> str:
        if self.is_at_end():
            return '\0'
        return self.source[self.current]

    def peek_next(self) -> str:
        if self.current + 1 >= len(self.source):
            return '\0'
        return self.source[self.current + 1]

    def match(self, expected: str) -> bool:
        if self.is_at_end():
            return False
        if self.source[self.current] != expected:
            return False
        self.current += 1
        return True

    def add_token(self, type: str, literal: str = None):
        text = self.source[self.start:self.current]
        # Calculate column
        col = self.start - self.line_start + 1
        self.tokens.append(Token(type, text, self.line, col))

    def scan_token(self):
        c = self.advance()
        
        if c == '(': self.add_token('LPAREN')
        elif c == ')': self.add_token('RPAREN')
        elif c == '{': self.add_token('LBRACE')
        elif c == '}': self.add_token('RBRACE')
        elif c == '[': self.add_token('LBRACK')
        elif c == ']': self.add_token('RBRACK')
        elif c == ',': self.add_token('COMMA')
        elif c == ';': self.add_token('SEMI')
        elif c == '+': self.add_token('PLUS')
        elif c == '-': self.add_token('MINUS')
        elif c == '*': self.add_token('STAR')
        elif c == '/':
            if self.match('/'):
                # Comment //
                while self.peek() != '\n' and not self.is_at_end():
                    self.advance()
            elif self.match('*'):
                # Block comment /* ... */
                while not self.is_at_end():
                    if self.peek() == '*' and self.peek_next() == '/':
                        self.advance()
                        self.advance()
                        break
                    if self.peek() == '\n':
                        self.line += 1
                        self.line_start = self.current + 1
                    self.advance()
            else:
                self.add_token('SLASH')
        elif c == '!':
            self.add_token('NE' if self.match('=') else 'NOT')
        elif c == '=':
            self.add_token('EQ' if self.match('=') else 'EQUAL')
        elif c == '<':
            self.add_token('LE' if self.match('=') else 'LT')
        elif c == '>':
            self.add_token('GE' if self.match('=') else 'GT')
        elif c == '&':
            if self.match('&'): self.add_token('AND')
        elif c == '|':
            if self.match('|'): self.add_token('OR')
        
        elif c in [' ', '\r', '\t']:
            pass # Ignore whitespace
        
        elif c == '\n':
            self.line += 1
            self.line_start = self.current
            
        elif c == '"':
            self.string()
        
        else:
            if c.isdigit():
                self.number()
            elif c.isalpha() or c == '_':
                self.identifier()
            else:
                # Unexpected character
                pass 

    def string(self):
        while self.peek() != '"' and not self.is_at_end():
            if self.peek() == '\n':
                self.line += 1
                self.line_start = self.current + 1
            self.advance()
        
        if self.is_at_end():
            # Unterminated string
            return

        self.advance() # The closing "
        
        # Trim quotes
        value = self.source[self.start+1 : self.current-1]
        self.add_token('STRING', value)

    def number(self):
        is_float = False
        while self.peek().isdigit():
            self.advance()

        if self.peek() == '.' and self.peek_next().isdigit():
            is_float = True
            self.advance() # Consume .
            while self.peek().isdigit():
                self.advance()
        
        self.add_token('NUM')

    def identifier(self):
        while self.peek().isalnum() or self.peek() == '_':
            self.advance()

        text = self.source[self.start:self.current]
        type = self.keywords.get(text, 'ID')
        self.add_token(type)

    keywords = {
        'int': 'INT',
        'float': 'FLOAT',
        'char': 'CHAR',
        'void': 'VOID',
        'if': 'IF',
        'else': 'ELSE',
        'for': 'FOR',
        'while': 'WHILE',
        'return': 'RETURN',
        'true': 'TRUE',
        'false': 'FALSE'
    }