diff --git a/vyper/ast/nodes.py b/vyper/ast/nodes.py index dbc2e36014..6ab1bb96e1 100644 --- a/vyper/ast/nodes.py +++ b/vyper/ast/nodes.py @@ -208,11 +208,13 @@ class VyperNode: _description : str, optional A human-readable description of the node. Used to give more verbose error messages. + _is_prefoldable : str, optional + If `True`, indicates that pre-folding should be attempted on the node. _only_empty_fields : Tuple, optional Field names that, if present, must be set to None or a `SyntaxException` is raised. This attribute is used to exclude syntax that is valid in Python but not in Vyper. - _terminus : bool, optional + _is_terminus : bool, optional If `True`, indicates that execution halts upon reaching this node. _translated_fields : Dict, optional Field names that are reassigned if encountered. Used to normalize fields @@ -883,6 +885,7 @@ def s(self): class List(ExprNode): __slots__ = ("elements",) + _is_prefoldable = True _translated_fields = {"elts": "elements"} def prefold(self) -> Optional[ExprNode]: @@ -920,6 +923,7 @@ class Name(ExprNode): class UnaryOp(ExprNode): __slots__ = ("op", "operand") + _is_prefoldable = True def prefold(self) -> Optional[ExprNode]: operand = self.operand._metadata.get("folded_value") @@ -975,6 +979,7 @@ def _op(self, value): class BinOp(ExprNode): __slots__ = ("left", "op", "right") + _is_prefoldable = True def prefold(self) -> Optional[ExprNode]: left = self.left._metadata.get("folded_value") @@ -1134,6 +1139,7 @@ class RShift(Operator): class BoolOp(ExprNode): __slots__ = ("op", "values") + _is_prefoldable = True def prefold(self) -> Optional[ExprNode]: values = [i._metadata.get("folded_value") for i in self.values] @@ -1190,6 +1196,7 @@ class Compare(ExprNode): """ __slots__ = ("left", "op", "right") + _is_prefoldable = True def __init__(self, *args, **kwargs): if len(kwargs["ops"]) > 1 or len(kwargs["comparators"]) > 1: @@ -1310,6 +1317,7 @@ class Attribute(ExprNode): class Subscript(ExprNode): __slots__ = ("slice", "value") + _is_prefoldable = True def prefold(self) -> Optional[ExprNode]: slice_ = self.slice.value._metadata.get("folded_value") diff --git a/vyper/semantics/analysis/pre_typecheck.py b/vyper/semantics/analysis/pre_typecheck.py index 6564440c1e..797fe3bbd3 100644 --- a/vyper/semantics/analysis/pre_typecheck.py +++ b/vyper/semantics/analysis/pre_typecheck.py @@ -46,23 +46,15 @@ def pre_typecheck(node: vy_ast.Module) -> None: def prefold(node: vy_ast.VyperNode, constants: dict[str, vy_ast.VyperNode]) -> None: - if isinstance( - node, - ( - vy_ast.BinOp, - vy_ast.BoolOp, - vy_ast.Compare, - vy_ast.List, - vy_ast.Subscript, - vy_ast.UnaryOp, - ), - ): + if getattr(node, "_is_prefoldable", None): node._metadata["folded_value"] = node.prefold() + return if isinstance(node, vy_ast.Name): var_name = node.id if var_name in constants: node._metadata["folded_value"] = constants[var_name] + return if isinstance(node, vy_ast.Call): if isinstance(node.func, vy_ast.Name): @@ -73,3 +65,4 @@ def prefold(node: vy_ast.VyperNode, constants: dict[str, vy_ast.VyperNode]) -> N call_type = DISPATCH_TABLE.get(func_name) if call_type: node._metadata["folded_value"] = call_type.prefold(node) # type: ignore + return