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