add _is_prefoldable attribute to nodes

tserg · Dec 22, 2023 · 59d2a04 · 59d2a04
1 parent 83a8b66
commit 59d2a04
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 12 deletions.
diff --git 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
@@ -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