Improve NTX4 and NTX5

natrix/ast_node.py

@@ -250,8 +250,9 @@ def is_internal(self) -> bool:
     @cached_property
     def memory_accesses(self) -> list[MemoryAccess]:
         """
-        Returns all read/write accesses inside this function by scanning for
-        variable_reads/variable_writes in nodes.
+        Returns all storage read/write accesses inside this function by scanning for
+        variable_reads/variable_writes in nodes. Excludes function arguments which
+        are read from calldata, not storage.
         """
         # Get all nodes that might have variable_reads or variable_writes
         all_nodes = self.get_descendants()
@@ -264,15 +265,24 @@ def memory_accesses(self) -> list[MemoryAccess]:
             for access_type in ("variable_reads", "variable_writes"):
                 if access_type in node.node_dict:
                     for item in node.get(access_type):
-                        accesses.append(
-                            MemoryAccess(
-                                node=node,
-                                type="read"
-                                if access_type == "variable_reads"
-                                else "write",
-                                var=item.get("name"),
+                        # Only include storage accesses (accessed via self.variable)
+                        # Check if this is an Attribute node where value.id is 'self'
+                        # This automatically excludes:
+                        # - Function arguments (accessed directly as Name nodes)
+                        # - Local variables (accessed directly as Name nodes)
+                        if (
+                            node.ast_type == "Attribute"
+                            and node.node_dict.get("value", {}).get("id") == "self"
+                        ):
+                            accesses.append(
+                                MemoryAccess(
+                                    node=node,
+                                    type="read"
+                                    if access_type == "variable_reads"
+                                    else "write",
+                                    var=item.get("name"),
+                                )
                             )
-                        )
         return accesses
 
     @cached_property

natrix/rules/common.py

@@ -279,3 +279,57 @@ def add_issue(self, node: Node, *message_args: Any) -> None:
             end_position=(end_line, end_character),
         )
         self.issues.append(issue)
+
+
+def _get_staticcall_function_mutability(staticcall_node: Node) -> str | None:
+    """
+    Extract the mutability (pure/view) of a function being called via staticcall.
+    # TODO: does not work for imported interfaces
+    Returns:
+        "pure" if the interface function is pure
+        "view" if the interface function is view
+        None if mutability cannot be determined
+    """
+    try:
+        # The staticcall_node.node_dict contains the StaticCall structure
+        # Navigate: StaticCall -> value.func (Attribute) -> type.type_decl_node
+        value_dict = staticcall_node.node_dict.get("value", {})
+        func_dict = value_dict.get("func", {})
+
+        if not func_dict or func_dict.get("ast_type") != "Attribute":
+            return None
+
+        # Get the type information which contains the declaration node reference
+        type_info = func_dict.get("type", {})
+        type_decl_node_info = type_info.get("type_decl_node", {})
+
+        if not type_decl_node_info:
+            return None
+
+        # Find the interface function definition in the module
+        module_node = staticcall_node.module_node
+        if not module_node:
+            return None
+
+        # Search for the function definition with matching node_id
+        target_node_id = type_decl_node_info.get("node_id")
+        if target_node_id is None:
+            return None
+
+        # Search through all interface definitions
+        for interface_def in module_node.get_descendants(node_type="InterfaceDef"):
+            for func_def in interface_def.get_descendants(node_type="FunctionDef"):
+                if func_def.node_dict.get("node_id") == target_node_id:
+                    # Check the function body for mutability indicators
+                    for stmt in func_def.node_dict.get("body", []):
+                        if stmt.get("ast_type") == "Expr":
+                            value = stmt.get("value", {})
+                            if value.get("ast_type") == "Name":
+                                name_id: str = value.get("id")
+                                if name_id in ["pure", "view"]:
+                                    return name_id
+
+        return None
+    except Exception:
+        # If we can't determine mutability, err on the side of caution
+        return None

natrix/rules/implicit_pure.py

@@ -2,7 +2,11 @@
 
 from typing import TYPE_CHECKING
 
-from natrix.rules.common import BaseRule, RuleRegistry
+from natrix.rules.common import (
+    BaseRule,
+    RuleRegistry,
+    _get_staticcall_function_mutability,
+)
 
 if TYPE_CHECKING:
     from natrix.ast_node import FunctionDefNode
@@ -33,6 +37,25 @@ def visit_FunctionDef(self, node: FunctionDefNode) -> None:
         accesses = node.memory_accesses
         read = any(access.type == "read" for access in accesses)
         write = any(access.type == "write" for access in accesses)
+        extcalls = node.get_descendants(node_type="ExtCall")
+        staticcalls = node.get_descendants(node_type="StaticCall")
+
+        # If there are extcalls, the function is not pure
+        if extcalls:
+            return
+
+        # If there are staticcalls, check if ALL of them are to pure functions
+        if staticcalls:
+            all_pure = True
+            for staticcall in staticcalls:
+                mutability = _get_staticcall_function_mutability(staticcall)
+                if mutability != "pure":
+                    all_pure = False
+                    break
+
+            # If not all staticcalls are pure, this function is not pure
+            if not all_pure:
+                return
 
         if not read and not write:
             self.add_issue(node, node.get("name"))

