✨ feat: support logical_simplification

SigureMo · SigureMo · commit edc198b4cb7d · 2024-09-02T02:45:16.000+08:00
diff --git a/README.md b/README.md
@@ -13,22 +13,27 @@
    <a href="https://gitmoji.dev"><img alt="Gitmoji" src="https://img.shields.io/badge/gitmoji-%20😜%20😍-FFDD67?style=flat-square"></a>
 </p>
 
-## Installation
+## Usage
 
-TODO...
+### Quick start in CLI
 
-## Usage
+We recommend using the tool `uv` to run the without manually installing the package:
 
 ```console
-$ expr_simplifier cse "a * 4 + (a * 4)"
+$ uvx expr_simplifier cse "a * 4 + (a * 4)"
 (___t_0 := (a * 4)) + ___t_0
-$ expr_simplifier constant_folding "(___x := 1 + 1) + ___x" --max-iter=1
+$ uvx expr_simplifier constant_folding "(___x := 1 + 1) + ___x" --max-iter=1
 (___x := 2) + ___x
-$ expr_simplifier constant_folding "(___x := 1 + 1) + ___x" --max-iter=2
+$ uvx expr_simplifier constant_folding "(___x := 1 + 1) + ___x" --max-iter=2
 4
+# uvx expr_simplifier logical_simplification "a and b and a"
+a and b
 ```
 
+### As a library
+
+TODO...
+
 ## TODOs
 
--  [ ] Fold same logic operations (`a and b and a` -> `a and b`)
 -  [ ] Add runtime checks in uts
diff --git a/src/expr_simplifier/__main__.py b/src/expr_simplifier/__main__.py
@@ -5,7 +5,7 @@
 from collections.abc import Callable
 
 from expr_simplifier import __version__
-from expr_simplifier.transforms import apply_constant_folding, apply_cse
+from expr_simplifier.transforms import apply_constant_folding, apply_cse, apply_logical_simplification
 from expr_simplifier.typing import Pass
 from expr_simplifier.utils import loop_until_stable
 
@@ -39,7 +39,10 @@ def main() -> None:
 
     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)
+    create_pass_parser("logical_simplification", [apply_logical_simplification], "Logical Simplification", sub_parsers)
+    create_pass_parser(
+        "auto", [apply_constant_folding, apply_logical_simplification, apply_cse], "Auto Simplification", sub_parsers
+    )
 
     args = parser.parse_args()
     args.func(args)
diff --git a/src/expr_simplifier/transforms/__init__.py b/src/expr_simplifier/transforms/__init__.py
@@ -6,6 +6,11 @@
     apply_constant_propagation as apply_constant_propagation,
     apply_inline_all_named_expr as apply_inline_all_named_expr,
 )
+from expr_simplifier.transforms.logical_simplification import (
+    apply_logical_short_circuiting as apply_logical_short_circuiting,
+    apply_logical_simplification as apply_logical_simplification,
+    apply_remove_same_subexpression_in_logical_op as apply_remove_same_subexpression_in_logical_op,
+)
 from expr_simplifier.transforms.remove_unused_named_expr import (
     apply_remove_unused_named_expr as apply_remove_unused_named_expr,
 )
