Pratt Parsing

parse.py

from dataclasses import dataclass, field
from enum import Enum, auto
from abc import ABC, abstractmethod
from typing import Dict, Optional
import random


class TokenType(Enum):
    NUMBER = auto()
    IDENT = auto()
    OPERTOR = auto()
    EOF = auto()


@dataclass
class Token:
    """Token 对象

    Attributes:
        TokenType: 类型
        TokenLiteral: 字面量
    """

    token_type: TokenType
    literal: str


class AstNode(ABC):
    """Ast 结点接口"""

    @abstractmethod
    def to_string(self) -> str:
        pass


@dataclass
class Atom(AstNode):
    """Ast 的叶子结点

    Attributes:
        token: 叶子结点存储的 token
    """

    token: Token

    def to_string(self) -> str:
        return self.token.literal


@dataclass
class BinaryExpression(AstNode):
    """Ast 二元运算符结点

    Attributes:
        operator: 运算符
        left: 左儿子
        right: 右儿子
    """

    operator: str
    left: Optional[AstNode] = None
    right: Optional[AstNode] = None

    def to_string(self) -> str:
        assert self.left is not None and self.right is not None
        return f"({self.left.to_string()} {self.operator} {self.right.to_string()})"


@dataclass
class Parser:
    tokens: list[Token]
    token_pos: int = 0
    expression: Optional[AstNode] = None
    precedence_map: Dict[str, int] = field(
        default_factory=lambda: {
            "s": 0,
            "+": 1,
            "-": 1,
            "*": 2,
            "/": 2,
        }
    )

    def has_token(self) -> bool:
        return self.token_pos < len(self.tokens)

    def next_token(self) -> Token:
        token = self.tokens[self.token_pos]
        self.token_pos += 1
        return token

    def peek_token(self) -> Token:
        return self.tokens[self.token_pos]

    def parse_expression(self) -> None:
        self.expression = self.__pratt_parse(self.precedence_map["s"])

    def __pratt_parse(self, left_precedence: int) -> Optional[AstNode]:
        root = Atom(self.next_token())

        # 如果当前运算符的 binding power 一直不够大，被右边的运算符狠狠捕获♂
        while True:
            # 查看下一个是否是运算符
            peek_token = self.peek_token()
            if not self.has_token() or peek_token.token_type != TokenType.OPERTOR:
                break

            # 如果当前 token 左边运算符的 binding power 足够大，那么当前 token 就不用被右边的运算符捕获
            operator = peek_token.literal
            right_precedence = self.precedence_map[operator]
            if left_precedence >= right_precedence:
                break

            # 存储下一个运算符
            root = BinaryExpression(
                operator=self.next_token().literal,
                left=root,
                right=self.__pratt_parse(self.precedence_map[operator]),
            )
        return root


def fuzzing_test():
    def get_random_number() -> int:
        return random.randint(1, 9)

    operator_list = ["+", "-", "*", "/"]

    for _ in range(1145):
        tokens: list[Token] = []
        expression_str = ""

        init_number = get_random_number()
        expression_str += str(init_number)
        tokens.append(Token(TokenType.NUMBER, str(init_number)))
        for _ in range(random.randint(10, 20)):
            random_operator = random.choice(operator_list)
            random_number = get_random_number()

            expression_str += random_operator
            expression_str += str(random_number)

            tokens.append(Token(TokenType.OPERTOR, random_operator))
            tokens.append(Token(TokenType.NUMBER, str(random_number)))
        tokens.append(Token(TokenType.EOF, ""))

        parser = Parser(tokens)
        parser.parse_expression()  # 使用 Pratt Parsing 进行分析

        if parser.expression is not None:
            our_expression_str = parser.expression.to_string()
            assert eval(expression_str) == eval(our_expression_str)
    print("success!")


def main() -> None:
    tokens: list[Token] = [
        Token(TokenType.NUMBER, "1"),
        Token(TokenType.OPERTOR, "+"),
        Token(TokenType.NUMBER, "3"),
        Token(TokenType.OPERTOR, "*"),
        Token(TokenType.NUMBER, "4"),
        Token(TokenType.OPERTOR, "*"),
        Token(TokenType.NUMBER, "5"),
        Token(TokenType.OPERTOR, "-"),
        Token(TokenType.NUMBER, "6"),
        Token(TokenType.EOF, ""),
    ]

    parser = Parser(tokens)
    parser.parse_expression()  # 使用 Pratt Parsing 进行分析

    if parser.expression is not None:
        print(parser.expression.to_string())


if __name__ == "__main__":
    # main()
    fuzzing_test()