Repository: vyperlang/vyper

Issue:
typechecker does not check constants
### Version Information

* vyper Version (output of `vyper --version`): v0.3.1

### What's your issue about?

During constant folding, the annotated types for constants do not get propagated. Instead they are translated into literals and the literal types are inferred at the use site. For example the following example compiles even though it shouldn't really:
```python
MY_CONSTANT: constant(uint256[3]) = [1,2,3]

@external
def foo() -> uint8[3]:
    return MY_CONSTANT
```

Compare the error generated for the following

```python
@external
def foo() -> uint8[3]:
    xs: uint256[3] = [1,2,3]
    return xs
```
```
Error compiling: tmp/foo.vy
vyper.exceptions.TypeMismatch: Given reference has type uint256[3], expected uint8[3]
  contract "tmp/foo.vy", function "foo", line 7:11 
       6     xs: uint256[3] = [1,2,3]
  ---> 7     return xs
  ------------------^
       8
```

### How can it be fixed?

Propagate type annotations during constant folding


Discussion:
I am interested in taking this up next. I will start end of this week just in case anyone else has already started work.

Resolving patch:
diff --git a/vyper/ast/folding.py b/vyper/ast/folding.py
--- a/vyper/ast/folding.py
+++ b/vyper/ast/folding.py
@@ -4,6 +4,8 @@
 from vyper.ast import nodes as vy_ast
 from vyper.builtin_functions import DISPATCH_TABLE
 from vyper.exceptions import UnfoldableNode
