We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[BUG] Optimization that compacts multiple filters into `eval` generates unexpected node in graph
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**

Optimization that compacts multiple filters into eval generates unexpected node in graph.

**To Reproduce**
To help us reproducing this bug, please provide information below:
1. Your Python version
2. The version of Mars you use
3. Versions of crucial packages, such as numpy, scipy and pandas
4. Full stack of the error.
5. Minimized code to reproduce the error.

```python
@enter_mode(build=True)
def test_arithmetic_query(setup):
    df1 = md.DataFrame(raw, chunk_size=10)
    df2 = md.DataFrame(raw2, chunk_size=10)
    df3 = df1.merge(df2, on='A', suffixes=('', '_'))
    df3['K'] = df4 = df3["A"] * (1 - df3["B"])
    graph = TileableGraph([df3.data])
    next(TileableGraphBuilder(graph).build())
    records = optimize(graph)
    opt_df4 = records.get_optimization_result(df4.data)
    assert opt_df4.op.expr == "(`A`) * ((1) - (`B`))"
    assert len(graph) == 5  # for now len(graph) is 6
    assert len([n for n in graph if isinstance(n.op, DataFrameEval)]) == 1  # and 2 evals exist
```


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `mars/optimization/logical/core.py`
Content:
```
1 # Copyright 1999-2021 Alibaba Group Holding Ltd.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #      http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 import weakref
16 from abc import ABC, abstractmethod
17 from collections import defaultdict
18 from dataclasses import dataclass
19 from enum import Enum
20 from typing import Dict, List, Tuple, Type
21 
22 from ...core import OperandType, ChunkType, EntityType, enter_mode
23 from ...core.graph import EntityGraph
24 from ...core.operand import Operand
25 
26 
27 class OptimizationRecordType(Enum):
28     replace = 0
29     new = 1
30     delete = 2
31 
32 
33 @dataclass
34 class OptimizationRecord:
35     original_chunk: ChunkType = None
36     new_chunk: ChunkType = None
37     record_type: OptimizationRecordType = None
38 
39 
40 class OptimizationRecords:
41     _records: List[OptimizationRecord]
42     _original_chunk_to_records: Dict[ChunkType, OptimizationRecord]
43 
44     def __init__(self):
45         self._records = list()
46         self._original_chunk_to_records = dict()
47         self._optimized_chunk_to_records = dict()
48 
49     def append_record(self, record: OptimizationRecord):
50         self._records.append(record)
51         if record.record_type in (
52             OptimizationRecordType.replace,
53             OptimizationRecordType.delete,
54         ):
55             self._original_chunk_to_records[record.original_chunk] = record
56         if record.record_type in (
57             OptimizationRecordType.new,
58             OptimizationRecordType.replace,
59         ):
60             self._optimized_chunk_to_records[record.new_chunk] = record
61 
62     def get_optimization_result(self, original_chunk: ChunkType) -> ChunkType:
63         chunk = original_chunk
64         if chunk not in self._original_chunk_to_records:
65             return
66         while chunk in self._original_chunk_to_records:
67             record = self._original_chunk_to_records[chunk]
68             if record.record_type == OptimizationRecordType.replace:
69                 chunk = record.new_chunk
70             else:
71                 assert record.record_type == OptimizationRecordType.delete
72                 return None
73         return chunk
74 
75     def get_original_chunk(self, optimized_chunk: ChunkType) -> ChunkType:
76         chunk = optimized_chunk
77         if chunk not in self._optimized_chunk_to_records:
78             return
79         while chunk in self._optimized_chunk_to_records:
80             record = self._optimized_chunk_to_records[chunk]
81             if record.record_type == OptimizationRecordType.replace:
82                 chunk = record.original_chunk
83             else:
84                 assert record.record_type == OptimizationRecordType.new
85                 return None
86         return chunk
87 
88 
89 class OptimizationRule(ABC):
90     _instances: Dict[
91         Tuple[Type["OptimizationRule"], EntityGraph, OptimizationRecords],
92         "OptimizationRule",
93     ] = dict()
94     _preds_to_remove = weakref.WeakKeyDictionary()
95 
96     def __init__(
97         self,
98         graph: EntityGraph,
99         records: OptimizationRecords,
100         optimizer_cls: Type["Optimizer"],
101     ):
102         self._graph = graph
103         self._records = records
104         self._optimizer_cls = optimizer_cls
105 
106     def __new__(
107         cls,
108         graph: EntityGraph,
109         records: OptimizationRecords,
110         optimizer_cls: Type["Optimizer"],
111     ):
112         if (cls, graph, records) in cls._instances:
113             return cls._instances[cls, graph, records]
114         inst = cls._instances[cls, graph, records] = object.__new__(cls)
115         return inst
116 
117     @abstractmethod
118     def match(self, op: OperandType) -> bool:
119         """
120         If this operand matches this rule.
121 
122         Parameters
123         ----------
124         op : OperandType
125             Operand.
126 
127         Returns
128         -------
129         matched : bool
130             Matched rule or not.
131         """
132 
133     @abstractmethod
134     def apply(self, op: OperandType):
135         """
136         Apply rule to an operand.
137 
138         Parameters
139         ----------
140         op : OperandType
141             Operand
142         """
143 
144     def _replace_node(self, original_node: EntityType, new_node: EntityType):
145         predecessors = self._graph.predecessors(original_node)
146         successors = self._graph.successors(original_node)
147         self._graph.remove_node(original_node)
148         self._graph.add_node(new_node)
149         for pred in predecessors:
150             self._graph.add_edge(pred, new_node)
151         for succ in successors:
152             self._graph.add_edge(new_node, succ)
153 
154     @classmethod
155     def _add_collapsable_predecessor(cls, node: EntityType, predecessor: EntityType):
156         if predecessor not in cls._preds_to_remove:
157             cls._preds_to_remove[predecessor] = {node}
158         else:
159             cls._preds_to_remove[predecessor].add(node)
160 
161     def _remove_collapsable_predecessors(self, node: EntityType):
162         node = self._records.get_optimization_result(node) or node
163         preds_opt_to_remove = []
164         for pred in self._graph.predecessors(node):
165             pred_original = self._records.get_original_chunk(pred) or pred
166             pred_opt = self._records.get_optimization_result(pred) or pred
167             if pred_opt in self._graph.results or pred_original in self._graph.results:
168                 continue
169             affect_succ = self._preds_to_remove.get(pred_original) or []
170             affect_succ_opt = [
171                 self._records.get_optimization_result(s) or s for s in affect_succ
172             ]
173             if all(s in affect_succ_opt for s in self._graph.successors(pred)):
174                 preds_opt_to_remove.append((pred_original, pred_opt))
175 
176         for pred_original, pred_opt in preds_opt_to_remove:
177             self._graph.remove_node(pred_opt)
178             self._records.append_record(
179                 OptimizationRecord(pred_original, None, OptimizationRecordType.delete)
180             )
181 
182 
183 class Optimizer(ABC):
184     _rules: List[Type[OptimizationRule]]
185     _op_to_rules: Dict[Type[OperandType], List[Type[OptimizationRule]]]
186 
187     @classmethod
188     def register_rule(
189         cls, operand_types: List[Type[OperandType]], rule: Type[OptimizationRule]
190     ):
191         if not hasattr(cls, "_rules"):
192             cls._rules = []
193         cls._rules.append(rule)
194 
195         if not hasattr(cls, "_op_to_rules"):
196             cls._op_to_rules = defaultdict(list)
197         for operand_type in operand_types:
198             cls._op_to_rules[operand_type].append(rule)
199 
200     @classmethod
201     def get_rule_types(
202         cls, operand_type: Type[OperandType]
203     ) -> List[Type[OptimizationRule]]:
204         rule_types = cls._op_to_rules.get(operand_type, None)
205         if rule_types is None:
206             for op_cls in operand_type.__mro__:
207                 if op_cls is Operand:
208                     break
209                 rule_types = cls._op_to_rules.get(op_cls)
210                 if rule_types is not None:
211                     break
212             cls._op_to_rules[operand_type] = rule_types or []
213         return rule_types
214 
215     @classmethod
216     def _replace_inputs(cls, graph: EntityGraph, records: OptimizationRecords):
217         for node in graph:
218             for succ in graph.successors(node):
219                 input_optimized = False
220                 new_inputs = []
221                 for inp in succ.inputs:
222                     optimized = records.get_optimization_result(inp)
223                     if optimized is None:
224                         optimized = inp
225                     if optimized is not inp:
226                         input_optimized = True
227                     new_inputs.append(optimized)
228                 if input_optimized:
229                     succ.inputs = new_inputs
230 
231     @classmethod
232     @enter_mode(build=True)
233     def optimize(cls, graph: EntityGraph) -> OptimizationRecords:
234         """
235         Optimize a graph.
236 
237         Parameters
238         ----------
239         graph : EntityGraph
240             Tileable or chunk graph.
241 
242         Returns
243         -------
244         optimization_records : OptimizationRecords
245             Optimization records.
246         """
247         records = OptimizationRecords()
248         optimized = False
249         for rule_type in cls._rules:
250             visited = set()
251             for entity in list(graph.topological_iter()):
252                 op = entity.op
253                 if op in visited:
254                     continue
255                 visited.add(op)
256 
257                 rule_types = cls.get_rule_types(type(op)) or []
258                 if rule_type not in rule_types:
259                     continue
260 
261                 rule = rule_type(graph, records, cls)
262                 if entity not in graph:  # pragma: no cover
263                     # maybe removed during optimization
264                     continue
265                 if rule.match(op):
266                     optimized = True
267                     rule.apply(op)
268         if optimized:
269             cls._replace_inputs(graph, records)
270         return records
271
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/mars/optimization/logical/core.py b/mars/optimization/logical/core.py
--- a/mars/optimization/logical/core.py
+++ b/mars/optimization/logical/core.py
@@ -162,8 +162,12 @@
         node = self._records.get_optimization_result(node) or node
         preds_opt_to_remove = []
         for pred in self._graph.predecessors(node):
-            pred_original = self._records.get_original_chunk(pred) or pred
-            pred_opt = self._records.get_optimization_result(pred) or pred
+            pred_original = self._records.get_original_chunk(pred)
+            pred_original = pred_original if pred_original is not None else pred
+
+            pred_opt = self._records.get_optimization_result(pred)
+            pred_opt = pred_opt if pred_opt is not None else pred
+
             if pred_opt in self._graph.results or pred_original in self._graph.results:
                 continue
             affect_succ = self._preds_to_remove.get(pred_original) or []

{"golden_diff": "diff --git a/mars/optimization/logical/core.py b/mars/optimization/logical/core.py\n--- a/mars/optimization/logical/core.py\n+++ b/mars/optimization/logical/core.py\n@@ -162,8 +162,12 @@\n         node = self._records.get_optimization_result(node) or node\n         preds_opt_to_remove = []\n         for pred in self._graph.predecessors(node):\n-            pred_original = self._records.get_original_chunk(pred) or pred\n-            pred_opt = self._records.get_optimization_result(pred) or pred\n+            pred_original = self._records.get_original_chunk(pred)\n+            pred_original = pred_original if pred_original is not None else pred\n+\n+            pred_opt = self._records.get_optimization_result(pred)\n+            pred_opt = pred_opt if pred_opt is not None else pred\n+\n             if pred_opt in self._graph.results or pred_original in self._graph.results:\n                 continue\n             affect_succ = self._preds_to_remove.get(pred_original) or []\n", "issue": "[BUG] Optimization that compacts multiple filters into `eval` generates unexpected node in graph\n<!--\r\nThank you for your contribution!\r\n\r\nPlease review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.\r\n-->\r\n\r\n**Describe the bug**\r\n\r\nOptimization that compacts multiple filters into eval generates unexpected node in graph.\r\n\r\n**To Reproduce**\r\nTo help us reproducing this bug, please provide information below:\r\n1. Your Python version\r\n2. The version of Mars you use\r\n3. Versions of crucial packages, such as numpy, scipy and pandas\r\n4. Full stack of the error.\r\n5. Minimized code to reproduce the error.\r\n\r\n```python\r\n@enter_mode(build=True)\r\ndef test_arithmetic_query(setup):\r\n    df1 = md.DataFrame(raw, chunk_size=10)\r\n    df2 = md.DataFrame(raw2, chunk_size=10)\r\n    df3 = df1.merge(df2, on='A', suffixes=('', '_'))\r\n    df3['K'] = df4 = df3[\"A\"] * (1 - df3[\"B\"])\r\n    graph = TileableGraph([df3.data])\r\n    next(TileableGraphBuilder(graph).build())\r\n    records = optimize(graph)\r\n    opt_df4 = records.get_optimization_result(df4.data)\r\n    assert opt_df4.op.expr == \"(`A`) * ((1) - (`B`))\"\r\n    assert len(graph) == 5  # for now len(graph) is 6\r\n    assert len([n for n in graph if isinstance(n.op, DataFrameEval)]) == 1  # and 2 evals exist\r\n```\r\n\n", "before_files": [{"content": "# Copyright 1999-2021 Alibaba Group Holding Ltd.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport weakref\nfrom abc import ABC, abstractmethod\nfrom collections import defaultdict\nfrom dataclasses import dataclass\nfrom enum import Enum\nfrom typing import Dict, List, Tuple, Type\n\nfrom ...core import OperandType, ChunkType, EntityType, enter_mode\nfrom ...core.graph import EntityGraph\nfrom ...core.operand import Operand\n\n\nclass OptimizationRecordType(Enum):\n    replace = 0\n    new = 1\n    delete = 2\n\n\n@dataclass\nclass OptimizationRecord:\n    original_chunk: ChunkType = None\n    new_chunk: ChunkType = None\n    record_type: OptimizationRecordType = None\n\n\nclass OptimizationRecords:\n    _records: List[OptimizationRecord]\n    _original_chunk_to_records: Dict[ChunkType, OptimizationRecord]\n\n    def __init__(self):\n        self._records = list()\n        self._original_chunk_to_records = dict()\n        self._optimized_chunk_to_records = dict()\n\n    def append_record(self, record: OptimizationRecord):\n        self._records.append(record)\n        if record.record_type in (\n            OptimizationRecordType.replace,\n            OptimizationRecordType.delete,\n        ):\n            self._original_chunk_to_records[record.original_chunk] = record\n        if record.record_type in (\n            OptimizationRecordType.new,\n            OptimizationRecordType.replace,\n        ):\n            self._optimized_chunk_to_records[record.new_chunk] = record\n\n    def get_optimization_result(self, original_chunk: ChunkType) -> ChunkType:\n        chunk = original_chunk\n        if chunk not in self._original_chunk_to_records:\n            return\n        while chunk in self._original_chunk_to_records:\n            record = self._original_chunk_to_records[chunk]\n            if record.record_type == OptimizationRecordType.replace:\n                chunk = record.new_chunk\n            else:\n                assert record.record_type == OptimizationRecordType.delete\n                return None\n        return chunk\n\n    def get_original_chunk(self, optimized_chunk: ChunkType) -> ChunkType:\n        chunk = optimized_chunk\n        if chunk not in self._optimized_chunk_to_records:\n            return\n        while chunk in self._optimized_chunk_to_records:\n            record = self._optimized_chunk_to_records[chunk]\n            if record.record_type == OptimizationRecordType.replace:\n                chunk = record.original_chunk\n            else:\n                assert record.record_type == OptimizationRecordType.new\n                return None\n        return chunk\n\n\nclass OptimizationRule(ABC):\n    _instances: Dict[\n        Tuple[Type[\"OptimizationRule\"], EntityGraph, OptimizationRecords],\n        \"OptimizationRule\",\n    ] = dict()\n    _preds_to_remove = weakref.WeakKeyDictionary()\n\n    def __init__(\n        self,\n        graph: EntityGraph,\n        records: OptimizationRecords,\n        optimizer_cls: Type[\"Optimizer\"],\n    ):\n        self._graph = graph\n        self._records = records\n        self._optimizer_cls = optimizer_cls\n\n    def __new__(\n        cls,\n        graph: EntityGraph,\n        records: OptimizationRecords,\n        optimizer_cls: Type[\"Optimizer\"],\n    ):\n        if (cls, graph, records) in cls._instances:\n            return cls._instances[cls, graph, records]\n        inst = cls._instances[cls, graph, records] = object.__new__(cls)\n        return inst\n\n    @abstractmethod\n    def match(self, op: OperandType) -> bool:\n        \"\"\"\n        If this operand matches this rule.\n\n        Parameters\n        ----------\n        op : OperandType\n            Operand.\n\n        Returns\n        -------\n        matched : bool\n            Matched rule or not.\n        \"\"\"\n\n    @abstractmethod\n    def apply(self, op: OperandType):\n        \"\"\"\n        Apply rule to an operand.\n\n        Parameters\n        ----------\n        op : OperandType\n            Operand\n        \"\"\"\n\n    def _replace_node(self, original_node: EntityType, new_node: EntityType):\n        predecessors = self._graph.predecessors(original_node)\n        successors = self._graph.successors(original_node)\n        self._graph.remove_node(original_node)\n        self._graph.add_node(new_node)\n        for pred in predecessors:\n            self._graph.add_edge(pred, new_node)\n        for succ in successors:\n            self._graph.add_edge(new_node, succ)\n\n    @classmethod\n    def _add_collapsable_predecessor(cls, node: EntityType, predecessor: EntityType):\n        if predecessor not in cls._preds_to_remove:\n            cls._preds_to_remove[predecessor] = {node}\n        else:\n            cls._preds_to_remove[predecessor].add(node)\n\n    def _remove_collapsable_predecessors(self, node: EntityType):\n        node = self._records.get_optimization_result(node) or node\n        preds_opt_to_remove = []\n        for pred in self._graph.predecessors(node):\n            pred_original = self._records.get_original_chunk(pred) or pred\n            pred_opt = self._records.get_optimization_result(pred) or pred\n            if pred_opt in self._graph.results or pred_original in self._graph.results:\n                continue\n            affect_succ = self._preds_to_remove.get(pred_original) or []\n            affect_succ_opt = [\n                self._records.get_optimization_result(s) or s for s in affect_succ\n            ]\n            if all(s in affect_succ_opt for s in self._graph.successors(pred)):\n                preds_opt_to_remove.append((pred_original, pred_opt))\n\n        for pred_original, pred_opt in preds_opt_to_remove:\n            self._graph.remove_node(pred_opt)\n            self._records.append_record(\n                OptimizationRecord(pred_original, None, OptimizationRecordType.delete)\n            )\n\n\nclass Optimizer(ABC):\n    _rules: List[Type[OptimizationRule]]\n    _op_to_rules: Dict[Type[OperandType], List[Type[OptimizationRule]]]\n\n    @classmethod\n    def register_rule(\n        cls, operand_types: List[Type[OperandType]], rule: Type[OptimizationRule]\n    ):\n        if not hasattr(cls, \"_rules\"):\n            cls._rules = []\n        cls._rules.append(rule)\n\n        if not hasattr(cls, \"_op_to_rules\"):\n            cls._op_to_rules = defaultdict(list)\n        for operand_type in operand_types:\n            cls._op_to_rules[operand_type].append(rule)\n\n    @classmethod\n    def get_rule_types(\n        cls, operand_type: Type[OperandType]\n    ) -> List[Type[OptimizationRule]]:\n        rule_types = cls._op_to_rules.get(operand_type, None)\n        if rule_types is None:\n            for op_cls in operand_type.__mro__:\n                if op_cls is Operand:\n                    break\n                rule_types = cls._op_to_rules.get(op_cls)\n                if rule_types is not None:\n                    break\n            cls._op_to_rules[operand_type] = rule_types or []\n        return rule_types\n\n    @classmethod\n    def _replace_inputs(cls, graph: EntityGraph, records: OptimizationRecords):\n        for node in graph:\n            for succ in graph.successors(node):\n                input_optimized = False\n                new_inputs = []\n                for inp in succ.inputs:\n                    optimized = records.get_optimization_result(inp)\n                    if optimized is None:\n                        optimized = inp\n                    if optimized is not inp:\n                        input_optimized = True\n                    new_inputs.append(optimized)\n                if input_optimized:\n                    succ.inputs = new_inputs\n\n    @classmethod\n    @enter_mode(build=True)\n    def optimize(cls, graph: EntityGraph) -> OptimizationRecords:\n        \"\"\"\n        Optimize a graph.\n\n        Parameters\n        ----------\n        graph : EntityGraph\n            Tileable or chunk graph.\n\n        Returns\n        -------\n        optimization_records : OptimizationRecords\n            Optimization records.\n        \"\"\"\n        records = OptimizationRecords()\n        optimized = False\n        for rule_type in cls._rules:\n            visited = set()\n            for entity in list(graph.topological_iter()):\n                op = entity.op\n                if op in visited:\n                    continue\n                visited.add(op)\n\n                rule_types = cls.get_rule_types(type(op)) or []\n                if rule_type not in rule_types:\n                    continue\n\n                rule = rule_type(graph, records, cls)\n                if entity not in graph:  # pragma: no cover\n                    # maybe removed during optimization\n                    continue\n                if rule.match(op):\n                    optimized = True\n                    rule.apply(op)\n        if optimized:\n            cls._replace_inputs(graph, records)\n        return records\n", "path": "mars/optimization/logical/core.py"}], "after_files": [{"content": "# Copyright 1999-2021 Alibaba Group Holding Ltd.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport weakref\nfrom abc import ABC, abstractmethod\nfrom collections import defaultdict\nfrom dataclasses import dataclass\nfrom enum import Enum\nfrom typing import Dict, List, Tuple, Type\n\nfrom ...core import OperandType, ChunkType, EntityType, enter_mode\nfrom ...core.graph import EntityGraph\nfrom ...core.operand import Operand\n\n\nclass OptimizationRecordType(Enum):\n    replace = 0\n    new = 1\n    delete = 2\n\n\n@dataclass\nclass OptimizationRecord:\n    original_chunk: ChunkType = None\n    new_chunk: ChunkType = None\n    record_type: OptimizationRecordType = None\n\n\nclass OptimizationRecords:\n    _records: List[OptimizationRecord]\n    _original_chunk_to_records: Dict[ChunkType, OptimizationRecord]\n\n    def __init__(self):\n        self._records = list()\n        self._original_chunk_to_records = dict()\n        self._optimized_chunk_to_records = dict()\n\n    def append_record(self, record: OptimizationRecord):\n        self._records.append(record)\n        if record.record_type in (\n            OptimizationRecordType.replace,\n            OptimizationRecordType.delete,\n        ):\n            self._original_chunk_to_records[record.original_chunk] = record\n        if record.record_type in (\n            OptimizationRecordType.new,\n            OptimizationRecordType.replace,\n        ):\n            self._optimized_chunk_to_records[record.new_chunk] = record\n\n    def get_optimization_result(self, original_chunk: ChunkType) -> ChunkType:\n        chunk = original_chunk\n        if chunk not in self._original_chunk_to_records:\n            return\n        while chunk in self._original_chunk_to_records:\n            record = self._original_chunk_to_records[chunk]\n            if record.record_type == OptimizationRecordType.replace:\n                chunk = record.new_chunk\n            else:\n                assert record.record_type == OptimizationRecordType.delete\n                return None\n        return chunk\n\n    def get_original_chunk(self, optimized_chunk: ChunkType) -> ChunkType:\n        chunk = optimized_chunk\n        if chunk not in self._optimized_chunk_to_records:\n            return\n        while chunk in self._optimized_chunk_to_records:\n            record = self._optimized_chunk_to_records[chunk]\n            if record.record_type == OptimizationRecordType.replace:\n                chunk = record.original_chunk\n            else:\n                assert record.record_type == OptimizationRecordType.new\n                return None\n        return chunk\n\n\nclass OptimizationRule(ABC):\n    _instances: Dict[\n        Tuple[Type[\"OptimizationRule\"], EntityGraph, OptimizationRecords],\n        \"OptimizationRule\",\n    ] = dict()\n    _preds_to_remove = weakref.WeakKeyDictionary()\n\n    def __init__(\n        self,\n        graph: EntityGraph,\n        records: OptimizationRecords,\n        optimizer_cls: Type[\"Optimizer\"],\n    ):\n        self._graph = graph\n        self._records = records\n        self._optimizer_cls = optimizer_cls\n\n    def __new__(\n        cls,\n        graph: EntityGraph,\n        records: OptimizationRecords,\n        optimizer_cls: Type[\"Optimizer\"],\n    ):\n        if (cls, graph, records) in cls._instances:\n            return cls._instances[cls, graph, records]\n        inst = cls._instances[cls, graph, records] = object.__new__(cls)\n        return inst\n\n    @abstractmethod\n    def match(self, op: OperandType) -> bool:\n        \"\"\"\n        If this operand matches this rule.\n\n        Parameters\n        ----------\n        op : OperandType\n            Operand.\n\n        Returns\n        -------\n        matched : bool\n            Matched rule or not.\n        \"\"\"\n\n    @abstractmethod\n    def apply(self, op: OperandType):\n        \"\"\"\n        Apply rule to an operand.\n\n        Parameters\n        ----------\n        op : OperandType\n            Operand\n        \"\"\"\n\n    def _replace_node(self, original_node: EntityType, new_node: EntityType):\n        predecessors = self._graph.predecessors(original_node)\n        successors = self._graph.successors(original_node)\n        self._graph.remove_node(original_node)\n        self._graph.add_node(new_node)\n        for pred in predecessors:\n            self._graph.add_edge(pred, new_node)\n        for succ in successors:\n            self._graph.add_edge(new_node, succ)\n\n    @classmethod\n    def _add_collapsable_predecessor(cls, node: EntityType, predecessor: EntityType):\n        if predecessor not in cls._preds_to_remove:\n            cls._preds_to_remove[predecessor] = {node}\n        else:\n            cls._preds_to_remove[predecessor].add(node)\n\n    def _remove_collapsable_predecessors(self, node: EntityType):\n        node = self._records.get_optimization_result(node) or node\n        preds_opt_to_remove = []\n        for pred in self._graph.predecessors(node):\n            pred_original = self._records.get_original_chunk(pred)\n            pred_original = pred_original if pred_original is not None else pred\n\n            pred_opt = self._records.get_optimization_result(pred)\n            pred_opt = pred_opt if pred_opt is not None else pred\n\n            if pred_opt in self._graph.results or pred_original in self._graph.results:\n                continue\n            affect_succ = self._preds_to_remove.get(pred_original) or []\n            affect_succ_opt = [\n                self._records.get_optimization_result(s) or s for s in affect_succ\n            ]\n            if all(s in affect_succ_opt for s in self._graph.successors(pred)):\n                preds_opt_to_remove.append((pred_original, pred_opt))\n\n        for pred_original, pred_opt in preds_opt_to_remove:\n            self._graph.remove_node(pred_opt)\n            self._records.append_record(\n                OptimizationRecord(pred_original, None, OptimizationRecordType.delete)\n            )\n\n\nclass Optimizer(ABC):\n    _rules: List[Type[OptimizationRule]]\n    _op_to_rules: Dict[Type[OperandType], List[Type[OptimizationRule]]]\n\n    @classmethod\n    def register_rule(\n        cls, operand_types: List[Type[OperandType]], rule: Type[OptimizationRule]\n    ):\n        if not hasattr(cls, \"_rules\"):\n            cls._rules = []\n        cls._rules.append(rule)\n\n        if not hasattr(cls, \"_op_to_rules\"):\n            cls._op_to_rules = defaultdict(list)\n        for operand_type in operand_types:\n            cls._op_to_rules[operand_type].append(rule)\n\n    @classmethod\n    def get_rule_types(\n        cls, operand_type: Type[OperandType]\n    ) -> List[Type[OptimizationRule]]:\n        rule_types = cls._op_to_rules.get(operand_type, None)\n        if rule_types is None:\n            for op_cls in operand_type.__mro__:\n                if op_cls is Operand:\n                    break\n                rule_types = cls._op_to_rules.get(op_cls)\n                if rule_types is not None:\n                    break\n            cls._op_to_rules[operand_type] = rule_types or []\n        return rule_types\n\n    @classmethod\n    def _replace_inputs(cls, graph: EntityGraph, records: OptimizationRecords):\n        for node in graph:\n            for succ in graph.successors(node):\n                input_optimized = False\n                new_inputs = []\n                for inp in succ.inputs:\n                    optimized = records.get_optimization_result(inp)\n                    if optimized is None:\n                        optimized = inp\n                    if optimized is not inp:\n                        input_optimized = True\n                    new_inputs.append(optimized)\n                if input_optimized:\n                    succ.inputs = new_inputs\n\n    @classmethod\n    @enter_mode(build=True)\n    def optimize(cls, graph: EntityGraph) -> OptimizationRecords:\n        \"\"\"\n        Optimize a graph.\n\n        Parameters\n        ----------\n        graph : EntityGraph\n            Tileable or chunk graph.\n\n        Returns\n        -------\n        optimization_records : OptimizationRecords\n            Optimization records.\n        \"\"\"\n        records = OptimizationRecords()\n        optimized = False\n        for rule_type in cls._rules:\n            visited = set()\n            for entity in list(graph.topological_iter()):\n                op = entity.op\n                if op in visited:\n                    continue\n                visited.add(op)\n\n                rule_types = cls.get_rule_types(type(op)) or []\n                if rule_type not in rule_types:\n                    continue\n\n                rule = rule_type(graph, records, cls)\n                if entity not in graph:  # pragma: no cover\n                    # maybe removed during optimization\n                    continue\n                if rule.match(op):\n                    optimized = True\n                    rule.apply(op)\n        if optimized:\n            cls._replace_inputs(graph, records)\n        return records\n", "path": "mars/optimization/logical/core.py"}]}