diff --git a/src/expr_simplifier/transforms/logical_simplification.py b/src/expr_simplifier/transforms/logical_simplification.py
@@ -0,0 +1,132 @@
+from __future__ import annotations
+
+import ast
+
+
+def create_logical_values_node(op: ast.boolop, values: list[ast.expr], default_value: bool) -> ast.expr:
+    if not values:
+        return ast.Constant(value=default_value)
+    if len(values) == 1:
+        return values[0]
+    return ast.BoolOp(op=op, values=values)
+
+
+class LogicalShortCircuiting(ast.NodeTransformer):
+    def visit_BoolOp(self, node: ast.BoolOp) -> ast.expr:
+        if isinstance(node.op, ast.And):
+            return self._visit_and(node)
+        elif isinstance(node.op, ast.Or):
+            return self._visit_or(node)
+        return node
+
+    def _visit_and(self, node: ast.BoolOp) -> ast.expr:
+        and_values: list[ast.expr] = []
+        for value in node.values:
+            new_value = self.visit(value)
+            if isinstance(new_value, ast.Constant):
+                if new_value.value is False:
+                    return new_value
+                if new_value.value is True:
+                    continue
+            and_values.append(new_value)
+        return create_logical_values_node(ast.And(), and_values, True)
+
+    def _visit_or(self, node: ast.BoolOp) -> ast.expr:
+        or_values: list[ast.expr] = []
+        for value in node.values:
+            new_value = self.visit(value)
+            if isinstance(new_value, ast.Constant):
+                if new_value.value is True:
+                    return new_value
+                if new_value.value is False:
+                    continue
+            or_values.append(new_value)
+        return create_logical_values_node(ast.Or(), or_values, False)
+
+
+class LogicCluster:
+    def __init__(self, ids: set[int]) -> None:
+        self.ids = ids
+        self.subexpressions = set[str]()
+
+    def match(self, other: LogicCluster) -> bool:
+        return bool(self.ids & other.ids)
+
+    def add_expression(self, expr: str) -> None:
+        self.subexpressions.add(expr)
+
+    def update(self, other: LogicCluster) -> None:
+        self.ids.update(other.ids)
+        self.subexpressions.update(other.subexpressions)
+
+    def __repr__(self) -> str:
+        return f"LogicCluster(ids={self.ids}, subexpressions={self.subexpressions})"
+
+
+class RemoveSameSubExpressionInLogicalOp(ast.NodeTransformer):
+    def __init__(self):
+        self.and_clusters: list[LogicCluster] = []
+        self.or_clusters: list[LogicCluster] = []
+
+    @staticmethod
+    def update_cluster(regisitry: list[LogicCluster], new_cluster: LogicCluster) -> LogicCluster:
+        for cluster in regisitry:
+            if cluster.match(new_cluster):
+                cluster.update(new_cluster)
+                return cluster
+        regisitry.append(new_cluster)
+        return new_cluster
+
+    def visit_BoolOp(self, node: ast.BoolOp) -> ast.expr:
+        if isinstance(node.op, ast.And):
+            return self._visit_and(node)
+        elif isinstance(node.op, ast.Or):
+            return self._visit_or(node)
+        return node
+
+    def _visit_and(self, node: ast.BoolOp) -> ast.expr:
+        node_id = id(node)
+        cluster_ids = {id(value) for value in node.values}
+        cluster = LogicCluster({*cluster_ids, node_id})
+        cluster = self.update_cluster(self.and_clusters, cluster)
+        transformed_values: list[ast.expr] = []
+        for value in node.values:
+            known_exprs = cluster.subexpressions.copy()
+            transformed_value = self.visit(value)
+            stringified_value = ast.unparse(transformed_value)
+            if stringified_value not in known_exprs:
+                cluster.add_expression(stringified_value)
+                transformed_values.append(transformed_value)
+        return create_logical_values_node(ast.And(), transformed_values, True)
+
+    def _visit_or(self, node: ast.BoolOp) -> ast.expr:
+        node_id = id(node)
+        cluster_ids = {id(value) for value in node.values}
+        cluster = LogicCluster({*cluster_ids, node_id})
+        cluster = self.update_cluster(self.or_clusters, cluster)
+        transformed_values: list[ast.expr] = []
+        for value in node.values:
+            known_exprs = cluster.subexpressions.copy()
+            transformed_value = self.visit(value)
+            stringified_value = ast.unparse(transformed_value)
+            if stringified_value not in known_exprs:
+                cluster.add_expression(stringified_value)
+                transformed_values.append(transformed_value)
+        return create_logical_values_node(ast.Or(), transformed_values, False)
+
+
+def apply_logical_short_circuiting(expr: ast.AST) -> ast.AST:
+    logical_short_circuiting = LogicalShortCircuiting()
+    return logical_short_circuiting.visit(expr)
+
+
+def apply_remove_same_subexpression_in_logical_op(expr: ast.AST) -> ast.AST:
+    remove_same_subexpression_in_logical_op = RemoveSameSubExpressionInLogicalOp()
+    res = remove_same_subexpression_in_logical_op.visit(expr)
+    return res
+
+
+def apply_logical_simplification(expr: ast.AST) -> ast.AST:
+    expr = apply_remove_same_subexpression_in_logical_op(expr)
+    expr = apply_logical_short_circuiting(expr)
+    return expr
diff --git a/tests/test_transforms/test_logical_simplification.py b/tests/test_transforms/test_logical_simplification.py
@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+import ast
+
+import pytest
+
+from expr_simplifier.transforms import (
+    apply_logical_short_circuiting,
+    apply_logical_simplification,
+    apply_remove_same_subexpression_in_logical_op,
+)
+
+
+@pytest.mark.parametrize(
+    ["expr", "expected"],
+    [
+        ("a and False and b", "False"),
+        ("(a and False) and b", "False"),
+        ("a and (False and b)", "False"),
+        ("a or True or b", "True"),
+        ("(a or True) or b", "True"),
+        ("a or (True or b)", "True"),
+        ("a and True and b", "a and b"),
+        ("a or False or b", "a or b"),
+        ("a and True or b", "a or b"),
+        ("not (a and False)", "not False"),
+        ("True and True", "True"),
+        ("True or True", "True"),
+        ("False and False", "False"),
+        ("False or False", "False"),
+    ],
+)
+def test_logical_short_circuiting(expr: str, expected: str):
+    tree = ast.parse(expr, mode="eval")
+    transformed_tree = apply_logical_short_circuiting(tree)
+    transformed_expr = ast.unparse(transformed_tree)
+    assert transformed_expr == expected
+
+
+@pytest.mark.parametrize(
+    ["expr", "expected"],
+    [
+        ("a and b and a", "a and b"),
+        ("(a and b) and a", "a and b"),
+        ("a and (b and a)", "a and b"),
+        ("a and b and a and b", "a and b"),
+        ("(a and b) and (a and b)", "(a and b) and True"),
+        ("a or b or a", "a or b"),
+        ("(a or b) or a", "a or b"),
+        ("a or (b or a)", "a or b"),
+        ("a or b or a or b", "a or b"),
+        ("(a or b) or (a or b)", "(a or b) or False"),
+        ("a and b and c and a and b and c", "a and b and c"),
+        ("(a and b) and (c and a) and b and c", "(a and b) and c"),
+        ("a and (b and c) and a and b and c", "a and (b and c)"),
+    ],
+)
+def test_remove_same_subexpression_in_logical_op(expr: str, expected: str):
+    tree = ast.parse(expr, mode="eval")
+    transformed_tree = apply_remove_same_subexpression_in_logical_op(tree)
+    transformed_expr = ast.unparse(transformed_tree)
+    assert transformed_expr == expected
+
+
+@pytest.mark.parametrize(
+    ["expr", "expected"],
+    [
+        ("(a and b) and (a and b)", "a and b"),
+        ("(a or b) or (a or b)", "a or b"),
+    ],
+)
+def test_logical_simplification(expr: str, expected: str):
+    tree = ast.parse(expr, mode="eval")
+    transformed_tree = apply_logical_simplification(tree)
+    transformed_expr = ast.unparse(transformed_tree)
+    assert transformed_expr == expected

Original file line number	Diff line number	Diff line change
`@@ -6,6 +6,11 @@`
`6`	`6`	`apply_constant_propagation as apply_constant_propagation,`
`7`	`7`	`apply_inline_all_named_expr as apply_inline_all_named_expr,`
`8`	`8`	`)`
	`9`	`+from expr_simplifier.transforms.logical_simplification import (`
	`10`	`+ apply_logical_short_circuiting as apply_logical_short_circuiting,`
	`11`	`+ apply_logical_simplification as apply_logical_simplification,`
	`12`	`+ apply_remove_same_subexpression_in_logical_op as apply_remove_same_subexpression_in_logical_op,`
	`13`	`+)`
`9`	`14`	`from expr_simplifier.transforms.remove_unused_named_expr import (`
`10`	`15`	`apply_remove_unused_named_expr as apply_remove_unused_named_expr,`
`11`	`16`	`)`