10943311926

Committed 19 Sep 2024 02:41PM CUT coverage: 98.89%. Remained the same

Build # 10943311926

Build Type

Pull #1062

github

Committed by

web-flow

Commit Message

Merge 1e881487b into 86c59ead1

Pull Request Pull Request #1062: Prepare for release v0.2.3

Run Details

1903 of 1910 branches covered (99.63%)

Branch coverage included in aggregate %.

8695 of 8807 relevant lines covered (98.73%)

0.99 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

96.24

/src/basilisp/lang/compiler/optimizer.py

import ast
import functools
from collections import deque
from contextlib import contextmanager
from typing import Deque, Iterable, List, Optional, Set

from basilisp.lang.compiler.constants import OPERATOR_ALIAS
from basilisp.lang.compiler.utils import ast_FunctionDef, ast_index


def _filter_dead_code(nodes: Iterable[ast.stmt]) -> List[ast.stmt]:
    """Return a list of body nodes, trimming out unreachable code (any
    statements appearing after `break`, `continue`, `raise`, and `return`
    nodes)."""
    new_nodes: List[ast.stmt] = []
    for node in nodes:
        if isinstance(node, (ast.Break, ast.Continue, ast.Raise, ast.Return)):
            new_nodes.append(node)
            break
        new_nodes.append(node)
    return new_nodes


def _needs_eq_operator(arg: ast.expr) -> bool:
    return isinstance(arg, ast.Constant) and all(
        arg.value is not v for v in (True, False, None, ...)
    )


@functools.singledispatch
def _optimize_operator_call(  # pylint: disable=unused-argument
    fn: ast.AST, node: ast.Call
) -> ast.AST:
    return node


@_optimize_operator_call.register(ast.Attribute)
def _optimize_operator_call_attr(  # pylint: disable=too-many-return-statements
    fn: ast.Attribute, node: ast.Call
) -> ast.AST:
    """Optimize calls to the Python `operator` module down to use the raw Python
    operators.

    Using Python operators directly will allow for more direct bytecode to be
    emitted by the Python compiler and take advantage of any additional performance
    improvements in future versions of Python."""
    if isinstance(fn.value, ast.Name) and fn.value.id == OPERATOR_ALIAS:
        binop = {
            "add": ast.Add,
            "and_": ast.BitAnd,
            "floordiv": ast.FloorDiv,
            "lshift": ast.LShift,
            "mod": ast.Mod,
            "mul": ast.Mult,
            "matmul": ast.MatMult,
            "or_": ast.BitOr,
            "pow": ast.Pow,
            "rshift": ast.RShift,
            "sub": ast.Sub,
            "truediv": ast.Div,
            "xor": ast.BitXor,
        }.get(fn.attr)
        if binop is not None:
            arg1, arg2 = node.args
            assert len(node.args) == 2
            return ast.BinOp(arg1, binop(), arg2)

        unaryop = {"not_": ast.Not, "inv": ast.Invert, "invert": ast.Invert}.get(
            fn.attr
        )
        if unaryop is not None:
            arg = node.args[0]
            assert len(node.args) == 1
            return ast.UnaryOp(unaryop(), arg)

        compareop = {
            "lt": ast.Lt,
            "le": ast.LtE,
            "eq": ast.Eq,
            "ne": ast.NotEq,
            "gt": ast.Gt,
            "ge": ast.GtE,
        }.get(fn.attr)
        if compareop is not None:
            arg1, arg2 = node.args
            assert len(node.args) == 2
            return ast.Compare(arg1, [compareop()], [arg2])

        isop = {
            "is_": (ast.Is, ast.Eq),
            "is_not": (ast.IsNot, ast.NotEq),
        }.get(fn.attr)
        if isop is not None:
            isoper, eqoper = isop
            arg1, arg2 = node.args
            assert len(node.args) == 2
            oper = (
                eqoper if any(_needs_eq_operator(arg) for arg in node.args) else isoper
            )
            return ast.Compare(arg1, [oper()], [arg2])

        if fn.attr == "contains":
            arg1, arg2 = node.args
            assert len(node.args) == 2
            return ast.Compare(arg2, [ast.In()], [arg1])

        if fn.attr == "delitem":
            target, index = node.args
            assert len(node.args) == 2
            return ast.Delete(
                targets=[
                    ast.Subscript(value=target, slice=ast_index(index), ctx=ast.Del())
                ]
            )

        if fn.attr == "getitem":
            target, index = node.args
            assert len(node.args) == 2
            return ast.Subscript(value=target, slice=ast_index(index), ctx=ast.Load())

    return node


