optimizer.py 1.6 KB

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