12468686896

Committed 23 Dec 2024 02:53PM CUT coverage: 98.68% (-0.06%) from 98.738%

Build # 12468686896

Build Type

Pull #1176

github

Committed by

web-flow

Commit Message

Merge 454fa0188 into 61e86f30e

Pull Request Pull Request #1176: Switch Python import machinery to focus primarily on Python modules

Run Details

1031 of 1040 branches covered (99.13%)

Branch coverage included in aggregate %.

60 of 61 new or added lines in 5 files covered. (98.36%)

5 existing lines in 1 file now uncovered.

8837 of 8960 relevant lines covered (98.63%)

0.99 hits per line

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

96.06

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

import ast
import functools
from collections import deque
from collections.abc import Iterable
from contextlib import contextmanager
from typing import Deque, Optional

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


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=index, ctx=ast.Del())]
            )

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

basilisp-lang / basilisp / 12468686896

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