Resolve references lazily

crates/metaslang/bindings/src/builder/loader/mod.rs

@@ -386,18 +386,16 @@ pub(crate) struct Loader<'a, KT: KindTypes + 'static> {
     graph: Graph<KT>,
     remapped_nodes: HashMap<usize, NodeID>,
     injected_node_count: usize,
-    cursors: HashMap<Handle<Node>, Cursor<KT>>,
     definitions_info: HashMap<Handle<Node>, DefinitionBindingInfo<KT>>,
-    references_info: HashMap<Handle<Node>, ReferenceBindingInfo>,
+    references_info: HashMap<Handle<Node>, ReferenceBindingInfo<KT>>,
     extension_hooks: HashSet<Handle<Node>>,
 }
 
 pub(crate) struct LoadResult<KT: KindTypes + 'static> {
     #[cfg(feature = "__private_testing_utils")]
     pub graph: Graph<KT>,
-    pub cursors: HashMap<Handle<Node>, Cursor<KT>>,
     pub definitions_info: HashMap<Handle<Node>, DefinitionBindingInfo<KT>>,
-    pub references_info: HashMap<Handle<Node>, ReferenceBindingInfo>,
+    pub references_info: HashMap<Handle<Node>, ReferenceBindingInfo<KT>>,
     // Nodes where we want to inject extensions
     pub extension_hooks: HashSet<Handle<Node>>,
 }
@@ -419,7 +417,6 @@ impl<'a, KT: KindTypes + 'static> Loader<'a, KT> {
             graph: Graph::new(),
             remapped_nodes: HashMap::new(),
             injected_node_count: 0,
-            cursors: HashMap::new(),
             definitions_info: HashMap::new(),
             references_info: HashMap::new(),
             extension_hooks: HashSet::new(),
@@ -509,7 +506,6 @@ impl<'a, KT: KindTypes + 'static> Loader<'a, KT> {
         Ok(LoadResult {
             #[cfg(feature = "__private_testing_utils")]
             graph: self.graph,
-            cursors: self.cursors,
             definitions_info: self.definitions_info,
             references_info: self.references_info,
             extension_hooks: self.extension_hooks,
@@ -833,15 +829,6 @@ impl<'a, KT: KindTypes> Loader<'a, KT> {
     ) -> Result<(), BuildError> {
         let node = &self.graph[node_ref];
 
-        // For every added graph node which links to a corresponding source
-        // node, save the corresponding CST cursor so our caller can extract
-        // that info later.
-        if let Some(source_node) = node.attributes.get(SOURCE_NODE_ATTR) {
-            let syntax_node_ref = source_node.as_syntax_node_ref()?;
-            let source_node = &self.graph[syntax_node_ref];
-            self.cursors.insert(node_handle, source_node.clone());
-        }
-
         if let Some(syntax_type) = node.attributes.get(SYNTAX_TYPE_ATTR) {
             let syntax_type = syntax_type.as_str()?;
             let syntax_type = self.stack_graph.add_string(syntax_type);
@@ -865,6 +852,14 @@ impl<'a, KT: KindTypes> Loader<'a, KT> {
         let node_handle = self.node_handle_for_graph_node(node_ref);
         let stack_graph_node = &self.stack_graph[node_handle];
 
+        let cursor = if let Some(source_node) = node.attributes.get(SOURCE_NODE_ATTR) {
+            let syntax_node_ref = source_node.as_syntax_node_ref()?;
+            let source_node = &self.graph[syntax_node_ref];
+            Some(source_node.clone())
+        } else {
+            None
+        };
+
         let parents = match node.attributes.get(PARENTS_ATTR) {
             Some(parents) => {
                 parents
@@ -911,15 +906,21 @@ impl<'a, KT: KindTypes> Loader<'a, KT> {
             self.definitions_info.insert(
                 node_handle,
                 DefinitionBindingInfo {
+                    cursor: cursor.expect("Definition to have a valid source node"),
                     definiens,
                     parents,
                     extension_scope,
                     inherit_extensions,
                 },
             );
         } else if stack_graph_node.is_reference() {
-            self.references_info
-                .insert(node_handle, ReferenceBindingInfo { parents });
+            self.references_info.insert(
+                node_handle,
+                ReferenceBindingInfo {
+                    cursor: cursor.expect("Reference to have a valid source node"),
+                    parents,
+                },
+            );
         }
 
         if Self::load_flag(node, EXTENSION_HOOK_ATTR)? {

crates/metaslang/bindings/src/builder/mod.rs

@@ -1,5 +1,4 @@
 mod loader;
-mod resolver;
 
 use std::collections::{HashMap, HashSet};
 use std::fmt::Debug;
@@ -10,7 +9,6 @@ use metaslang_cst::cursor::Cursor;
 use metaslang_cst::kinds::KindTypes;
 use metaslang_graph_builder::ast::File;
 use metaslang_graph_builder::functions::Functions;
-use resolver::Resolver;
 use semver::Version;
 use stack_graphs::graph::StackGraph;
 
@@ -21,25 +19,30 @@ pub(crate) type CursorID = usize;
 pub use crate::graph::BindingGraph;
 
 pub(crate) struct DefinitionBindingInfo<KT: KindTypes + 'static> {
-    pub(crate) definiens: Cursor<KT>,
+    cursor: Cursor<KT>,
+    definiens: Cursor<KT>,
     parents: Vec<GraphHandle>,
     extension_scope: Option<GraphHandle>,
     inherit_extensions: bool,
 }
 
-pub(crate) struct ReferenceBindingInfo {
+pub(crate) struct ReferenceBindingInfo<KT: KindTypes + 'static> {
+    cursor: Cursor<KT>,
     parents: Vec<GraphHandle>,
 }
 
 pub struct BindingGraphBuilder<KT: KindTypes + 'static> {
     graph_builder_file: File<KT>,
     functions: Functions<KT>,
+    graph: ExtendedStackGraph<KT>,
+}
+
+pub(crate) struct ExtendedStackGraph<KT: KindTypes + 'static> {
     pub(crate) stack_graph: StackGraph,
-    pub(crate) cursors: HashMap<GraphHandle, Cursor<KT>>,
-    pub(crate) definitions_info: HashMap<GraphHandle, DefinitionBindingInfo<KT>>,
-    pub(crate) references_info: HashMap<GraphHandle, ReferenceBindingInfo>,
-    pub(crate) cursor_to_definitions: HashMap<CursorID, GraphHandle>,
-    pub(crate) cursor_to_references: HashMap<CursorID, GraphHandle>,
+    definitions_info: HashMap<GraphHandle, DefinitionBindingInfo<KT>>,
+    references_info: HashMap<GraphHandle, ReferenceBindingInfo<KT>>,
+    cursor_to_definitions: HashMap<CursorID, GraphHandle>,
+    cursor_to_references: HashMap<CursorID, GraphHandle>,
     extension_hooks: HashSet<GraphHandle>,
 }
 
@@ -118,25 +121,26 @@ impl<KT: KindTypes + 'static> BindingGraphBuilder<KT> {
         Self {
             graph_builder_file,
             functions,
-            stack_graph,
-            cursors: HashMap::new(),
-            definitions_info: HashMap::new(),
-            references_info: HashMap::new(),
-            cursor_to_definitions: HashMap::new(),
-            cursor_to_references: HashMap::new(),
-            extension_hooks: HashSet::new(),
+            graph: ExtendedStackGraph {
+                stack_graph,
+                definitions_info: HashMap::new(),
+                references_info: HashMap::new(),
+                cursor_to_definitions: HashMap::new(),
+                cursor_to_references: HashMap::new(),
+                extension_hooks: HashSet::new(),
+            },
         }
     }
 
     pub fn add_system_file(&mut self, file_path: &str, tree_cursor: Cursor<KT>) {
         let file_kind = FileDescriptor::System(file_path.into());
-        let file = self.stack_graph.get_or_create_file(&file_kind.as_string());
+        let file = self.graph.get_or_create_file(&file_kind);
         _ = self.add_file_internal(file, tree_cursor);
     }
 
     pub fn add_user_file(&mut self, file_path: &str, tree_cursor: Cursor<KT>) {
         let file_kind = FileDescriptor::User(file_path.into());
-        let file = self.stack_graph.get_or_create_file(&file_kind.as_string());
+        let file = self.graph.get_or_create_file(&file_kind);
         _ = self.add_file_internal(file, tree_cursor);
     }
 
@@ -147,7 +151,7 @@ impl<KT: KindTypes + 'static> BindingGraphBuilder<KT> {
         tree_cursor: Cursor<KT>,
     ) -> metaslang_graph_builder::graph::Graph<KT> {
         let file_kind = FileDescriptor::User(file_path.into());
-        let file = self.stack_graph.get_or_create_file(&file_kind.as_string());
+        let file = self.graph.get_or_create_file(&file_kind);
         let result = self.add_file_internal(file, tree_cursor);
         result.graph
     }
@@ -156,35 +160,76 @@ impl<KT: KindTypes + 'static> BindingGraphBuilder<KT> {
         let loader = Loader::new(
             &self.graph_builder_file,
             &self.functions,
-            &mut self.stack_graph,
+            &mut self.graph.stack_graph,
             file,
             tree_cursor,
         );
         let mut result = loader
             .execute(&loader::NoCancellation)
             .expect("Internal error while building bindings");
 
-        for (handle, cursor) in result.cursors.drain() {
-            let cursor_id = cursor.node().id();
-            if self.stack_graph[handle].is_definition() {
-                self.cursor_to_definitions.insert(cursor_id, handle);
-            } else {
-                self.cursor_to_references.insert(cursor_id, handle);
-            }
-            self.cursors.insert(handle, cursor);
-        }
-        self.definitions_info
+        result
+            .definitions_info
+            .iter()
+            .for_each(|(handle, definition_info)| {
+                let cursor_id = definition_info.cursor.node().id();
+                self.graph.cursor_to_definitions.insert(cursor_id, *handle);
+            });
+        result
+            .references_info
+            .iter()
+            .for_each(|(handle, reference_info)| {
+                let cursor_id = reference_info.cursor.node().id();
+                self.graph.cursor_to_references.insert(cursor_id, *handle);
+            });
+
+        self.graph
+            .definitions_info
             .extend(result.definitions_info.drain());
-        self.references_info.extend(result.references_info.drain());
-        self.extension_hooks.extend(result.extension_hooks.drain());
+        self.graph
+            .references_info
+            .extend(result.references_info.drain());
+        self.graph
+            .extension_hooks
+            .extend(result.extension_hooks.drain());
 
         result
     }
 
     pub fn build(self) -> Rc<BindingGraph<KT>> {
-        let resolver = Resolver::new(&self);
-        let resolved_references = resolver.resolve();
-        BindingGraph::build(self, resolved_references)
+        BindingGraph::build(self.graph)
+    }
+}
+
+impl<KT: KindTypes + 'static> ExtendedStackGraph<KT> {
+    fn get_or_create_file(&mut self, file_kind: &FileDescriptor) -> FileHandle {
+        self.stack_graph.get_or_create_file(&file_kind.as_string())
+    }
+
+    pub(crate) fn iter_definitions(&self) -> impl Iterator<Item = GraphHandle> + '_ {
+        self.stack_graph
+            .iter_nodes()
+            .filter(|handle| self.is_definition(*handle))
+    }
+
+    pub(crate) fn iter_references(&self) -> impl Iterator<Item = GraphHandle> + '_ {
+        self.stack_graph
+            .iter_nodes()
+            .filter(|handle| self.is_reference(*handle))
+    }
+
+    pub(crate) fn iter_files(&self) -> impl Iterator<Item = FileHandle> + '_ {
+        self.stack_graph.iter_files()
+    }
+
+    pub(crate) fn definition_at(&self, cursor: &Cursor<KT>) -> Option<GraphHandle> {
+        let cursor_id = cursor.node().id();
+        self.cursor_to_definitions.get(&cursor_id).copied()
+    }
+
+    pub(crate) fn reference_at(&self, cursor: &Cursor<KT>) -> Option<GraphHandle> {
+        let cursor_id = cursor.node().id();
+        self.cursor_to_references.get(&cursor_id).copied()
     }
 
     pub(crate) fn get_parents(&self, handle: GraphHandle) -> Vec<GraphHandle> {
@@ -201,27 +246,27 @@ impl<KT: KindTypes + 'static> BindingGraphBuilder<KT> {
         }
     }
 