natrix/rules/implicit_view.py

@@ -2,7 +2,11 @@
 
 from typing import TYPE_CHECKING
 
-from natrix.rules.common import BaseRule, RuleRegistry
+from natrix.rules.common import (
+    BaseRule,
+    RuleRegistry,
+    _get_staticcall_function_mutability,
+)
 
 if TYPE_CHECKING:
     from natrix.ast_node import FunctionDefNode
@@ -38,6 +42,28 @@ def visit_FunctionDef(self, node: FunctionDefNode) -> None:
         accesses = node.memory_accesses
         read = any(access.type == "read" for access in accesses)
         write = any(access.type == "write" for access in accesses)
+        extcalls = node.get_descendants(node_type="ExtCall")
+        staticcalls = node.get_descendants(node_type="StaticCall")
 
-        if read and not write:
+        # Check staticcalls for view function calls
+        has_view_staticcalls = False
+        if staticcalls:
+            all_pure = True
+            for staticcall in staticcalls:
+                mutability = _get_staticcall_function_mutability(staticcall)
+                if mutability == "view":
+                    has_view_staticcalls = True
+                    all_pure = False
+                elif mutability != "pure":
+                    all_pure = False
+
+            # If all staticcalls are pure, let implicit_pure handle this
+            if all_pure:
+                return
+
+        # Function needs @view if it:
+        # 1. Reads from storage, OR
+        # 2. Makes staticcalls to view functions
+        # AND doesn't write to storage or make extcalls
+        if (read or has_view_staticcalls) and not (write or extcalls):
             self.add_issue(node, node.get("name"))

tests/contracts/test_implicit_pure.vy

@@ -0,0 +1,32 @@
+interface ITest:
+    def some_interface_pure_function() -> uint256: pure
+    def some_interface_other_pure_function() -> uint256: pure
+    def some_interface_view_function() -> uint256: view
+
+
+# This calls a pure function and should raise
+@external
+def get_value_pure(_token: address) -> uint256:
+    return staticcall ITest(_token).some_interface_pure_function()
+
+# This calls multiple pure function and should raise
+@external
+def get_values_pure(_token: address) -> uint256:
+    a: uint256 = staticcall ITest(_token).some_interface_pure_function()
+    b: uint256 = staticcall ITest(_token).some_interface_other_pure_function()
+    return a + b
+
+# This calls a view function and should not raise
+@external
+@view
+def get_value_view(_token: address) -> uint256:
+    return staticcall ITest(_token).some_interface_view_function()
+
+
+# This calls a mix of pure and view function and should not raise
+@external
+@view
+def get_value_view_and_pure(_token: address) -> uint256:
+    a: uint256 = staticcall ITest(_token).some_interface_view_function()
+    b: uint256 = staticcall ITest(_token).some_interface_pure_function()
+    return a + b

tests/contracts/test_implicit_view.vy

@@ -0,0 +1,17 @@
+interface IERC20:
+    def transfer(_to: address, _amount: uint256) -> bool: nonpayable
+    def totalSupply() -> uint256: view
+
+
+# This has no memory access but uses an extcall, so it is not a view and should not raise
+@external
+def transfer(_token: address, _target: address, _amount: uint256):
+    assert extcall IERC20(_token).transfer(
+        _target, _amount, default_return_value=True
+    )
+
+
+# This is just a staticcall to a view function and should raise
+@external
+def get_supply(_token: address) -> uint256:
+    return staticcall IERC20(_token).totalSupply()

tests/rules/test_implicit_pure.py

@@ -27,3 +27,20 @@ def test_implicit_pure(test_project_context):
         issues[3].message
         == "Function 'pure_internal_marked_as_view' does not access state but is not marked as 'pure'."
     )
+
+
+def test_implicit_pure_with_staticcalls(test_project_context):
+    rule = ImplicitPureRule()
+
+    issues = run_rule_on_file(rule, "test_implicit_pure.vy", test_project_context)
+    assert len(issues) == 2
+    assert issues[0].position == "9:0"
+    assert issues[1].position == "14:0"
+    assert (
+        issues[0].message
+        == "Function 'get_value_pure' does not access state but is not marked as 'pure'."
+    )
+    assert (
+        issues[1].message
+        == "Function 'get_values_pure' does not access state but is not marked as 'pure'."
+    )

tests/rules/test_implicit_view.py

@@ -17,3 +17,15 @@ def test_implicit_view(test_project_context):
         issues[1].message
         == "Function 'view_internal_marked_as_nothing' reads contract state but is not marked as 'view'."
     )
+
+
+def test_implicit_view_with_extcalls(test_project_context):
+    rule = ImplicitViewRule()
+
+    issues = run_rule_on_file(rule, "test_implicit_view.vy", test_project_context)
+    assert len(issues) == 1
+    assert issues[0].position == "16:0"
+    assert (
+        issues[0].message
+        == "Function 'get_supply' reads contract state but is not marked as 'view'."
+    )