✨ feat: remove unused NamedExpr in CSE

Author: SigureMo

pyproject.toml

@@ -4,9 +4,7 @@ version = "0.1.0"
 description = ""
 readme = "README.md"
 requires-python = ">=3.9"
-dependencies = [
-    "typing-extensions>=4.12.2",
-]
+dependencies = ["typing-extensions>=4.12.2"]
 authors = [{ name = "Nyakku Shigure", email = "sigure.qaq@gmail.com" }]
 keywords = []
 license = { text = "MIT" }
@@ -77,6 +75,7 @@ ignore = [
 [tool.ruff.lint.isort]
 required-imports = ["from __future__ import annotations"]
 known-first-party = ["expr_simplifier"]
+combine-as-imports = true
 
 [tool.ruff.lint.per-file-ignores]
 "setup.py" = ["I"]

src/expr_simplifier/transforms/__init__.py

@@ -2,4 +2,10 @@
 
 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.inline_named_expr import apply_inline_named_expr as apply_inline_named_expr
+from expr_simplifier.transforms.inline_named_expr import (
+    apply_constant_propagation as apply_constant_propagation,
+    apply_inline_all_named_expr as apply_inline_all_named_expr,
+)
+from expr_simplifier.transforms.remove_unused_named_expr import (
+    apply_remove_unused_named_expr as apply_remove_unused_named_expr,
+)

src/expr_simplifier/transforms/constant_folding.py

@@ -4,7 +4,7 @@
 
 from typing_extensions import TypeAlias
 
-from expr_simplifier.transforms.inline_named_expr import apply_inline_named_expr
+from expr_simplifier.transforms.inline_named_expr import apply_constant_propagation
 
 SubExpressionTable: TypeAlias = dict[str, tuple[str, int]]
 
@@ -35,7 +35,7 @@ def visit(self, node: ast.AST) -> ast.AST:
 
 def apply_constant_folding(expr: ast.AST) -> ast.AST:
     # Constant propagation
-    expr = apply_inline_named_expr(expr, constant_only=True)
+    expr = apply_constant_propagation(expr)
     # Constant folding
     expr = ConstantFolding().visit(expr)
     return expr

src/expr_simplifier/transforms/cse.py

@@ -5,7 +5,8 @@
 from typing_extensions import TypeAlias
 
 from expr_simplifier.symbol_table import SymbolTable
-from expr_simplifier.transforms.inline_named_expr import apply_inline_named_expr
+from expr_simplifier.transforms.inline_named_expr import apply_inline_all_named_expr
+from expr_simplifier.transforms.remove_unused_named_expr import apply_remove_unused_named_expr
 
 SubExpressionTable: TypeAlias = dict[str, tuple[str, int]]
 
@@ -45,9 +46,9 @@ def visit(self, node: ast.AST) -> ast.expr:
             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)
+                    assign_node = ast.NamedExpr(target=ast.Name(id=symbol, ctx=ast.Store()), value=transformed_node)
                     return assign_node
-                return ast.Name(id=symbol)
+                return ast.Name(id=symbol, ctx=ast.Load())
         return transformed_node
 
 
@@ -57,8 +58,9 @@ def show_subexpressions(subexpressions: SubExpressionTable) -> None:
 
 
 def apply_cse(expr: ast.AST) -> ast.AST:
-    expr = apply_inline_named_expr(expr)
+    expr = apply_inline_all_named_expr(expr)
     cse_pre_analyzer = CSEPreAnalyzer()
     cse_pre_analyzer.visit(expr)
     cse = CommonSubexpressionElimination(cse_pre_analyzer.subexpressions)
-    return cse.visit(expr)
+    expr = cse.visit(expr)
+    return apply_remove_unused_named_expr(expr)

src/expr_simplifier/transforms/inline_named_expr.py

@@ -2,33 +2,43 @@
 
 import ast
 import copy
+from collections.abc import Callable
+
+from typing_extensions import TypeAlias
+
+NamedExpressions: TypeAlias = dict[str, ast.expr]
+ShouldReplaceFn: TypeAlias = Callable[[str, NamedExpressions], bool]
 
 
 class InlineNamedExpr(ast.NodeTransformer):
-    def __init__(self, constant_only: bool = False) -> None:
+    def __init__(self, should_replace_fn: ShouldReplaceFn) -> None:
         super().__init__()
+        self.should_replace_fn = should_replace_fn
         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):
+        if not self.should_replace_fn(name, self.named_expressions):
             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):
+            if not self.should_replace_fn(node.id, self.named_expressions):
                 return node
             return copy.deepcopy(self.named_expressions[node.id])
         return node
 
 
-def apply_inline_named_expr(expr: ast.AST, constant_only: bool = False) -> ast.AST:
-    inline_named_expr = InlineNamedExpr(constant_only)
+def apply_inline_all_named_expr(expr: ast.AST) -> ast.AST:
+    inline_named_expr = InlineNamedExpr(lambda symbol, named_expressions: True)
+    return inline_named_expr.visit(expr)
+
+
+def apply_constant_propagation(expr: ast.AST) -> ast.AST:
+    inline_named_expr = InlineNamedExpr(
+        lambda symbol, named_expressions: isinstance(named_expressions[symbol], ast.Constant)
+    )
     return inline_named_expr.visit(expr)