-    fn is_definition(&self, handle: GraphHandle) -> bool {
+    pub(crate) fn is_definition(&self, handle: GraphHandle) -> bool {
         self.stack_graph[handle].is_definition()
     }
 
-    fn is_reference(&self, handle: GraphHandle) -> bool {
+    pub(crate) fn is_reference(&self, handle: GraphHandle) -> bool {
         self.stack_graph[handle].is_reference()
     }
 
-    fn get_extension_scope(&self, handle: GraphHandle) -> Option<GraphHandle> {
+    pub(crate) fn get_extension_scope(&self, handle: GraphHandle) -> Option<GraphHandle> {
         self.definitions_info
             .get(&handle)
             .and_then(|info| info.extension_scope)
     }
 
-    fn inherits_extensions(&self, handle: GraphHandle) -> bool {
+    pub(crate) fn inherits_extensions(&self, handle: GraphHandle) -> bool {
         self.definitions_info
             .get(&handle)
             .is_some_and(|info| info.inherit_extensions)
     }
 
-    fn get_file(&self, handle: GraphHandle) -> Option<FileDescriptor> {
+    pub(crate) fn get_file_descriptor(&self, handle: GraphHandle) -> Option<FileDescriptor> {
         self.stack_graph[handle]
             .file()
             .map(|file| FileDescriptor::from(self.stack_graph[file].name()))
@@ -230,4 +275,24 @@ impl<KT: KindTypes + 'static> BindingGraphBuilder<KT> {
     pub(crate) fn is_extension_hook(&self, node_handle: GraphHandle) -> bool {
         self.extension_hooks.contains(&node_handle)
     }
+
+    pub(crate) fn get_cursor(&self, handle: GraphHandle) -> Option<&Cursor<KT>> {
+        if self.is_definition(handle) {
+            self.definitions_info.get(&handle).map(|info| &info.cursor)
+        } else if self.is_reference(handle) {
+            self.references_info.get(&handle).map(|info| &info.cursor)
+        } else {
+            None
+        }
+    }
+
+    pub(crate) fn get_definiens_cursor(&self, handle: GraphHandle) -> Option<&Cursor<KT>> {
+        if self.is_definition(handle) {
+            self.definitions_info
+                .get(&handle)
+                .map(|info| &info.definiens)
+        } else {
+            None
+        }
+    }
 }

crates/metaslang/bindings/src/graph/definition.rs

@@ -38,33 +38,24 @@ impl<KT: KindTypes + 'static> Definition<KT> {
     }
 
     pub fn get_cursor(&self) -> &Cursor<KT> {
-        &self
-            .owner
-            .definitions
-            .get(&self.handle)
+        self.owner
+            .graph
+            .get_cursor(self.handle)
             .expect("Definition handle is valid")
-            .cursor
     }
 
     pub fn get_definiens_cursor(&self) -> &Cursor<KT> {
-        &self
-            .owner
-            .definitions
-            .get(&self.handle)
+        self.owner
+            .graph
+            .get_definiens_cursor(self.handle)
             .expect("Definition handle is valid")
-            .definiens
     }
 
     pub fn get_file(&self) -> FileDescriptor {
         self.owner
-            .get_file(
-                self.owner
-                    .definitions
-                    .get(&self.handle)
-                    .expect("Definition handle is valid")
-                    .file,
-            )
-            .expect("Definition does not have a valid file descriptor")
+            .graph
+            .get_file_descriptor(self.handle)
+            .expect("Definition to have a valid file descriptor")
     }
 }