fix[tool]: fix layout export with nonreentrancy pragma on

tests/unit/cli/storage_layout/test_storage_layout.py

@@ -380,3 +380,27 @@ def test_multiple_compile_codes(make_input_bundle):
 
     assert layout1 == {"x": {"slot": start_slot, "type": "uint256", "n_slots": 1}}
     assert layout2 == {"x": {"slot": start_slot + 1, "type": "uint256", "n_slots": 1}}
+
+
+# test that the nonreentrancy lock gets excported when the nonreentrant pragma
+# is on and the public getters are nonreentrant
+def test_lock_export_with_nonreentrant_pragma(make_input_bundle):
+    main = """
+# pragma nonreentrancy on
+a: public(uint256)
+    """
+    out = compile_code(main, output_formats=["layout"])["layout"]
+
+    if version_check(begin="cancun"):
+        storage_layout = {"a": {"type": "uint256", "n_slots": 1, "slot": 0}}
+        transient_layout = {
+            "$.nonreentrant_key": {"type": "nonreentrant lock", "slot": 0, "n_slots": 1}
+        }
+        assert transient_layout == out["transient_storage_layout"]
+    else:
+        storage_layout = {
+            "a": {"type": "uint256", "n_slots": 1, "slot": 1},
+            "$.nonreentrant_key": {"type": "nonreentrant lock", "slot": 0, "n_slots": 1},
+        }
+
+    assert storage_layout == out["storage_layout"]

tests/unit/cli/storage_layout/test_storage_layout_overrides.py

@@ -435,3 +435,21 @@ def test_module_overlap(make_input_bundle):
             input_bundle=input_bundle,
             storage_layout_override=override,
         )
+
+
+def test_override_with_nonreentrant_pragma(make_input_bundle):
+    code = """
+# pragma nonreentrancy on
+a: public(uint256)
+    """
+
+    if version_check(begin="cancun"):
+        override = {"a": {"type": "uint256", "n_slots": 1, "slot": 0}}
+    else:
+        override = {
+            "a": {"type": "uint256", "n_slots": 1, "slot": 0},
+            "$.nonreentrant_key": {"type": "nonreentrant lock", "n_slots": 1, "slot": 20},
+        }
+
+    # note: compile_code checks roundtrip of the override
+    compile_code(code, storage_layout_override=override)

vyper/semantics/analysis/data_positions.py

@@ -182,6 +182,17 @@ def _allocate_with_overrides(vyper_module: vy_ast.Module, layout: StorageLayout)
     _allocate_with_overrides_r(vyper_module, layout, allocator, nonreentrant_slot, [])
 
 
+def _get_func_defs(vyper_module: vy_ast.Module):
+    funcdefs = vyper_module.get_children(vy_ast.FunctionDef)
+    for vardecl in vyper_module.get_children(vy_ast.VariableDecl):
+        if not vardecl.is_public:
+            # no getter
+            continue
+        funcdefs.append(vardecl._expanded_getter)
+
+    return funcdefs
+
+
 def _allocate_with_overrides_r(
     vyper_module: vy_ast.Module,
     layout: StorageLayout,
@@ -190,12 +201,7 @@ def _allocate_with_overrides_r(
     path: list[str],
 ):
     # Search through function definitions to find non-reentrant functions
-    funcdefs = vyper_module.get_children(vy_ast.FunctionDef)
-    for vardecl in vyper_module.get_children(vy_ast.VariableDecl):
-        if not vardecl.is_public:
-            # no getter
-            continue
-        funcdefs.append(vardecl._expanded_getter)
+    funcdefs = _get_func_defs(vyper_module)
 
     for node in funcdefs:
         fn_t = node._metadata["func_type"]
@@ -391,7 +397,8 @@ def _generate_layout_export_r(vyper_module):
             raise CompilerPanic("unreachable")
         ret[layout_key][node.target.id] = item
 
-    for fn in vyper_module.get_children(vy_ast.FunctionDef):
+    funcdefs = _get_func_defs(vyper_module)
+    for fn in funcdefs:
         fn_t = fn._metadata["func_type"]
         if not fn_t.nonreentrant:
             continue