src/expr_simplifier/transforms/remove_unused_named_expr.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import ast
+
+
+class UsedSymbolsAnalyzer(ast.NodeVisitor):
+    def __init__(self) -> None:
+        super().__init__()
+        self.used_symbols = set[str]()
+
+    def visit_Name(self, node: ast.Name) -> None:
+        symbol = node.id
+        if isinstance(node.ctx, ast.Load):
+            self.used_symbols.add(symbol)
+
+
+class RemoveUnusedNamedExpr(ast.NodeTransformer):
+    def __init__(self, used_symbols: set[str]) -> None:
+        super().__init__()
+        self.used_symbols = used_symbols
+
+    def visit_NamedExpr(self, node: ast.NamedExpr) -> ast.expr:
+        value = self.visit(node.value)
+        name = node.target.id
+        if name not in self.used_symbols:
+            return value
+        return node
+
+
+def apply_remove_unused_named_expr(expr: ast.AST) -> ast.AST:
+    used_symbols_analyzer = UsedSymbolsAnalyzer()
+    used_symbols_analyzer.visit(expr)
+    used_symbols = used_symbols_analyzer.used_symbols
+
+    remove_unused_named_expr = RemoveUnusedNamedExpr(used_symbols)
+    return remove_unused_named_expr.visit(expr)

tests/test_transforms/test_cse.py

@@ -19,15 +19,15 @@
         ("1 + 1 + 2", "1 + 1 + 2"),
         (
             "(___x := a.b).c + ___x.d.e + (___y := (___z := a.b).c) + ___z.d.e + ___y.f.g",
-            "(___t_1 := (___t_0 := a.b).c) + (___t_3 := (___t_2 := ___t_0.d).e) + ___t_1 + ___t_3 + ___t_1.f.g",
+            "(___t_1 := (___t_0 := a.b).c) + (___t_3 := ___t_0.d.e) + ___t_1 + ___t_3 + ___t_1.f.g",
         ),
         (
             "a.b.c and (fn1(fn2((___x := a.b.c.d.e.f.g)))) and (___x == False) and (___y := a.b.c)",
-            "(___t_1 := (___t_0 := a.b).c) and fn1(fn2((___t_5 := (___t_4 := (___t_3 := (___t_2 := ___t_1.d).e).f).g))) and (___t_5 == False) and ___t_1",
+            "(___t_1 := a.b.c) and fn1(fn2((___t_5 := ___t_1.d.e.f.g))) and (___t_5 == False) and ___t_1",
         ),
         (
             "(___x := a.b.c) + ___x",
-            "(___t_1 := (___t_0 := a.b).c) + ___t_1",
+            "(___t_1 := a.b.c) + ___t_1",
         ),
     ],
 )

tests/test_transforms/test_inline_named_expr.py

@@ -4,7 +4,10 @@
 
 import pytest
 
-from expr_simplifier.transforms import apply_inline_named_expr
+from expr_simplifier.transforms import (
+    apply_constant_propagation,
+    apply_inline_all_named_expr,
+)
 
 
 @pytest.mark.parametrize(
@@ -19,7 +22,7 @@
 )
 def test_inline_named_expr(expr: str, expected: str):
     tree = ast.parse(expr, mode="eval")
-    transformed_tree = apply_inline_named_expr(tree)
+    transformed_tree = apply_inline_all_named_expr(tree)
     transformed_expr = ast.unparse(transformed_tree)
     assert transformed_expr == expected
 
@@ -37,6 +40,6 @@ def test_inline_named_expr(expr: str, expected: str):
 )
 def test_constant_propagation(expr: str, expected: str):
     tree = ast.parse(expr, mode="eval")
-    transformed_tree = apply_inline_named_expr(tree, constant_only=True)
+    transformed_tree = apply_constant_propagation(tree)
     transformed_expr = ast.unparse(transformed_tree)
     assert transformed_expr == expected

tests/test_transforms/test_remove_unused_named_expr.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+import ast
+
+import pytest
+
+from expr_simplifier.transforms import apply_remove_unused_named_expr
+
+
+@pytest.mark.parametrize(
+    ["expr", "expected"],
+    [
+        ("(___x := a.b) + ___x", "a.b + ___x"),
+        ("(___y := (___x := a.b)) + ___y", "(___y := a.b) + ___y"),
+        ("(___y := (___x := a.b))", "a.b"),
+    ],
+)
+def test_inline_named_expr(expr: str, expected: str):
+    tree = ast.parse(expr, mode="eval")
+    transformed_tree = apply_remove_unused_named_expr(tree)
+    transformed_expr = ast.unparse(transformed_tree)
+    assert transformed_expr == expected