class PythonASTOptimizer(ast.NodeTransformer):
    __slots__ = ("_global_ctx",)

    def __init__(self):
        self._global_ctx: Deque[Set[str]] = deque([set()])

    @contextmanager
    def _new_global_context(self):
        """Context manager which sets a new Python `global` context."""
        self._global_ctx.append(set())
        try:
            yield
        finally:
            self._global_ctx.pop()

    @property
    def _global_context(self) -> Set[str]:
        """Return the current Python `global` context."""
        return self._global_ctx[-1]

    def visit_Call(self, node: ast.Call) -> ast.AST:
        """Eliminate most calls to Python's `operator` module in favor of using native
        operators."""
        new_node = self.generic_visit(node)
        if isinstance(new_node, ast.Call):
            return ast.copy_location(
                _optimize_operator_call(node.func, new_node), new_node
            )
        return new_node

    def visit_ExceptHandler(self, node: ast.ExceptHandler) -> Optional[ast.AST]:
        """Eliminate dead code from except handler bodies."""
        new_node = self.generic_visit(node)
        assert isinstance(new_node, ast.ExceptHandler)
        return ast.copy_location(
            ast.ExceptHandler(
                type=new_node.type,
                name=new_node.name,
                body=_filter_dead_code(new_node.body),
            ),
            new_node,
        )

    def visit_Expr(self, node: ast.Expr) -> Optional[ast.Expr]:
        """Eliminate no-op constant expressions which are in the tree
        as standalone statements."""
        if isinstance(node.value, (ast.Constant, ast.Name)):
            return None
        return node

    def visit_FunctionDef(self, node: ast.FunctionDef) -> Optional[ast.AST]:
        """Eliminate dead code from function bodies."""
        with self._new_global_context():
            new_node = self.generic_visit(node)
        assert isinstance(new_node, ast.FunctionDef)
        return ast.copy_location(
            ast_FunctionDef(
                name=new_node.name,
                args=new_node.args,
                body=_filter_dead_code(new_node.body),
                decorator_list=new_node.decorator_list,
                returns=new_node.returns,
            ),
            new_node,
        )

    def visit_Global(self, node: ast.Global) -> Optional[ast.Global]:
        """Eliminate redundant name declarations inside a Python `global` statement.

        Python `global` statements may only refer to a name prior to its declaration.
        Global contexts track names in prior `global` declarations and eliminate
        redundant names in `global` declarations. If all of the names in the current
        `global` statement are redundant, the entire node will be omitted."""
        new_names = set(node.names) - self._global_context
        self._global_context.update(new_names)
        return (
            ast.copy_location(ast.Global(names=list(new_names)), node)
            if new_names
            else None
        )

    def visit_If(self, node: ast.If) -> Optional[ast.AST]:
        """Eliminate dead code from if/elif bodies.

        If the new `if` statement `body` is empty after eliminating dead code, replace
        the body with the `orelse` body and negate the `if` condition.

        If both the `body` and `orelse` body are empty, eliminate the node from the
        tree."""
        new_node = self.generic_visit(node)
        assert isinstance(new_node, ast.If)

        new_body = _filter_dead_code(new_node.body)
        new_orelse = _filter_dead_code(new_node.orelse)

        if new_body:
            ifstmt = ast.If(
                test=new_node.test,
                body=new_body,
                orelse=new_orelse,
            )
        elif new_orelse:
            ifstmt = ast.If(
                test=ast.UnaryOp(op=ast.Not(), operand=new_node.test),
                body=new_orelse,
                orelse=[],
            )
        else:
            return None

        return ast.copy_location(ifstmt, new_node)

    def visit_While(self, node: ast.While) -> Optional[ast.AST]:
        """Eliminate dead code from while bodies."""
        new_node = self.generic_visit(node)
        assert isinstance(new_node, ast.While)
        return ast.copy_location(
            ast.While(
                test=new_node.test,
                body=_filter_dead_code(new_node.body),
                orelse=_filter_dead_code(new_node.orelse),
            ),
            new_node,
        )

    def visit_Try(self, node: ast.Try) -> Optional[ast.AST]:
        """Eliminate dead code from except try bodies."""
        new_node = self.generic_visit(node)
        assert isinstance(new_node, ast.Try)
        return ast.copy_location(
            ast.Try(
                body=_filter_dead_code(new_node.body),
                handlers=new_node.handlers,
                orelse=_filter_dead_code(new_node.orelse),
                finalbody=_filter_dead_code(new_node.finalbody),
            ),
            new_node,
        )

