✨ feat: add cse and constant_folding

Author: SigureMo

.vscode/settings.json

@@ -5,5 +5,6 @@
   "python.analysis.inlayHints.variableTypes": true,
   "[python]": {
     "editor.defaultFormatter": "charliermarsh.ruff"
-  }
+  },
+  "ruff.interpreter": ["${workspaceFolder}/.venv/bin/python"]
 }

README.md

@@ -12,3 +12,18 @@
    <a href="https://github.com/astral-sh/ruff"><img alt="ruff" src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json&style=flat-square"></a>
    <a href="https://gitmoji.dev"><img alt="Gitmoji" src="https://img.shields.io/badge/gitmoji-%20😜%20😍-FFDD67?style=flat-square"></a>
 </p>
+
+## Installation
+
+TODO...
+
+## Usage
+
+```console
+$ expr_simplifier cse "a * 4 + (a * 4)"
+(___t_0 := (a * 4)) + ___t_0
+$ expr_simplifier constant_folding "(___x := 1 + 1) + ___x" --max-iter=2
+(___x := 2) + ___x
+$ expr_simplifier constant_folding "(___x := 1 + 1) + ___x" --max-iter=2
+4
+```

pyproject.toml

@@ -4,7 +4,9 @@ version = "0.1.0"
 description = ""
 readme = "README.md"
 requires-python = ">=3.9"
-dependencies = []
+dependencies = [
+    "typing-extensions>=4.12.2",
+]
 authors = [{ name = "Nyakku Shigure", email = "sigure.qaq@gmail.com" }]
 keywords = []
 license = { text = "MIT" }

src/expr_simplifier/__main__.py

@@ -1,14 +1,48 @@
 from __future__ import annotations
 
 import argparse
+import ast
+from collections.abc import Callable
 
 from expr_simplifier import __version__
+from expr_simplifier.transforms import apply_constant_folding, apply_cse
+from expr_simplifier.typing import Pass
+from expr_simplifier.utils import loop_until_stable
+
+
+def create_pass_command(name: str, passes: list[Pass]) -> Callable[[argparse.Namespace], None]:
+    def pass_command(args: argparse.Namespace) -> None:
+        expr = ast.parse(args.input, mode="eval")
+        simplified_expr = loop_until_stable(expr, passes, args.max_iter)
+        print(ast.unparse(simplified_expr))
+
+    pass_command.__name__ = name
+    return pass_command
+
+
+def create_pass_parser(
+    name: str,
+    passes: list[Pass],
+    description: str,
+    subparser: argparse._SubParsersAction[argparse.ArgumentParser],  # pyright: ignore [reportPrivateUsage]
+) -> None:
+    parser = subparser.add_parser(name, help=description)
+    parser.add_argument("input", help="The expression to simplify")
+    parser.add_argument("--max-iter", type=int, default=100, help="The maximum number of iterations")
+    parser.set_defaults(func=create_pass_command(name, passes))
 
 
 def main() -> None:
     parser = argparse.ArgumentParser(prog="moelib", description="A moe moe project")
     parser.add_argument("-v", "--version", action="version", version=__version__)
-    args = parser.parse_args()  # type: ignore
+    sub_parsers = parser.add_subparsers(help="sub-command help", dest="sub_command")
+
+    create_pass_parser("cse", [apply_cse], "Common Subexpression Elimination", sub_parsers)
+    create_pass_parser("constant_folding", [apply_constant_folding], "Constant Folding", sub_parsers)
+    create_pass_parser("auto", [apply_constant_folding, apply_cse], "Auto Simplification", sub_parsers)
+
+    args = parser.parse_args()
+    args.func(args)
 
 
 if __name__ == "__main__":

src/expr_simplifier/symbol_table.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+
+class UniqueNameGenerator:
+    def __init__(self, prefix: str):
+        self.prefix = prefix
+        self._counter = 0
+
+    def generate_name(self) -> str:
+        name = f"{self.prefix}{self._counter}"
+        self._counter += 1
+        return name
+
+
+class SymbolTable:
+    def __init__(self):
+        self._symbols = set[str]()
+        self._name_generator = UniqueNameGenerator("___t_")
+
+    def define_symbol(self, symbol: str):
+        self._symbols.add(symbol)
+
+    def request_new_symbol(self) -> str:
+        while True:
+            new_symbol = self._name_generator.generate_name()
+            if self.is_symbol_defined(new_symbol):
+                continue
+            self.define_symbol(new_symbol)
+            return new_symbol
+
+    def is_symbol_defined(self, symbol: str) -> bool:
+        return symbol in self._symbols
+
+    def __contains__(self, symbol: str) -> bool:
+        return self.is_symbol_defined(symbol)

src/expr_simplifier/transforms/__init__.py

@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+from expr_simplifier.transforms.constant_folding import apply_constant_folding as apply_constant_folding
+from expr_simplifier.transforms.cse import apply_cse as apply_cse
+from expr_simplifier.transforms.remove_named_expr import apply_remove_named_expr as apply_remove_named_expr

