refactor: reimplement AST folding

vyper/ast/folding.py

@@ -1,10 +1,9 @@
-from typing import Optional, Union
+from typing import Union
 
 from vyper.ast import nodes as vy_ast
 from vyper.builtins.functions import DISPATCH_TABLE
-from vyper.exceptions import UnfoldableNode, UnknownType
+from vyper.exceptions import UnfoldableNode
 from vyper.semantics.types.base import VyperType
-from vyper.semantics.types.utils import type_from_annotation
 
 
 def fold(vyper_module: vy_ast.Module) -> None:
@@ -49,6 +48,12 @@ def replace_literal_ops(vyper_module: vy_ast.Module) -> int:
         except UnfoldableNode:
             continue
 
+        # type may not be available if it is within a type's annotation
+        # e.g. DynArray[uint256, 2 ** 8]
+        typ = node._metadata.get("type")
+        if typ:
+            new_node._metadata["type"] = node._metadata["type"]
+
         changed_nodes += 1
         vyper_module.replace_in_tree(node, new_node)
 
@@ -111,6 +116,7 @@ def replace_builtin_functions(vyper_module: vy_ast.Module) -> int:
             continue
         try:
             new_node = func.evaluate(node)  # type: ignore
+            new_node._metadata["type"] = node._metadata["type"]
         except UnfoldableNode:
             continue
 
@@ -146,37 +152,26 @@ def replace_user_defined_constants(vyper_module: vy_ast.Module) -> int:
             continue
 
         # Extract type definition from propagated annotation
-        type_ = None
-        try:
-            type_ = type_from_annotation(node.annotation)
-        except UnknownType:
-            # handle user-defined types e.g. structs - it's OK to not
-            # propagate the type annotation here because user-defined
-            # types can be unambiguously inferred at typechecking time
-            pass
-
-        changed_nodes += replace_constant(
-            vyper_module, node.target.id, node.value, False, type_=type_
-        )
+        type_ = node._metadata["type"]
+
+        changed_nodes += replace_constant(vyper_module, node.target.id, node.value, type_, False)
 
     return changed_nodes
 
 
 # TODO constant folding on log events
 
 
-def _replace(old_node, new_node, type_=None):
+def _replace(old_node, new_node, type_):
     if isinstance(new_node, vy_ast.Constant):
         new_node = new_node.from_node(old_node, value=new_node.value)
-        if type_:
-            new_node._metadata["type"] = type_
+        new_node._metadata["type"] = type_
         return new_node
     elif isinstance(new_node, vy_ast.List):
         base_type = type_.value_type if type_ else None
         list_values = [_replace(old_node, i, type_=base_type) for i in new_node.elements]
         new_node = new_node.from_node(old_node, elements=list_values)
-        if type_:
-            new_node._metadata["type"] = type_
+        new_node._metadata["type"] = type_
         return new_node
     elif isinstance(new_node, vy_ast.Call):
         # Replace `Name` node with `Call` node
@@ -188,6 +183,7 @@ def _replace(old_node, new_node, type_=None):
         new_node = new_node.from_node(
             old_node, func=new_node.func, args=new_node.args, keyword=keyword, keywords=keywords
         )
+        new_node._metadata["type"] = type_
         return new_node
     else:
         raise UnfoldableNode
@@ -197,8 +193,8 @@ def replace_constant(
     vyper_module: vy_ast.Module,
     id_: str,
     replacement_node: Union[vy_ast.Constant, vy_ast.List, vy_ast.Call],
+    type_: VyperType,
     raise_on_error: bool,
-    type_: Optional[VyperType] = None,
 ) -> int:
     """
     Replace references to a variable name with a literal value.

vyper/ast/nodes.py

@@ -912,6 +912,14 @@
 class UnaryOp(ExprNode):
     __slots__ = ("op", "operand")
 
+    def prefold(self) -> ExprNode:
+        operand = self.operand._metadata.get("folded_value")
+        if operand is None:
+            return
+
+        value = self.op._op(operand.value)
+        return type(self.operand).from_node(self, value=value)
+
     def evaluate(self) -> ExprNode:
         """
         Attempt to evaluate the unary operation.
@@ -960,6 +968,22 @@
 class BinOp(ExprNode):
     __slots__ = ("left", "op", "right")
 