1	import ast	1✔
2	import functools	1✔
3	from collections import deque	1✔
4	from contextlib import contextmanager	1✔
5	from typing import Deque, Iterable, List, Optional, Set	1✔
6
7	from basilisp.lang.compiler.constants import OPERATOR_ALIAS	1✔
8	from basilisp.lang.compiler.utils import ast_FunctionDef, ast_index	1✔
9
10
11	def _filter_dead_code(nodes: Iterable[ast.stmt]) -> List[ast.stmt]:	1✔
12	"""Return a list of body nodes, trimming out unreachable code (any
13	statements appearing after `break`, `continue`, `raise`, and `return`
14	nodes)."""
15	new_nodes: List[ast.stmt] = []	1✔
16	for node in nodes:	1✔
17	if isinstance(node, (ast.Break, ast.Continue, ast.Raise, ast.Return)):	1✔
18	new_nodes.append(node)	1✔
19	break	1✔
20	new_nodes.append(node)	1✔
21	return new_nodes	1✔
22
23
24	def _needs_eq_operator(arg: ast.expr) -> bool:	1✔
25	return isinstance(arg, ast.Constant) and all(	1✔
26	arg.value is not v for v in (True, False, None, ...)
27	)
28
29
30	@functools.singledispatch	1✔
31	def _optimize_operator_call( # pylint: disable=unused-argument	1✔
32	fn: ast.AST, node: ast.Call
33	) -> ast.AST:
34	return node	1✔
35
36
37	@_optimize_operator_call.register(ast.Attribute)	1✔
38	def _optimize_operator_call_attr( # pylint: disable=too-many-return-statements	1✔
39	fn: ast.Attribute, node: ast.Call
40	) -> ast.AST:
41	"""Optimize calls to the Python `operator` module down to use the raw Python
42	operators.
43
44	Using Python operators directly will allow for more direct bytecode to be
45	emitted by the Python compiler and take advantage of any additional performance
46	improvements in future versions of Python."""
47	if isinstance(fn.value, ast.Name) and fn.value.id == OPERATOR_ALIAS:	1✔
48	binop = {	1✔
49	"add": ast.Add,
50	"and_": ast.BitAnd,
51	"floordiv": ast.FloorDiv,
52	"lshift": ast.LShift,
53	"mod": ast.Mod,
54	"mul": ast.Mult,
55	"matmul": ast.MatMult,
56	"or_": ast.BitOr,
57	"pow": ast.Pow,
58	"rshift": ast.RShift,
59	"sub": ast.Sub,
60	"truediv": ast.Div,
61	"xor": ast.BitXor,
62	}.get(fn.attr)
63	if binop is not None:	1✔
64	arg1, arg2 = node.args	1✔
65	assert len(node.args) == 2	1✔
66	return ast.BinOp(arg1, binop(), arg2)	1✔
67
68	unaryop = {"not_": ast.Not, "inv": ast.Invert, "invert": ast.Invert}.get(	1✔
69	fn.attr
70	)
71	if unaryop is not None:	1✔
72	arg = node.args[0]	1✔
73	assert len(node.args) == 1	1✔
74	return ast.UnaryOp(unaryop(), arg)	1✔
75
76	compareop = {	1✔
77	"lt": ast.Lt,
78	"le": ast.LtE,
79	"eq": ast.Eq,
80	"ne": ast.NotEq,
81	"gt": ast.Gt,
82	"ge": ast.GtE,
83	}.get(fn.attr)
84	if compareop is not None:	1✔
85	arg1, arg2 = node.args	1✔
86	assert len(node.args) == 2	1✔
87	return ast.Compare(arg1, [compareop()], [arg2])	1✔
88
89	isop = {	1✔
90	"is_": (ast.Is, ast.Eq),
91	"is_not": (ast.IsNot, ast.NotEq),
92	}.get(fn.attr)
93	if isop is not None:	1✔
94	isoper, eqoper = isop	1✔
95	arg1, arg2 = node.args	1✔
96	assert len(node.args) == 2	1✔
97	oper = (	1✔
98	eqoper if any(_needs_eq_operator(arg) for arg in node.args) else isoper
99	)
100	return ast.Compare(arg1, [oper()], [arg2])	1✔
101
102	if fn.attr == "contains":	1✔
103	arg1, arg2 = node.args	1✔
104	assert len(node.args) == 2	1✔
105	return ast.Compare(arg2, [ast.In()], [arg1])	1✔
106
107	if fn.attr == "delitem":	1✔
108	target, index = node.args	×
109	assert len(node.args) == 2	×
110	return ast.Delete(	×
111	targets=[
112	ast.Subscript(value=target, slice=ast_index(index), ctx=ast.Del())
113	]
114	)
115
116	if fn.attr == "getitem":	1✔
117	target, index = node.args	1✔
118	assert len(node.args) == 2	1✔
119	return ast.Subscript(value=target, slice=ast_index(index), ctx=ast.Load())	1✔
120
121	return node	1✔
122
123
124	class PythonASTOptimizer(ast.NodeTransformer):	1✔
125	__slots__ = ("_global_ctx",)	1✔
126
127	def __init__(self):	1✔
128	self._global_ctx: Deque[Set[str]] = deque([set()])	1✔
129
130	@contextmanager	1✔
131	def _new_global_context(self):	1✔
132	"""Context manager which sets a new Python `global` context."""
133	self._global_ctx.append(set())	1✔
134	try:	1✔
135	yield	1✔
136	finally:
137	self._global_ctx.pop()	1✔
138
139	@property	1✔
140	def _global_context(self) -> Set[str]:	1✔
141	"""Return the current Python `global` context."""
142	return self._global_ctx[-1]	1✔
143
144	def visit_Call(self, node: ast.Call) -> ast.AST:	1✔
145	"""Eliminate most calls to Python's `operator` module in favor of using native
146	operators."""
147	new_node = self.generic_visit(node)	1✔
148	if isinstance(new_node, ast.Call):	1✔
149	return ast.copy_location(	1✔
150	_optimize_operator_call(node.func, new_node), new_node
151	)
152	return new_node	×
153
154	def visit_ExceptHandler(self, node: ast.ExceptHandler) -> Optional[ast.AST]:	1✔
155	"""Eliminate dead code from except handler bodies."""
156	new_node = self.generic_visit(node)	1✔
157	assert isinstance(new_node, ast.ExceptHandler)	1✔
158	return ast.copy_location(	1✔
159	ast.ExceptHandler(
160	type=new_node.type,
161	name=new_node.name,
162	body=_filter_dead_code(new_node.body),
163	),
164	new_node,
165	)
166
167	def visit_Expr(self, node: ast.Expr) -> Optional[ast.Expr]:	1✔
168	"""Eliminate no-op constant expressions which are in the tree
169	as standalone statements."""
170	if isinstance(node.value, (ast.Constant, ast.Name)):	1✔
171	return None	1✔
172	return node	1✔
173
174	def visit_FunctionDef(self, node: ast.FunctionDef) -> Optional[ast.AST]:	1✔
175	"""Eliminate dead code from function bodies."""
176	with self._new_global_context():	1✔
177	new_node = self.generic_visit(node)	1✔
178	assert isinstance(new_node, ast.FunctionDef)	1✔
179	return ast.copy_location(	1✔
180	ast_FunctionDef(
181	name=new_node.name,
182	args=new_node.args,
183	body=_filter_dead_code(new_node.body),
184	decorator_list=new_node.decorator_list,
185	returns=new_node.returns,
186	),
187	new_node,
188	)
189
190	def visit_Global(self, node: ast.Global) -> Optional[ast.Global]:	1✔
191	"""Eliminate redundant name declarations inside a Python `global` statement.
192
193	Python `global` statements may only refer to a name prior to its declaration.
194	Global contexts track names in prior `global` declarations and eliminate
195	redundant names in `global` declarations. If all of the names in the current
196	`global` statement are redundant, the entire node will be omitted."""
197	new_names = set(node.names) - self._global_context	1✔
198	self._global_context.update(new_names)	1✔
199	return (	1✔
200	ast.copy_location(ast.Global(names=list(new_names)), node)
201	if new_names
202	else None
203	)
204
205	def visit_If(self, node: ast.If) -> Optional[ast.AST]:	1✔
206	"""Eliminate dead code from if/elif bodies.
207
208	If the new `if` statement `body` is empty after eliminating dead code, replace
209	the body with the `orelse` body and negate the `if` condition.
210
211	If both the `body` and `orelse` body are empty, eliminate the node from the
212	tree."""
213	new_node = self.generic_visit(node)	1✔
214	assert isinstance(new_node, ast.If)	1✔
215
216	new_body = _filter_dead_code(new_node.body)	1✔
217	new_orelse = _filter_dead_code(new_node.orelse)	1✔
218
219	if new_body:	1✔
220	ifstmt = ast.If(	1✔
221	test=new_node.test,
222	body=new_body,
223	orelse=new_orelse,
224	)
225	elif new_orelse:	1✔
226	ifstmt = ast.If(	1✔
227	test=ast.UnaryOp(op=ast.Not(), operand=new_node.test),
228	body=new_orelse,
229	orelse=[],
230	)
231	else:
232	return None	×
233
234	return ast.copy_location(ifstmt, new_node)	1✔
235
236	def visit_While(self, node: ast.While) -> Optional[ast.AST]:	1✔
237	"""Eliminate dead code from while bodies."""
238	new_node = self.generic_visit(node)	1✔
239	assert isinstance(new_node, ast.While)	1✔
240	return ast.copy_location(	1✔
241	ast.While(
242	test=new_node.test,
243	body=_filter_dead_code(new_node.body),
244	orelse=_filter_dead_code(new_node.orelse),
245	),
246	new_node,
247	)
248
249	def visit_Try(self, node: ast.Try) -> Optional[ast.AST]:	1✔
250	"""Eliminate dead code from except try bodies."""
251	new_node = self.generic_visit(node)	1✔
252	assert isinstance(new_node, ast.Try)	1✔
253	return ast.copy_location(	1✔
254	ast.Try(
255	body=_filter_dead_code(new_node.body),
256	handlers=new_node.handlers,
257	orelse=_filter_dead_code(new_node.orelse),
258	finalbody=_filter_dead_code(new_node.finalbody),
259	),
260	new_node,
261	)

basilisp-lang / basilisp / 10943311926

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous