fix codegen failure with nonreentrant keys

tests/parser/features/decorators/test_nonreentrant.py

@@ -23,6 +23,9 @@ def updated_protected():
 interface Callback:
     def updated(): nonpayable
     def updated_protected(): nonpayable
+interface Self:
+    def protected_function(val: String[100], do_callback: bool) -> uint256: nonpayable
+    def protected_function2(val: String[100], do_callback: bool) -> uint256: nonpayable
 
 special_value: public(String[100])
 callback: public(Callback)
@@ -42,6 +45,16 @@ def protected_function(val: String[100], do_callback: bool) -> uint256:
     else:
         return 2
 
+@external
+@nonreentrant('protect_special_value')
+def protected_function2(val: String[100], do_callback: bool) -> uint256:
+    self.special_value = val
+    if do_callback:
+        # call other function with same nonreentrancy key
+        Self(self).protected_function(val, False)
+        return 1
+    return 2
+
 @external
 def unprotected_function(val: String[100], do_callback: bool):
     self.special_value = val
@@ -66,6 +79,11 @@ def unprotected_function(val: String[100], do_callback: bool):
 
     assert_tx_failed(lambda: reentrant_contract.protected_function("zzz value", True, transact={}))
 
+    reentrant_contract.protected_function2("another value", False, transact={})
+    assert reentrant_contract.special_value() == "another value"
+
+    assert_tx_failed(lambda: reentrant_contract.protected_function2("zzz value", True, transact={}))
+
 
 def test_disallow_on_init_function(get_contract):
     # nonreentrant has no effect when used on the __init__ fn

vyper/semantics/validation/data_positions.py

@@ -1,5 +1,6 @@
 # TODO this doesn't really belong in "validation"
 import math
+from typing import Dict
 
 from vyper import ast as vy_ast
 from vyper.semantics.types.bases import StorageSlot
@@ -28,7 +29,7 @@ def set_storage_slots(vyper_module: vy_ast.Module) -> StorageLayout:
     # note storage is word-addressable, not byte-addressable
     storage_slot = 0
 
-    ret = {}
+    ret: Dict[str, Dict] = {}
 
     for node in vyper_module.get_children(vy_ast.FunctionDef):
         type_ = node._metadata["type"]
@@ -38,8 +39,11 @@ def set_storage_slots(vyper_module: vy_ast.Module) -> StorageLayout:
         variable_name = f"nonreentrant.{type_.nonreentrant}"
 
         # a nonreentrant key can appear many times in a module but it
-        # only takes one slot. ignore it after the first time we see it.
+        # only takes one slot. after the first time we see it, do not
+        # increment the storage slot.
         if variable_name in ret:
+            _slot = ret[variable_name]["slot"]
+            type_.set_reentrancy_key_position(StorageSlot(_slot))
             continue
 
         type_.set_reentrancy_key_position(StorageSlot(storage_slot))