src/expr_simplifier/transforms/constant_folding.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+import ast
+
+from typing_extensions import TypeAlias
+
+from expr_simplifier.transforms.remove_named_expr import apply_remove_named_expr
+
+SubExpressionTable: TypeAlias = dict[str, tuple[str, int]]
+
+
+def fold_to_constant(node: ast.AST) -> ast.Constant:
+    return ast.Constant(value=eval(ast.unparse(node)))
+
+
+class ConstantFolding(ast.NodeTransformer):
+    def visit(self, node: ast.AST) -> ast.AST:
+        transformed_node = super().visit(node)
+        if isinstance(node, ast.BinOp) and isinstance(node.left, ast.Constant) and isinstance(node.right, ast.Constant):
+            return fold_to_constant(node)
+        if isinstance(node, ast.UnaryOp) and isinstance(node.operand, ast.Constant):
+            return fold_to_constant(node)
+        if isinstance(node, ast.BoolOp) and all(isinstance(value, ast.Constant) for value in node.values):
+            return fold_to_constant(node)
+        if isinstance(node, ast.Compare) and all(isinstance(comp, ast.Constant) for comp in node.comparators):
+            return fold_to_constant(node)
+        if isinstance(node, ast.JoinedStr) and all(
+            isinstance(value, ast.Constant)
+            or (isinstance(value, ast.FormattedValue) and isinstance(value.value, ast.Constant))
+            for value in node.values
+        ):
+            return fold_to_constant(node)
+        return transformed_node
+
+
+def apply_constant_folding(expr: ast.AST) -> ast.AST:
+    # Constant propagation
+    expr = apply_remove_named_expr(expr, constant_only=True)
+    # Constant folding
+    expr = ConstantFolding().visit(expr)
+    return expr

src/expr_simplifier/transforms/cse.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+import ast
+
+from typing_extensions import TypeAlias
+
+from expr_simplifier.symbol_table import SymbolTable
+from expr_simplifier.transforms.remove_named_expr import apply_remove_named_expr
+
+SubExpressionTable: TypeAlias = dict[str, tuple[str, int]]
+
+
+class CSEPreAnalyzer(ast.NodeVisitor):
+    def __init__(self):
+        self.subexpressions: SubExpressionTable = {}
+        self.symbols = SymbolTable()
+        super().__init__()
+
+    def visit(self, node: ast.AST) -> None:
+        super().visit(node)
+        expr_string = ast.unparse(node)
+        if isinstance(node, ast.expr):
+            if isinstance(node, ast.Name):
+                self.symbols.define_symbol(node.id)
+                return
+            if isinstance(node, ast.Constant):
+                return
+            if expr_string not in self.subexpressions:
+                self.subexpressions[expr_string] = (self.symbols.request_new_symbol(), 0)
+            symbol, count = self.subexpressions[expr_string]
+            self.subexpressions[expr_string] = (symbol, count + 1)
+
+
+class CommonSubexpressionElimination(ast.NodeTransformer):
+    def __init__(self, subexpressions: dict[str, tuple[str, int]]):
+        self.subexpressions = subexpressions
+        self.declared_symbols = set[str]()
+        super().__init__()
+
+    def visit(self, node: ast.AST) -> ast.expr:
+        expr_string = ast.unparse(node)
+        transformed_node = super().visit(node)
+        if isinstance(node, ast.expr) and expr_string in self.subexpressions:
+            symbol, count = self.subexpressions[expr_string]
+            if count > 1:
+                if symbol not in self.declared_symbols:
+                    self.declared_symbols.add(symbol)
+                    assign_node = ast.NamedExpr(target=ast.Name(id=symbol), value=transformed_node)
+                    return assign_node
+                return ast.Name(id=symbol)
+        return transformed_node
+
+
+def show_subexpressions(subexpressions: SubExpressionTable) -> None:
+    for subexpression, (symbol, count) in subexpressions.items():
+        print(f"{symbol}: {subexpression} ({count})")
+
+
+def apply_cse(expr: ast.AST) -> ast.AST:
+    expr = apply_remove_named_expr(expr)
+    cse_pre_analyzer = CSEPreAnalyzer()
+    cse_pre_analyzer.visit(expr)
+    cse = CommonSubexpressionElimination(cse_pre_analyzer.subexpressions)
+    return cse.visit(expr)

src/expr_simplifier/transforms/remove_named_expr.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+import ast
+import copy
+
+
+class RemoveNamedExpr(ast.NodeTransformer):
+    def __init__(self, constant_only: bool = False) -> None:
+        super().__init__()
+        self.named_expressions = dict[str, ast.expr]()
+        self.constant_only = constant_only
+
+    def should_replace(self, symbol: str) -> bool:
+        return not self.constant_only or isinstance(self.named_expressions[symbol], ast.Constant)
+
+    def visit_NamedExpr(self, node: ast.NamedExpr) -> ast.expr:
+        value = self.visit(node.value)
+        name = node.target.id
+        self.named_expressions[name] = value
+        if not self.should_replace(name):
+            return node
+        return value
+
+    def visit_Name(self, node: ast.Name) -> ast.expr:
+        if node.id in self.named_expressions:
+            if not self.should_replace(node.id):
+                return node
+            return copy.deepcopy(self.named_expressions[node.id])
+        return node
+
+
+def apply_remove_named_expr(expr: ast.AST, constant_only: bool = False) -> ast.AST:
+    remove_named_expr = RemoveNamedExpr(constant_only)
+    return remove_named_expr.visit(expr)

src/expr_simplifier/typing.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+import ast
+from collections.abc import Callable
+
+from typing_extensions import TypeAlias
+
+Pass: TypeAlias = Callable[[ast.AST], ast.AST]