Handle instance variable references

Author: vinistock

rust/rubydex/src/indexing/local_graph.rs

@@ -5,9 +5,11 @@ use crate::model::definitions::Definition;
 use crate::model::document::Document;
 use crate::model::graph::NameDependent;
 use crate::model::identity_maps::IdentityHashMap;
-use crate::model::ids::{ConstantReferenceId, DefinitionId, MethodReferenceId, NameId, StringId, UriId};
+use crate::model::ids::{
+    ConstantReferenceId, DefinitionId, InstanceVariableReferenceId, MethodReferenceId, NameId, StringId, UriId,
+};
 use crate::model::name::{Name, NameRef};
-use crate::model::references::{ConstantReference, MethodRef};
+use crate::model::references::{ConstantReference, InstanceVariableRef, MethodRef};
 use crate::model::string_ref::StringRef;
 use crate::offset::Offset;
 
@@ -19,6 +21,7 @@ type LocalGraphParts = (
     IdentityHashMap<NameId, NameRef>,
     IdentityHashMap<ConstantReferenceId, ConstantReference>,
     IdentityHashMap<MethodReferenceId, MethodRef>,
+    IdentityHashMap<InstanceVariableReferenceId, InstanceVariableRef>,
     IdentityHashMap<NameId, Vec<NameDependent>>,
 );
 
@@ -31,6 +34,7 @@ pub struct LocalGraph {
     names: IdentityHashMap<NameId, NameRef>,
     constant_references: IdentityHashMap<ConstantReferenceId, ConstantReference>,
     method_references: IdentityHashMap<MethodReferenceId, MethodRef>,
+    instance_variable_references: IdentityHashMap<InstanceVariableReferenceId, InstanceVariableRef>,
     name_dependents: IdentityHashMap<NameId, Vec<NameDependent>>,
 }
 
@@ -45,6 +49,7 @@ impl LocalGraph {
             names: IdentityHashMap::default(),
             constant_references: IdentityHashMap::default(),
             method_references: IdentityHashMap::default(),
+            instance_variable_references: IdentityHashMap::default(),
             name_dependents: IdentityHashMap::default(),
         }
     }
@@ -187,6 +192,28 @@ impl LocalGraph {
         reference_id
     }
 
+    // Instance variable references
+
+    #[must_use]
+    pub fn instance_variable_references(&self) -> &IdentityHashMap<InstanceVariableReferenceId, InstanceVariableRef> {
+        &self.instance_variable_references
+    }
+
+    pub fn add_instance_variable_reference(&mut self, reference: InstanceVariableRef) -> InstanceVariableReferenceId {
+        let reference_id = reference.id();
+
+        if self
+            .instance_variable_references
+            .insert(reference_id, reference)
+            .is_some()
+        {
+            debug_assert!(false, "InstanceVariableReferenceId collision in local graph");
+        }
+
+        self.document.add_instance_variable_reference(reference_id);
+        reference_id
+    }
+
     // Diagnostics
 
     #[must_use]
@@ -218,6 +245,7 @@ impl LocalGraph {
             self.names,
             self.constant_references,
             self.method_references,
+            self.instance_variable_references,
             self.name_dependents,
         )
     }

rust/rubydex/src/indexing/ruby_indexer.rs

