wip(solidity/references)

astrodevs-labs · Feb 26, 2024 · 3e65d70 · 3e65d70
1 parent 270fb24
commit 3e65d70
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 8 deletions.
diff --git a/libs/solc-references/src/node_finder.rs b/libs/solc-references/src/node_finder.rs
@@ -13,7 +13,13 @@ pub struct NodeVisitor {
 }
 
 impl <'ast> Visit<'ast> for NodeVisitor {
-
+    fn visit_user_defined_type_name(&mut self,_udt: &'ast UserDefinedTypeName) {
+        if is_node_in_range(&_udt.src, &self.position, &self.source) {
+            self.node = Some(InteractableNode::UserDefinedTypeName(_udt.clone()));
+        }
+        self.above_node = Some(InteractableNode::UserDefinedTypeName(_udt.clone()));
+        visit::visit_user_defined_type_name(self, _udt);
+    }
     fn visit_contract_definition(&mut self,contract: &'ast ContractDefinition) {
         if is_node_in_range(&contract.src, &self.position, &self.source) {
             self.above_node = self.node.clone();

diff --git a/libs/solc-references/src/types.rs b/libs/solc-references/src/types.rs
@@ -152,6 +152,9 @@ pub fn get_reference_id(node: &InteractableNode) -> Option<i64> {
 
             }
         },
+        InteractableNode::UserDefinedTypeName(node) => {
+            return Some(node.referenced_declaration);
+        },
         InteractableNode::Identifier(node) => node.referenced_declaration,
         InteractableNode::MemberAccess(node) => node.referenced_declaration,
 //        InteractableNode::NewExpression(node, _) => node.type_descriptions,

diff --git a/libs/solc-references/src/usages.rs b/libs/solc-references/src/usages.rs
@@ -7,12 +7,29 @@ use crate::types::InteractableNode;
 pub struct UsagesFinder {
     pub id: i64,
     pub to_find: Vec<InteractableNode>,
-    pub is_looking_for_type: bool,
+    is_looking_for_type: bool,
+    above_node: Option<InteractableNode>,
 }
 
 impl <'ast> Visit<'ast> for UsagesFinder {
-    fn visit_using_for_directive(&mut self, using_for: &'ast UsingForDirective) {
-        visit::visit_using_for_directive(self, using_for);
+
+    fn visit_type_descriptions(&mut self, type_descriptions: &'ast TypeDescriptions) {
+        if self.above_node.clone().is_some_and(|node| {
+            match node {
+                InteractableNode::VariableDeclaration(_) => false,
+                InteractableNode::ImportDirective(_) => false,
+                _ => true
+            }
+        }) {
+            return;
+        }
+        if let Some(ident) = &type_descriptions.type_identifier {
+            if let Some(str_id) = ident.split("$").collect::<Vec<&str>>().get(2) {
+                if Ok(self.id) == str_id.parse::<i64>() {
+                    self.to_find.push(self.above_node.clone().unwrap());
+                }
+            }
+        }
     }
 
     fn visit_identifier_path(&mut self,_path: &'ast IdentifierPath) {
@@ -22,22 +39,34 @@ impl <'ast> Visit<'ast> for UsagesFinder {
         visit::visit_identifier_path(self, _path);
     }
 
+    fn visit_variable_declaration(&mut self,variable: &'ast VariableDeclaration) {
+        self.above_node = Some(InteractableNode::VariableDeclaration(variable.clone()));
+        if let Some(type_name) = &variable.type_name {
+            visit::visit_type_name(self, type_name);
+        }
+        if let Some(value) = &variable.value {
+            visit::visit_expression(self, value);
+        }
+    }
+    fn visit_import_directive(&mut self,import: &'ast ImportDirective) {
+        self.above_node = Some(InteractableNode::ImportDirective(import.clone()));
+        visit::visit_import_directive(self, import);
+    }
+
     fn visit_user_defined_type_name(&mut self,_udt: &'ast UserDefinedTypeName) {
         if _udt.referenced_declaration == self.id {
             self.to_find.push(InteractableNode::UserDefinedTypeName(_udt.clone()));
         }
+        self.above_node = Some(InteractableNode::UserDefinedTypeName(_udt.clone()));
         visit::visit_user_defined_type_name(self, _udt);
     }
 
-    fn visit_import_directive(&mut self, import: &'ast ImportDirective) {
-        visit::visit_import_directive(self, import);
-    }
-
     fn visit_identifier(&mut self,ident: &'ast Identifier) {
         let node = InteractableNode::Identifier(ident.clone());
         if ident.referenced_declaration.is_some_and(|id| self.id == id) {
             self.to_find.push(node.clone());
         }
+        self.above_node = Some(node.clone());
         visit::visit_identifier(self, ident);
     }
 
@@ -46,6 +75,7 @@ impl <'ast> Visit<'ast> for UsagesFinder {
         if member.referenced_declaration.is_some_and(|id| self.id == id) {
             self.to_find.push(node.clone());
         }
+        self.above_node = Some(node.clone());
         visit::visit_member_access(self, member);
     }
 }
@@ -56,11 +86,13 @@ impl UsagesFinder {
             id,
             to_find: Vec::new(),
             is_looking_for_type,
+            above_node: None,
         }
     }
 
     pub fn find(&mut self, ast: &SourceUnit) -> Vec<InteractableNode> {
         self.to_find = Vec::new();
+        self.above_node = None;
         eprintln!("Finding usages inf file: {}", ast.id);
         self.visit_source_unit(ast);
         eprintln!("Nodes found for file with id {}: ", ast.id);

diff --git a/libs/solc-references/src/utils.rs b/libs/solc-references/src/utils.rs
@@ -63,6 +63,7 @@ pub fn is_node_a_type(node: &InteractableNode) -> bool {
         InteractableNode::StructDefinition(_) => true,
         InteractableNode::ModifierDefinition(_) => true,
         InteractableNode::EventDefinition(_) => true,
+        InteractableNode::UserDefinedTypeName(_) => true,
         _ => false
     }
 }