+    def prefold(self) -> ExprNode:
+        left = self.left._metadata.get("folded_value")
+        right = self.right._metadata.get("folded_value")
+
+        if None in (left, right):
+            return
+
+        # this validation is performed to prevent the compiler from hanging
+        # on very large shifts and improve the error message for negative
+        # values.
+        if isinstance(self.op, (LShift, RShift)) and not (0 <= right.value <= 256):
+            raise InvalidLiteral("Shift bits must be between 0 and 256", right)
+
+        value = self.op._op(left.value, right.value)
+        return type(left).from_node(self, value=value)
+
     def evaluate(self) -> ExprNode:
         """
         Attempt to evaluate the arithmetic operation.
@@ -1110,6 +1134,14 @@
 class BoolOp(ExprNode):
     __slots__ = ("op", "values")
 
+    def prefold(self) -> ExprNode:
+        values = [i._metadata.get("folded_value") for i in self.values]
+        if None in values:
+            return
+
+        value = self.op._op(values)
+        return NameConstant.from_node(self, value=value)
+
     def evaluate(self) -> ExprNode:
         """
         Attempt to evaluate the boolean operation.
@@ -1166,6 +1198,16 @@
         kwargs["right"] = kwargs.pop("comparators")[0]
         super().__init__(*args, **kwargs)
 
+    def prefold(self) -> ExprNode:
+        left = self.left._metadata.get("folded_value")
+        right = self.right._metadata.get("folded_value")
+
+        if None in (left, right):
+            return
+
+        value = self.op._op(left.value, right.value)
+        return NameConstant.from_node(self, value=value)
+
     def evaluate(self) -> ExprNode:
         """
         Attempt to evaluate the comparison.

vyper/builtins/_signatures.py

@@ -1,12 +1,16 @@
 import functools
 from typing import Dict
 
-from vyper.ast import nodes as vy_ast
+
 from vyper.ast.validation import validate_call_args
 from vyper.codegen.expr import Expr
 from vyper.codegen.ir_node import IRnode
 from vyper.exceptions import CompilerPanic, TypeMismatch
-from vyper.semantics.analysis.utils import get_exact_type_from_node, validate_expected_type
+from vyper.semantics.analysis.utils import (
+    check_kwargable,
+    get_exact_type_from_node,
+    validate_expected_type,
+)
 from vyper.semantics.types import TYPE_T, KwargSettings, VyperType
 from vyper.semantics.types.utils import type_from_annotation
 
@@ -97,14 +101,14 @@
             # note special meaning for -1 in validate_call_args API
             expect_num_args = (num_args, -1)
 
-        validate_call_args(node, expect_num_args, self._kwargs)
+        validate_call_args(node, expect_num_args, self._kwargs.keys())
 
         for arg, (_, expected) in zip(node.args, self._inputs):
             self._validate_single(arg, expected)
 
         for kwarg in node.keywords:
             kwarg_settings = self._kwargs[kwarg.arg]
-            if kwarg_settings.require_literal and not isinstance(kwarg.value, vy_ast.Constant):
+            if kwarg_settings.require_literal and not check_kwargable(kwarg.value):
                 raise TypeMismatch("Value for kwarg must be a literal", kwarg.value)
             self._validate_single(kwarg.value, kwarg_settings.typ)
 

vyper/builtins/functions.py

@@ -111,6 +111,8 @@ class FoldedFunction(BuiltinFunction):
     # Since foldable builtin functions are not folded before semantics validation,
     # this flag is used for `check_kwargable` in semantics validation.
     _kwargable = True
+    # Skip annotation of builtins if it will be folded before codegen
+    _is_folded_before_codegen = True
 
 
 class TypenameFoldedFunction(FoldedFunction):
@@ -2260,6 +2262,9 @@ def build_IR(self, expr, args, kwargs, context):
 
 class Empty(TypenameFoldedFunction):
     _id = "empty"
+    # Since `empty` is not folded in the AST, `is_folded` is set to False
+    # so that it will be properly annotated.
+    _is_folded_before_codegen = False
 
     def fetch_call_return(self, node):
         type_ = self.infer_arg_types(node)[0].typedef

vyper/codegen/function_definitions/external_function.py

@@ -1,3 +1,4 @@
+from vyper import ast as vy_ast
 from vyper.codegen.abi_encoder import abi_encoding_matches_vyper
 from vyper.codegen.context import Context, VariableRecord
 from vyper.codegen.core import get_element_ptr, getpos, make_setter, needs_clamp
@@ -50,7 +51,7 @@ def _register_function_args(func_t: ContractFunctionT, context: Context) -> list
 
 
 def _generate_kwarg_handlers(
-    func_t: ContractFunctionT, context: Context
+    func_t: ContractFunctionT, context: Context, code: vy_ast.FunctionDef
 ) -> dict[str, tuple[int, IRnode]]:
     # generate kwarg handlers.
     # since they might come in thru calldata or be default,
@@ -62,7 +63,7 @@ def _generate_kwarg_handlers(
     #    write default args to memory
     #    goto external_function_common_ir
 
-    def handler_for(calldata_kwargs, default_kwargs):
+    def handler_for(calldata_kwargs, folded_default_kwargs, original_default_kwargs):
         calldata_args = func_t.positional_args + calldata_kwargs
         # create a fake type so that get_element_ptr works
         calldata_args_t = TupleT(list(arg.typ for arg in calldata_args))
@@ -81,7 +82,7 @@ def handler_for(calldata_kwargs, default_kwargs):
         calldata_min_size = args_abi_t.min_size() + 4
 
         # TODO optimize make_setter by using
-        # TupleT(list(arg.typ for arg in calldata_kwargs + default_kwargs))
+        # TupleT(list(arg.typ for arg in calldata_kwargs + folded_default_kwargs))
         # (must ensure memory area is contiguous)
 
         for i, arg_meta in enumerate(calldata_kwargs):
@@ -97,15 +98,15 @@ def handler_for(calldata_kwargs, default_kwargs):
             copy_arg.source_pos = getpos(arg_meta.ast_source)
             ret.append(copy_arg)
 
-        for x in default_kwargs:
+        for x, y in zip(original_default_kwargs, folded_default_kwargs):
             dst = context.lookup_var(x.name).pos
             lhs = IRnode(dst, location=MEMORY, typ=x.typ)
-            lhs.source_pos = getpos(x.ast_source)
-            kw_ast_val = func_t.default_values[x.name]  # e.g. `3` in x: int = 3
+            lhs.source_pos = getpos(y)
+            kw_ast_val = y
             rhs = Expr(kw_ast_val, context).ir_node
 
             copy_arg = make_setter(lhs, rhs)
-            copy_arg.source_pos = getpos(x.ast_source)
+            copy_arg.source_pos = getpos(y)
             ret.append(copy_arg)
 
         ret.append(["goto", func_t._ir_info.external_function_base_entry_label])
@@ -116,19 +117,25 @@ def handler_for(calldata_kwargs, default_kwargs):
     ret = {}
 
     keyword_args = func_t.keyword_args
+    folded_keyword_args = code.args.defaults
 
     # allocate variable slots in memory
     for arg in keyword_args:
         context.new_variable(arg.name, arg.typ, is_mutable=False)
 
     for i, _ in enumerate(keyword_args):
         calldata_kwargs = keyword_args[:i]
-        default_kwargs = keyword_args[i:]
+        # folded ast
+        original_default_kwargs = keyword_args[i:]
+        # unfolded ast
+        folded_default_kwargs = folded_keyword_args[1:]
 
-        sig, calldata_min_size, ir_node = handler_for(calldata_kwargs, default_kwargs)
+        sig, calldata_min_size, ir_node = handler_for(
+            calldata_kwargs, folded_default_kwargs, original_default_kwargs
+        )
         ret[sig] = calldata_min_size, ir_node
 
-    sig, calldata_min_size, ir_node = handler_for(keyword_args, [])
+    sig, calldata_min_size, ir_node = handler_for(keyword_args, [], [])
 
     ret[sig] = calldata_min_size, ir_node
 
@@ -153,7 +160,7 @@ def generate_ir_for_external_function(code, func_t, context):
     handle_base_args = _register_function_args(func_t, context)
 
     # generate handlers for kwargs and register the variable records
-    kwarg_handlers = _generate_kwarg_handlers(func_t, context)
+    kwarg_handlers = _generate_kwarg_handlers(func_t, context, code)
 
     body = ["seq"]
     # once optional args have been handled,

vyper/compiler/phases.py

@@ -271,13 +271,15 @@ def generate_folded_ast(
 
     vy_ast.validation.validate_literal_nodes(vyper_module)
 
+    with input_bundle.search_path(contract_path.parent):
+        validate_semantics(vyper_module, input_bundle)
+
+    symbol_tables = set_data_positions(vyper_module, storage_layout_overrides)
+
     vyper_module_folded = copy.deepcopy(vyper_module)
     vy_ast.folding.fold(vyper_module_folded)
 
-    with input_bundle.search_path(contract_path.parent):
-        validate_semantics(vyper_module_folded, input_bundle)
-
-    symbol_tables = set_data_positions(vyper_module_folded, storage_layout_overrides)
+    vy_ast.expansion.remove_unused_statements(vyper_module_folded)
 
     return vyper_module_folded, symbol_tables
 

vyper/semantics/analysis/__init__.py

@@ -4,6 +4,7 @@
 from ..namespace import get_namespace
 from .local import validate_functions
 from .module import add_module_namespace
+from .pre_typecheck import pre_typecheck
 from .utils import _ExprAnalyser
 
 
@@ -12,6 +13,7 @@ def validate_semantics(vyper_ast, input_bundle):
     namespace = get_namespace()
 
     with namespace.enter_scope():
+        pre_typecheck(vyper_ast)
         add_module_namespace(vyper_ast, input_bundle)
-        vy_ast.expansion.expand_annotated_ast(vyper_ast)
+        vy_ast.expansion.generate_public_variable_getters(vyper_ast)
         validate_functions(vyper_ast)