find references for function-local variables

vyper_lsp/ast.py

@@ -107,6 +107,12 @@ def get_internal_functions(self):
 
         return inernal_nodes
 
+    def get_internal_functions_names(self):
+        if self.ast_data is None:
+            return []
+
+        return [node.name for node in self.get_internal_functions()]
+
     def find_nodes_referencing_internal_function(self, function: str):
         if self.ast_data is None:
             return []
@@ -244,3 +250,26 @@ def find_top_level_node_at_pos(self, pos: Position) -> Optional[VyperNode]:
         for node in self.ast_data.get_children():
             if node.lineno <= pos.line and node.end_lineno >= pos.line:
                 return node
+
+    def find_nodes_referencing_symbol(self, symbol: str):
+        if self.ast_data is None:
+            return []
+
+        return_nodes = []
+
+        for node in self.ast_data.get_descendants(nodes.Name, {"id": symbol}):
+            if isinstance(node.get_ancestor(), nodes.Dict):
+                if node in node.get_ancestor().keys:
+                    continue
+            else:
+                return_nodes.append(node)
+
+        return return_nodes
+
+    @classmethod
+    def create_new_instance(cls, ast):
+        # Create a new instance
+        new_instance = super(AST, cls).__new__(cls)
+        # Optionally, initialize the new instance
+        new_instance.ast_data = ast
+        return new_instance

vyper_lsp/navigation.py

@@ -4,7 +4,7 @@
 from typing import List, Optional
 
 from pygls.workspace import Document
-from vyper.ast import EnumDef
+from vyper.ast import EnumDef, FunctionDef
 from vyper_lsp.ast import AST
 from vyper_lsp.utils import get_expression_at_cursor, get_word_at_cursor
 
@@ -47,6 +47,7 @@ def find_type_declaration(self, word: str) -> Optional[Range]:
     def find_references(self, doc: Document, pos: Position) -> List[Range]:
         og_line = doc.lines[pos.line]
         word = get_word_at_cursor(og_line, pos.character)
+        expression = get_expression_at_cursor(og_line, pos.character)
         if self.ast.ast_data is None:
             return []
         references = []
@@ -70,6 +71,16 @@ def find_references(self, doc: Document, pos: Position) -> List[Range]:
                     end=Position(line=ref.end_lineno - 1, character=ref.end_col_offset),
                 )
                 references.append(range)
+        elif word in self.ast.get_internal_functions_names() and (
+            og_line.startswith("def") or expression.startswith("self.")
+        ):
+            refs = self.ast.find_nodes_referencing_internal_function(word)
+            for ref in refs:
+                range = Range(
+                    start=Position(line=ref.lineno - 1, character=ref.col_offset),
+                    end=Position(line=ref.end_lineno - 1, character=ref.end_col_offset),
+                )
+                references.append(range)
         elif not og_line[0].isspace() and word in self.ast.get_state_variables():
             if "constant(" in og_line:
                 refs = self.ast.find_nodes_referencing_constant(word)
@@ -92,6 +103,16 @@ def find_references(self, doc: Document, pos: Position) -> List[Range]:
                     end=Position(line=ref.end_lineno - 1, character=ref.end_col_offset),
                 )
                 references.append(range)
+        elif isinstance(top_level_node, FunctionDef):
+            refs = AST.create_new_instance(
+                top_level_node
+            ).find_nodes_referencing_symbol(word)
+            for ref in refs:
+                range = Range(
+                    start=Position(line=ref.lineno - 1, character=ref.col_offset),
+                    end=Position(line=ref.end_lineno - 1, character=ref.end_col_offset),
+                )
+                references.append(range)
         return references
 
     def find_declaration(self, document: Document, pos: Position) -> Optional[Range]: