Use cursor.node().id() to map cursors to bindings handles

This avoids a linear search of the full set of cursors whenever we want to find
a `Definition` or `Reference` from from a `Cursor`.

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

@@ -247,7 +247,6 @@
 mod cancellation;
 mod functions;
 
-use std::collections::hash_map::Drain;
 use std::collections::{HashMap, HashSet};
 use std::path::Path;
 
@@ -321,7 +320,7 @@ pub const ROOT_NODE_VAR: &str = "ROOT_NODE";
 /// Name of the variable used to pass the file path.
 pub const FILE_PATH_VAR: &str = "FILE_PATH";
 
-pub struct Builder<'a, KT: KindTypes> {
+pub(crate) struct Builder<'a, KT: KindTypes> {
  msgb: &'a GraphBuilderFile<KT>,
  functions: &'a Functions<KT>,
  stack_graph: &'a mut StackGraph,
@@ -334,6 +333,12 @@ pub struct Builder<'a, KT: KindTypes> {
  definiens: HashMap<Handle<Node>, Cursor<KT>>,
 }
 
+pub(crate) struct BuildResult<KT: KindTypes> {
+ pub graph: Graph<KT>,
+ pub cursors: HashMap<Handle<Node>, Cursor<KT>>,
+ pub definiens: HashMap<Handle<Node>, Cursor<KT>>,
+}
+
 impl<'a, KT: KindTypes + 'static> Builder<'a, KT> {
  pub fn new(
  msgb: &'a GraphBuilderFile<KT>,
@@ -386,13 +391,11 @@ impl<'a, KT: KindTypes + 'static> Builder<'a, KT> {
  variables
  }
 
- #[cfg(feature = "__private_testing_utils")]
- pub(crate) fn graph(self) -> Graph<KT> {
- self.graph
- }
-
  /// Executes this builder.
- pub fn build(&mut self, cancellation_flag: &dyn CancellationFlag) -> Result<(), BuildError> {
+ pub fn build(
+ mut self,
+ cancellation_flag: &dyn CancellationFlag,
+ ) -> Result<BuildResult<KT>, BuildError> {
  let variables = self.build_global_variables();
 
  let config = ExecutionConfig::new(self.functions, &variables)
@@ -419,7 +422,13 @@ impl<'a, KT: KindTypes + 'static> Builder<'a, KT> {
  &(cancellation_flag as &dyn CancellationFlag),
  )?;
 
- self.load(cancellation_flag)
+ self.load(cancellation_flag)?;
+
+ Ok(BuildResult {
+ graph: self.graph,
+ cursors: self.cursors,
+ definiens: self.definiens,
+ })
  }
 
  /// Create a graph node to represent the stack graph node. It is the callers responsibility to
@@ -430,14 +439,6 @@ impl<'a, KT: KindTypes + 'static> Builder<'a, KT> {
  self.injected_node_count += 1;
  node
  }
-
- pub fn extract_cursors(&mut self) -> Drain<'_, Handle<Node>, Cursor<KT>> {
- self.cursors.drain()
- }
-
- pub fn extract_definiens(&mut self) -> Drain<'_, Handle<Node>, Cursor<KT>> {
- self.definiens.drain()
- }
 }
 
 /// An error that can occur while loading a stack graph from a TSG file

crates/metaslang/bindings/src/lib.rs

@@ -6,6 +6,7 @@ use std::fmt::Debug;
 use std::iter::once;
 use std::sync::Arc;
 
+use builder::BuildResult;
 use metaslang_cst::cursor::Cursor;
 use metaslang_cst::KindTypes;
 use metaslang_graph_builder::ast::File;
@@ -19,13 +20,16 @@ use stack_graphs::stitching::{
 
 type Builder<'a, KT> = builder::Builder<'a, KT>;
 type GraphHandle = stack_graphs::arena::Handle<stack_graphs::graph::Node>;
+type CursorID = usize;
 
 pub struct Bindings<KT: KindTypes + 'static> {
  graph_builder_file: File<KT>,
  functions: Functions<KT>,
  stack_graph: StackGraph,
  cursors: HashMap<GraphHandle, Cursor<KT>>,
  definiens: HashMap<GraphHandle, Cursor<KT>>,
+ cursor_to_definitions: HashMap<CursorID, GraphHandle>,
+ cursor_to_references: HashMap<CursorID, GraphHandle>,
 }
 
 pub trait PathResolver {
@@ -57,6 +61,8 @@ impl<KT: KindTypes + 'static> Bindings<KT> {
  stack_graph,
  cursors: HashMap::new(),
  definiens: HashMap::new(),
+ cursor_to_definitions: HashMap::new(),
+ cursor_to_references: HashMap::new(),
  }
  }
 
@@ -70,26 +76,36 @@ impl<KT: KindTypes + 'static> Bindings<KT> {
  file_path: &str,
  tree_cursor: Cursor<KT>,
  ) -> metaslang_graph_builder::graph::Graph<KT> {
- let builder = self.add_file_internal(file_path, tree_cursor);
- builder.graph()
+ let result = self.add_file_internal(file_path, tree_cursor);
+ result.graph
  }
 
- fn add_file_internal(&mut self, file_path: &str, tree_cursor: Cursor<KT>) -> Builder<'_, KT> {
+ fn add_file_internal(&mut self, file_path: &str, tree_cursor: Cursor<KT>) -> BuildResult<KT> {
  let file = self.stack_graph.get_or_create_file(file_path);
 
- let mut builder = Builder::new(
+ let builder = Builder::new(
  &self.graph_builder_file,
  &self.functions,
  &mut self.stack_graph,
  file,
  tree_cursor,
  );
- builder
+ let mut result = builder
  .build(&builder::NoCancellation)
  .expect("Internal error while building bindings");
- self.cursors.extend(builder.extract_cursors());
- self.definiens.extend(builder.extract_definiens());
- builder
+
+ 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.definiens.extend(result.definiens.drain());
+
+ result
  }
 
  pub fn all_definitions(&self) -> impl Iterator<Item = Definition<'_, KT>> + '_ {
@@ -113,27 +129,23 @@ impl<KT: KindTypes + 'static> Bindings<KT> {
  }
 
  pub fn definition_at(&self, cursor: &Cursor<KT>) -> Option<Definition<'_, KT>> {
- for (handle, handle_cursor) in &self.cursors {
- if handle_cursor == cursor && self.stack_graph[*handle].is_definition() {
- return Some(Definition {
- owner: self,
- handle: *handle,
- });
- }
- }
- None
+ let cursor_id = cursor.node().id();
+ self.cursor_to_definitions
+ .get(&cursor_id)
+ .map(|handle| Definition {
+ owner: self,
+ handle: *handle,
+ })
  }
 
  pub fn reference_at(&self, cursor: &Cursor<KT>) -> Option<Reference<'_, KT>> {
- for (handle, handle_cursor) in &self.cursors {
- if handle_cursor == cursor && self.stack_graph[*handle].is_reference() {
- return Some(Reference {
- owner: self,
- handle: *handle,
- });
- }
- }
- None
+ let cursor_id = cursor.node().id();
+ self.cursor_to_references
+ .get(&cursor_id)
+ .map(|handle| Reference {
+ owner: self,
+ handle: *handle,
+ })
  }
 }