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"""
10from . import nodes
11from .visitor import NodeTransformer
12
13
14def optimize(node, environment):
15    """The context hint can be used to perform an static optimization
16    based on the context given."""
17    optimizer = Optimizer(environment)
18    return optimizer.visit(node)
19
20
21class Optimizer(NodeTransformer):
22    def __init__(self, environment):
23        self.environment = environment
24
25    def generic_visit(self, node, *args, **kwargs):
26        node = super().generic_visit(node, *args, **kwargs)
27
28        # Do constant folding. Some other nodes besides Expr have
29        # as_const, but folding them causes errors later on.
30        if isinstance(node, nodes.Expr):
31            try:
32                return nodes.Const.from_untrusted(
33                    node.as_const(args[0] if args else None),
34                    lineno=node.lineno,
35                    environment=self.environment,
36                )
37            except nodes.Impossible:
38                pass
39
40        return node
41