@@ -13,7 +13,7 @@ use crate::model::definitions::{
 use crate::model::document::Document;
 use crate::model::ids::{DefinitionId, NameId, StringId, UriId};
 use crate::model::name::{Name, ParentScope};
-use crate::model::references::{ConstantReference, MethodRef};
+use crate::model::references::{ConstantReference, InstanceVariableRef, MethodRef};
 use crate::model::visibility::Visibility;
 use crate::offset::Offset;
 
@@ -555,6 +555,15 @@ impl<'a> RubyIndexer<'a> {
         definition_id
     }
 
+    fn add_instance_variable_reference(&mut self, location: &ruby_prism::Location) {
+        let name = Self::location_to_string(location);
+        let str_id = self.local_graph.intern_string(name);
+        let offset = Offset::from_prism_location(location);
+        let parent_nesting_id = self.parent_nesting_id();
+        let reference = InstanceVariableRef::new(str_id, self.uri_id, offset, parent_nesting_id);
+        self.local_graph.add_instance_variable_reference(reference);
+    }
+
     /// Adds a class variable definition.
     ///
     /// Class variables use lexical scoping - they belong to the lexically enclosing class/module,
@@ -2193,7 +2202,7 @@ impl Visit<'_> for RubyIndexer<'_> {
     }
 
     fn visit_instance_variable_operator_write_node(&mut self, node: &ruby_prism::InstanceVariableOperatorWriteNode) {
-        self.add_instance_variable_definition(&node.name_loc());
+        self.add_instance_variable_reference(&node.name_loc());
         self.visit(&node.value());
     }
 
@@ -2207,6 +2216,10 @@ impl Visit<'_> for RubyIndexer<'_> {
         self.visit(&node.value());
     }
 
+    fn visit_instance_variable_read_node(&mut self, node: &ruby_prism::InstanceVariableReadNode) {
+        self.add_instance_variable_reference(&node.location());
+    }
+
     fn visit_class_variable_and_write_node(&mut self, node: &ruby_prism::ClassVariableAndWriteNode) {
         self.add_class_variable_definition(&node.name_loc());
         self.visit(&node.value());

rust/rubydex/src/indexing/ruby_indexer_tests.rs

@@ -103,6 +103,34 @@ macro_rules! assert_not_promotable {
     }};
 }
 
+macro_rules! assert_instance_variable_references_eq {
+    ($context:expr, $expected_names:expr) => {{
+        let mut actual_references = $context
+            .graph()
+            .instance_variable_references()
+            .values()
+            .map(|r| {
+                (
+                    r.offset().start(),
+                    $context.graph().strings().get(r.str_id()).unwrap().as_str(),
+                )
+            })
+            .collect::<Vec<_>>();
+
+        actual_references.sort();
+
+        let actual_names = actual_references.iter().map(|(_, name)| *name).collect::<Vec<_>>();
+
+        assert_eq!(
+            $expected_names,
+            actual_names.as_slice(),
+            "instance variable references mismatch: expected `{:?}`, got `{:?}`",
+            $expected_names,
+            actual_names
+        );
+    }};
+}
+
 fn index_source(source: &str) -> LocalGraphTest {
     LocalGraphTest::new("file:///foo.rb", source)
 }
@@ -1801,6 +1829,38 @@ fn constant_with_colon_colon_call_is_promotable() {
     });
 }
 
+#[test]
+fn instance_variable_reads_are_indexed_as_references() {
+    let context = index_source({
+        "
+        class Foo
+          def bar
+            @baz
+          end
+        end
+        "
+    });
+
+    assert_no_local_diagnostics!(&context);
+    assert_instance_variable_references_eq!(context, ["@baz"]);
+}
+
+#[test]
+fn instance_variable_operator_writes_are_references() {
+    let context = index_source({
+        "
+        class Foo
+          def bar
+            @baz += 1
+          end
+        end
+        "
+    });
+
+    assert_no_local_diagnostics!(&context);
+    assert_instance_variable_references_eq!(context, ["@baz"]);
+}
+
 mod constant_tests {
     use super::*;
 
@@ -2075,14 +2135,12 @@ mod variable_tests {
               @baz, @qux = 3, 4
             end
 
-            @bar &= 5
-            @baz &&= 6
-            @qux ||= 7
+            @bar &&= 5
+            @baz ||= 6
 
             class Bar
-              @foo &= 8
-              @bar &&= 9
-              @baz ||= 10
+              @foo &&= 8
+              @bar ||= 9
             end
             "
         });
@@ -2124,29 +2182,18 @@ mod variable_tests {
             assert!(def.lexical_nesting_id().is_none());
         });
 
-        assert_definition_at!(&context, "10:1-10:5", InstanceVariable, |def| {
-            assert_def_str_eq!(&context, def, "@qux");
-            assert!(def.lexical_nesting_id().is_none());
-        });
-
-        assert_definition_at!(&context, "12:1-16:4", Class, |bar_class_def| {
-            assert_definition_at!(&context, "13:3-13:7", InstanceVariable, |def| {
+        assert_definition_at!(&context, "11:1-14:4", Class, |bar_class_def| {
+            assert_definition_at!(&context, "12:3-12:7", InstanceVariable, |def| {
                 assert_def_str_eq!(&context, def, "@foo");
                 assert_eq!(bar_class_def.id(), def.lexical_nesting_id().unwrap());
                 assert_eq!(bar_class_def.members()[0], def.id());
             });
 
-            assert_definition_at!(&context, "14:3-14:7", InstanceVariable, |def| {
+            assert_definition_at!(&context, "13:3-13:7", InstanceVariable, |def| {
                 assert_def_str_eq!(&context, def, "@bar");
                 assert_eq!(bar_class_def.id(), def.lexical_nesting_id().unwrap());
                 assert_eq!(bar_class_def.members()[1], def.id());
             });
-
-            assert_definition_at!(&context, "15:3-15:7", InstanceVariable, |def| {
-                assert_def_str_eq!(&context, def, "@baz");
-                assert_eq!(bar_class_def.id(), def.lexical_nesting_id().unwrap());
-                assert_eq!(bar_class_def.members()[2], def.id());
-            });
         });
     }
 

rust/rubydex/src/model/declaration.rs

@@ -448,6 +448,43 @@ impl Declaration {
             _ => unreachable!("Cannot remove constant reference from {} declaration", self.kind()),
         }
     }
+
+    /// Returns the instance variable reference IDs for `InstanceVariable` declarations.
+    /// Returns `None` for other declaration types.
+    #[must_use]
+    pub fn instance_variable_references(&self) -> Option<&IdentityHashSet<InstanceVariableReferenceId>> {
+        match self {
+            Declaration::InstanceVariable(it) => Some(it.references()),
+            _ => None,
+        }
+    }
+
+    /// Adds an instance variable reference to this declaration.
+    ///
+    /// # Panics
+    ///
+    /// Panics if called on a declaration that doesn't track instance variable references.
+    pub fn add_instance_variable_reference(&mut self, reference_id: InstanceVariableReferenceId) {
+        match self {
+            Declaration::InstanceVariable(it) => it.add_reference(reference_id),
+            _ => unreachable!("Cannot add instance variable reference to {} declaration", self.kind()),
+        }
+    }
+
+    /// Removes an instance variable reference from this declaration.
+    ///
+    /// # Panics
+    ///
+    /// Panics if called on a declaration that doesn't track instance variable references.
+    pub fn remove_instance_variable_reference(&mut self, reference_id: &InstanceVariableReferenceId) {
+        match self {
+            Declaration::InstanceVariable(it) => it.remove_reference(reference_id),
+            _ => unreachable!(
+                "Cannot remove instance variable reference from {} declaration",
+                self.kind()
+            ),
+        }
+    }
 }
 
 #[derive(Debug)]

rust/rubydex/src/model/document.rs

@@ -4,7 +4,7 @@ use line_index::LineIndex;
 use url::Url;
 
 use crate::diagnostic::Diagnostic;
-use crate::model::ids::{ConstantReferenceId, DefinitionId, MethodReferenceId};
+use crate::model::ids::{ConstantReferenceId, DefinitionId, InstanceVariableReferenceId, MethodReferenceId};
 
 // Represents a document currently loaded into memory. Identified by its unique URI, it holds the edges to all
 // definitions and references discovered in it
@@ -15,6 +15,7 @@ pub struct Document {
     definition_ids: Vec<DefinitionId>,
     method_reference_ids: Vec<MethodReferenceId>,
     constant_reference_ids: Vec<ConstantReferenceId>,
+    instance_variable_reference_ids: Vec<InstanceVariableReferenceId>,
     diagnostics: Vec<Diagnostic>,
 }
 
@@ -27,6 +28,7 @@ impl Document {
             definition_ids: Vec::new(),
             method_reference_ids: Vec::new(),
             constant_reference_ids: Vec::new(),
+            instance_variable_reference_ids: Vec::new(),
             diagnostics: Vec::new(),
         }
     }
@@ -73,6 +75,15 @@ impl Document {
         self.constant_reference_ids.push(reference_id);
     }
 
+    #[must_use]
+    pub fn instance_variable_references(&self) -> &[InstanceVariableReferenceId] {
+        &self.instance_variable_reference_ids
+    }
+
+    pub fn add_instance_variable_reference(&mut self, reference_id: InstanceVariableReferenceId) {
+        self.instance_variable_reference_ids.push(reference_id);
+    }
+
     #[must_use]
     pub fn diagnostics(&self) -> &[Diagnostic] {
         &self.diagnostics