1"""The optimizer tries to constant fold expressions and modify the AST 2in place so that it should be faster to evaluate. 3 4Because the AST does not contain all the scoping information and the 5compiler has to find that out, we cannot do all the optimizations we 6want. For example, loop unrolling doesn't work because unrolled loops 7would have a different scope. The solution would be a second syntax tree 8that stored the scoping rules. 9""" 10import typing as t 11 12from . import nodes 13from .visitor import NodeTransformer 14 15if t.TYPE_CHECKING: 16 from .environment import Environment 17 18 19def optimize(node: nodes.Node, environment: "Environment") -> nodes.Node: 20 """The context hint can be used to perform an static optimization 21 based on the context given.""" 22 optimizer = Optimizer(environment) 23 return t.cast(nodes.Node, optimizer.visit(node)) 24 25 26class Optimizer(NodeTransformer): 27 def __init__(self, environment: "t.Optional[Environment]") -> None: 28 self.environment = environment 29 30 def generic_visit( 31 self, node: nodes.Node, *args: t.Any, **kwargs: t.Any 32 ) -> nodes.Node: 33 node = super().generic_visit(node, *args, **kwargs) 34 35 # Do constant folding. Some other nodes besides Expr have 36 # as_const, but folding them causes errors later on. 37 if isinstance(node, nodes.Expr): 38 try: 39 return nodes.Const.from_untrusted( 40 node.as_const(args[0] if args else None), 41 lineno=node.lineno, 42 environment=self.environment, 43 ) 44 except nodes.Impossible: 45 pass 46 47 return node 48