+from vyper.semantics.types.bases import BaseTypeDefinition, DataLocation
+from vyper.semantics.types.utils import get_type_from_annotation
 
 BUILTIN_CONSTANTS = {
     "EMPTY_BYTES32": (
@@ -172,7 +174,17 @@ def replace_user_defined_constants(vyper_module: vy_ast.Module) -> int:
             # annotation is not wrapped in `constant(...)`
             continue
 
-        changed_nodes += replace_constant(vyper_module, node.target.id, node.value, False)
+        # Extract type definition from propagated annotation
+        constant_annotation = node.get("annotation.args")[0]
+        type_ = (
+            get_type_from_annotation(constant_annotation, DataLocation.UNSET)
+            if constant_annotation
+            else None
+        )
+
+        changed_nodes += replace_constant(
+            vyper_module, node.target.id, node.value, False, type_=type_
+        )
 
     return changed_nodes
 
@@ -180,12 +192,19 @@ def replace_user_defined_constants(vyper_module: vy_ast.Module) -> int:
 # TODO constant folding on log events
 
 
-def _replace(old_node, new_node):
+def _replace(old_node, new_node, type_=None):
     if isinstance(new_node, vy_ast.Constant):
-        return new_node.from_node(old_node, value=new_node.value)
+        new_node = new_node.from_node(old_node, value=new_node.value)
+        if type_:
+            new_node._metadata["type"] = type_
+        return new_node
     elif isinstance(new_node, vy_ast.List):
-        list_values = [_replace(old_node, i) for i in new_node.elements]
-        return new_node.from_node(old_node, elements=list_values)
+        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_
+        return new_node
     else:
         raise UnfoldableNode
 
@@ -195,6 +214,7 @@ def replace_constant(
     id_: str,
     replacement_node: Union[vy_ast.Constant, vy_ast.List],
     raise_on_error: bool,
+    type_: BaseTypeDefinition = None,
 ) -> int:
     """
     Replace references to a variable name with a literal value.
@@ -209,6 +229,8 @@ def replace_constant(
         Vyper ast node representing the literal value to be substituted in.
     raise_on_error: bool
         Boolean indicating if `UnfoldableNode` exception should be raised or ignored.
+    type_ : BaseTypeDefinition, optional
+        Type definition to be propagated to type checker.
 
     Returns
     -------
@@ -236,7 +258,7 @@ def replace_constant(
                 continue
 
         try:
-            new_node = _replace(node, replacement_node)
+            new_node = _replace(node, replacement_node, type_=type_)
         except UnfoldableNode:
             if raise_on_error:
                 raise
diff --git a/vyper/semantics/validation/local.py b/vyper/semantics/validation/local.py
--- a/vyper/semantics/validation/local.py
+++ b/vyper/semantics/validation/local.py
@@ -33,7 +33,6 @@
 from vyper.semantics.types.value.address import AddressDefinition
 from vyper.semantics.types.value.array_value import StringDefinition
 from vyper.semantics.types.value.boolean import BoolDefinition
-from vyper.semantics.types.value.numeric import Uint256Definition
 from vyper.semantics.validation.annotation import StatementAnnotationVisitor
 from vyper.semantics.validation.base import VyperNodeVisitorBase
 from vyper.semantics.validation.utils import (
@@ -320,7 +319,7 @@ def visit_For(self, node):
                     raise StateAccessViolation("Value must be a literal", node)
                 if args[0].value <= 0:
                     raise StructureException("For loop must have at least 1 iteration", args[0])
-                validate_expected_type(args[0], Uint256Definition())
+                validate_expected_type(args[0], IntegerAbstractType())
                 type_list = get_possible_types_from_node(args[0])
             else:
                 validate_expected_type(args[0], IntegerAbstractType())
diff --git a/vyper/semantics/validation/utils.py b/vyper/semantics/validation/utils.py
--- a/vyper/semantics/validation/utils.py
+++ b/vyper/semantics/validation/utils.py
@@ -88,6 +88,8 @@ def get_exact_type_from_node(self, node, only_definitions=True):
     def get_possible_types_from_node(self, node, only_definitions=True):
         """
         Find all possible types for a given node.
+        If the node's metadata contains type information propagated from constant folding,
+        then that type is returned.
 
         Arguments
         ---------
@@ -102,6 +104,10 @@ def get_possible_types_from_node(self, node, only_definitions=True):
         List
             A list of type objects
         """
+        # Early termination if typedef is propagated in metadata
+        if "type" in node._metadata:
+            return [node._metadata["type"]]
+
         fn = self._find_fn(node)
         types_list = fn(node)
         if only_definitions:
@@ -113,6 +119,7 @@ def get_possible_types_from_node(self, node, only_definitions=True):
                     )
                 else:
                     raise InvalidReference("Expected a literal or variable", node)
+
         if all(isinstance(i, IntegerAbstractType) for i in types_list):
             # for numeric types, sort according by number of bits descending
             # we do this to ensure literals are cast with the largest possible type


Test patch:
diff --git a/tests/parser/types/numbers/test_constants.py b/tests/parser/types/numbers/test_constants.py
--- a/tests/parser/types/numbers/test_constants.py
+++ b/tests/parser/types/numbers/test_constants.py
@@ -1,6 +1,10 @@
+import itertools
 from decimal import Decimal
 
+import pytest
+
 from vyper.compiler import compile_code
+from vyper.exceptions import InvalidType
 
 
 def test_builtin_constants(get_contract_with_gas_estimation):
@@ -127,6 +131,22 @@ def test_add(a: uint256) -> uint256:
     assert c.test_add(7) == 40
 
 
+# Would be nice to put this somewhere accessible, like in vyper.types or something
+integer_types = ["uint8", "int128", "int256", "uint256"]
+
+
+@pytest.mark.parametrize("storage_type,return_type", itertools.permutations(integer_types, 2))
+def test_custom_constants_fail(get_contract, assert_compile_failed, storage_type, return_type):
+    code = f"""
+MY_CONSTANT: constant({storage_type}) = 1
+
+@external
+def foo() -> {return_type}:
+    return MY_CONSTANT
+    """
+    assert_compile_failed(lambda: get_contract(code), InvalidType)
+
+
 def test_constant_address(get_contract):
     code = """
 OWNER: constant(address) = 0x0000000000000000000000000000000000000012
diff --git a/tests/parser/types/test_dynamic_array.py b/tests/parser/types/test_dynamic_array.py
--- a/tests/parser/types/test_dynamic_array.py
+++ b/tests/parser/types/test_dynamic_array.py
@@ -1,6 +1,8 @@
+import itertools
+
 import pytest
 
-from vyper.exceptions import ArrayIndexException, OverflowException
+from vyper.exceptions import ArrayIndexException, InvalidType, OverflowException, TypeMismatch
 
 
 def test_list_tester_code(get_contract_with_gas_estimation):
@@ -481,3 +483,42 @@ def ix(i: uint256) -> decimal:
 
 
 # TODO test loops
+
+# Would be nice to put this somewhere accessible, like in vyper.types or something
+integer_types = ["uint8", "int128", "int256", "uint256"]
+
+
+@pytest.mark.parametrize("storage_type,return_type", itertools.permutations(integer_types, 2))
+def test_constant_list_fail(get_contract, assert_compile_failed, storage_type, return_type):
+    code = f"""
+MY_CONSTANT: constant(DynArray[{storage_type}, 3]) = [1, 2, 3]
+
+@external
+def foo() -> DynArray[{return_type}, 3]:
+    return MY_CONSTANT
+    """
+    assert_compile_failed(lambda: get_contract(code), InvalidType)
+
+
+@pytest.mark.parametrize("storage_type,return_type", itertools.permutations(integer_types, 2))
+def test_constant_list_fail_2(get_contract, assert_compile_failed, storage_type, return_type):
+    code = f"""
+MY_CONSTANT: constant(DynArray[{storage_type}, 3]) = [1, 2, 3]
+
+@external
+def foo() -> {return_type}:
+    return MY_CONSTANT[0]
+    """
+    assert_compile_failed(lambda: get_contract(code), InvalidType)
+
+
+@pytest.mark.parametrize("storage_type,return_type", itertools.permutations(integer_types, 2))
+def test_constant_list_fail_3(get_contract, assert_compile_failed, storage_type, return_type):
+    code = f"""
+MY_CONSTANT: constant(DynArray[{storage_type}, 3]) = [1, 2, 3]
+
+@external
+def foo(i: uint256) -> {return_type}:
+    return MY_CONSTANT[i]
+    """
+    assert_compile_failed(lambda: get_contract(code), TypeMismatch)
diff --git a/tests/parser/types/test_lists.py b/tests/parser/types/test_lists.py
--- a/tests/parser/types/test_lists.py
+++ b/tests/parser/types/test_lists.py
@@ -1,6 +1,8 @@
+import itertools
+
 import pytest
 
-from vyper.exceptions import ArrayIndexException, OverflowException
+from vyper.exceptions import ArrayIndexException, InvalidType, OverflowException, TypeMismatch
 
 
 def test_list_tester_code(get_contract_with_gas_estimation):
@@ -428,6 +430,46 @@ def ix(i: uint256) -> {type}:
     assert_tx_failed(lambda: c.ix(len(value) + 1))
 
 
+# Would be nice to put this somewhere accessible, like in vyper.types or something
+integer_types = ["uint8", "int128", "int256", "uint256"]
+
+
+@pytest.mark.parametrize("storage_type,return_type", itertools.permutations(integer_types, 2))
+def test_constant_list_fail(get_contract, assert_compile_failed, storage_type, return_type):
+    code = f"""
+MY_CONSTANT: constant({storage_type}[3]) = [1, 2, 3]
+
+@external
+def foo() -> {return_type}[3]:
+    return MY_CONSTANT
+    """
+    assert_compile_failed(lambda: get_contract(code), InvalidType)
+
+
+@pytest.mark.parametrize("storage_type,return_type", itertools.permutations(integer_types, 2))
+def test_constant_list_fail_2(get_contract, assert_compile_failed, storage_type, return_type):
+    code = f"""
+MY_CONSTANT: constant({storage_type}[3]) = [1, 2, 3]
+
+@external
+def foo() -> {return_type}:
+    return MY_CONSTANT[0]
+    """
+    assert_compile_failed(lambda: get_contract(code), InvalidType)
+
+
+@pytest.mark.parametrize("storage_type,return_type", itertools.permutations(integer_types, 2))
+def test_constant_list_fail_3(get_contract, assert_compile_failed, storage_type, return_type):
+    code = f"""
+MY_CONSTANT: constant({storage_type}[3]) = [1, 2, 3]
+
+@external
+def foo(i: uint256) -> {return_type}:
+    return MY_CONSTANT[i]
+    """
+    assert_compile_failed(lambda: get_contract(code), TypeMismatch)
+
+
 def test_constant_list_address(get_contract, assert_tx_failed):
     some_good_address = [
         "0x0000000000000000000000000000000000012345",
@@ -455,3 +497,63 @@ def ix(i: uint256) -> address:
         assert c.ix(i) == p
     # assert oob
     assert_tx_failed(lambda: c.ix(len(some_good_address) + 1))
+
+
+@pytest.mark.parametrize(
+    "type,value",
+    [
+        ("decimal", [[5.0, 11.0], [17.0, 29.0], [37.0, 41.0]]),
+        ("uint8", [[0, 1], [17, 250], [255, 2]]),
+        ("int128", [[0, -1], [1, -(2 ** 127)], [2 ** 127 - 1, -50]]),
+        ("int256", [[0, -1], [1, -(2 ** 255)], [2 ** 255 - 1, -50]]),
+        ("uint256", [[0, 1], [2 ** 8, 2 ** 255 + 1], [2 ** 256 - 1, 100]]),
+        (
+            "uint256",
+            [
+                [2 ** 255 + 1, 2 ** 255 + 2],
+                [2 ** 255 + 3, 2 ** 255 + 4],
+                [2 ** 255 + 5, 2 ** 255 + 6],
+            ],
+        ),
+        ("bool", [[True, False], [True, False], [True, False]]),
+    ],
+)
+def test_constant_nested_list(get_contract, assert_tx_failed, type, value):
+    code = f"""
+MY_LIST: constant({type}[{len(value[0])}][{len(value)}]) = {value}
+@external
+def ix(i: uint256, j: uint256) -> {type}:
+    return MY_LIST[i][j]
+    """
+    c = get_contract(code)
+    for i, p in enumerate(value):
+        for j, q in enumerate(p):
+            assert c.ix(i, j) == q
+    # assert oob
+    assert_tx_failed(lambda: c.ix(len(value) + 1, len(value[0]) + 1))
+
+
+@pytest.mark.parametrize("storage_type,return_type", itertools.permutations(integer_types, 2))
+def test_constant_nested_list_fail(get_contract, assert_compile_failed, storage_type, return_type):
+    code = f"""
+MY_CONSTANT: constant({storage_type}[2][3]) = [[1, 2], [3, 4], [5, 6]]
+
+@external
+def foo() -> {return_type}[2][3]:
+    return MY_CONSTANT
+    """
+    assert_compile_failed(lambda: get_contract(code), InvalidType)
+
+
+@pytest.mark.parametrize("storage_type,return_type", itertools.permutations(integer_types, 2))
+def test_constant_nested_list_fail_2(
+    get_contract, assert_compile_failed, storage_type, return_type
+):
+    code = f"""
+MY_CONSTANT: constant({storage_type}[2][3]) = [[1, 2], [3, 4], [5, 6]]
+
+@external
+def foo() -> {return_type}:
+    return MY_CONSTANT[0][0]
+    """
+    assert_compile_failed(lambda: get_contract(code), InvalidType)