Try use_def_maps: IndexVec<ScopeId, UseDefMapBuilder>

Closes #1156

Author: DavisVaughan

crates/oak_index/src/builder.rs

@@ -30,7 +30,6 @@ use crate::semantic_index::SymbolFlags;
 use crate::semantic_index::SymbolTableBuilder;
 use crate::semantic_index::Use;
 use crate::semantic_index::UseId;
-use crate::use_def_map::UseDefMap;
 use crate::use_def_map::UseDefMapBuilder;
 
 /// Build a [`SemanticIndex`] from a parsed R file.
@@ -49,9 +48,7 @@ struct SemanticIndexBuilder {
     symbol_tables: IndexVec<ScopeId, SymbolTableBuilder>,
     definitions: IndexVec<ScopeId, IndexVec<DefinitionId, Definition>>,
     uses: IndexVec<ScopeId, IndexVec<UseId, Use>>,
-    use_def_maps: IndexVec<ScopeId, UseDefMap>,
-    current_use_def: UseDefMapBuilder,
-    use_def_stack: Vec<UseDefMapBuilder>,
+    use_def_maps: IndexVec<ScopeId, UseDefMapBuilder>,
     current_scope: ScopeId,
 }
 
@@ -79,16 +76,14 @@ impl SemanticIndexBuilder {
         symbol_tables.push(SymbolTableBuilder::new());
         definitions.push(IndexVec::new());
         uses.push(IndexVec::new());
-        use_def_maps.push(UseDefMap::empty());
+        use_def_maps.push(UseDefMapBuilder::new());
 
         Self {
             scopes,
             symbol_tables,
             definitions,
             uses,
             use_def_maps,
-            current_use_def: UseDefMapBuilder::new(),
-            use_def_stack: Vec::new(),
             current_scope: file,
         }
     }
@@ -112,10 +107,7 @@ impl SemanticIndexBuilder {
         self.symbol_tables.push(SymbolTableBuilder::new());
         self.definitions.push(IndexVec::new());
         self.uses.push(IndexVec::new());
-        self.use_def_maps.push(UseDefMap::empty());
-
-        let parent_use_def = std::mem::replace(&mut self.current_use_def, UseDefMapBuilder::new());
-        self.use_def_stack.push(parent_use_def);
+        self.use_def_maps.push(UseDefMapBuilder::new());
 
         id
     }
@@ -128,13 +120,6 @@ impl SemanticIndexBuilder {
             Some(parent) => parent,
             None => panic!("`pop_scope()` called on the file scope"),
         };
-
-        let parent_use_def = match self.use_def_stack.pop() {
-            Some(builder) => builder,
-            None => panic!("`pop_scope()` called with empty use-def stack"),
-        };
-        let finalized = std::mem::replace(&mut self.current_use_def, parent_use_def).finish();
-        self.use_def_maps[id] = finalized;
     }
 
     fn add_definition(
@@ -150,9 +135,8 @@ impl SemanticIndexBuilder {
             kind,
             range,
         });
-
-        self.current_use_def.ensure_symbol(symbol_id);
-        self.current_use_def.record_definition(symbol_id, def_id);
+        self.use_def_maps[self.current_scope].ensure_symbol(symbol_id);
+        self.use_def_maps[self.current_scope].record_definition(symbol_id, def_id);
     }
 
     // Super-assignment is lexically in the current scope but binds in an
@@ -180,10 +164,8 @@ impl SemanticIndexBuilder {
             kind,
             range,
         });
-
-        let builder = self.use_def_builder_mut(target_scope);
-        builder.ensure_symbol(target_symbol);
-        builder.record_deferred_definition(target_symbol, target_def_id);
+        self.use_def_maps[target_scope].ensure_symbol(target_symbol);
+        self.use_def_maps[target_scope].record_deferred_definition(target_symbol, target_def_id);
     }
 
     // Walk up from the parent scope looking for a scope where `name` already
@@ -212,34 +194,14 @@ impl SemanticIndexBuilder {
         }
     }
 
-    fn use_def_builder_mut(&mut self, target: ScopeId) -> &mut UseDefMapBuilder {
-        if target == self.current_scope {
-            return &mut self.current_use_def;
-        }
-
-        let mut steps = 0;
-        let mut scope = self.current_scope;
-        while scope != target {
-            scope = match self.scopes[scope].parent {
-                Some(parent) => parent,
-                None => panic!("Target scope {target:?} is not an ancestor of current scope"),
-            };
-            steps += 1;
-        }
-
-        let index = self.use_def_stack.len() - steps;
-        &mut self.use_def_stack[index]
-    }
-
     fn add_use(&mut self, name: &str, range: TextRange) {
         let symbol_id = self.symbol_tables[self.current_scope].intern(name, SymbolFlags::IS_USED);
         let use_id = self.uses[self.current_scope].push(Use {
             symbol: symbol_id,
             range,
         });
-
-        self.current_use_def.ensure_symbol(symbol_id);
-        self.current_use_def.record_use(symbol_id, use_id);
+        self.use_def_maps[self.current_scope].ensure_symbol(symbol_id);
+        self.use_def_maps[self.current_scope].record_use(symbol_id, use_id);
     }
 
     // --- Recursive descent ---
@@ -349,15 +311,15 @@ impl SemanticIndexBuilder {
                     self.collect_expression(&sequence);
                 }
 
-                let pre_loop = self.current_use_def.snapshot();
+                let pre_loop = self.use_def_maps[self.current_scope].snapshot();
 
                 if let Ok(body) = stmt.body() {
                     let first_use = self.uses[self.current_scope].next_id();
                     self.collect_expression(&body);
-                    self.current_use_def.finish_loop_defs(&pre_loop, first_use);
+                    self.use_def_maps[self.current_scope].finish_loop_defs(&pre_loop, first_use);
                 }
 
-                self.current_use_def.merge(pre_loop);
+                self.use_def_maps[self.current_scope].merge(pre_loop);
             },
 
             AnyRExpression::RIfStatement(stmt) => {
@@ -366,15 +328,15 @@ impl SemanticIndexBuilder {
                     self.collect_expression(&condition);
                 }
 
-                let pre_if = self.current_use_def.snapshot();
+                let pre_if = self.use_def_maps[self.current_scope].snapshot();
 
                 // If-body (consequence)
                 if let Ok(consequence) = stmt.consequence() {
                     self.collect_expression(&consequence);
                 }
 
-                let post_if = self.current_use_def.snapshot();
-                self.current_use_def.restore(pre_if);
+                let post_if = self.use_def_maps[self.current_scope].snapshot();
+                self.use_def_maps[self.current_scope].restore(pre_if);
 
                 // Else-body (alternative), if present. If absent, the
                 // "else path" is just the pre-if state we restored to.
@@ -385,33 +347,33 @@ impl SemanticIndexBuilder {
                 }
 
                 // After: definitions from both branches are live
-                self.current_use_def.merge(post_if);
+                self.use_def_maps[self.current_scope].merge(post_if);
             },
 
             AnyRExpression::RWhileStatement(stmt) => {
                 if let Ok(condition) = stmt.condition() {
                     self.collect_expression(&condition);
                 }
 
-                let pre_loop = self.current_use_def.snapshot();
+                let pre_loop = self.use_def_maps[self.current_scope].snapshot();
 
                 if let Ok(body) = stmt.body() {
                     let first_use = self.uses[self.current_scope].next_id();
                     self.collect_expression(&body);
-                    self.current_use_def.finish_loop_defs(&pre_loop, first_use);
+                    self.use_def_maps[self.current_scope].finish_loop_defs(&pre_loop, first_use);
                 }
 
                 // Body may not execute
-                self.current_use_def.merge(pre_loop);
+                self.use_def_maps[self.current_scope].merge(pre_loop);
             },
 
             AnyRExpression::RRepeatStatement(stmt) => {
                 // Body always executes at least once, no snapshot needed
                 if let Ok(body) = stmt.body() {
-                    let pre_loop = self.current_use_def.snapshot();
+                    let pre_loop = self.use_def_maps[self.current_scope].snapshot();
                     let first_use = self.uses[self.current_scope].next_id();
                     self.collect_expression(&body);
-                    self.current_use_def.finish_loop_defs(&pre_loop, first_use);
+                    self.use_def_maps[self.current_scope].finish_loop_defs(&pre_loop, first_use);
                 }
             },
 
@@ -563,16 +525,15 @@ impl SemanticIndexBuilder {
     fn finish(mut self) -> SemanticIndex {
         self.scopes[ScopeId::from(0)].descendants.end = self.scopes.next_id();
 
-        let file_use_def_map = self.current_use_def.finish();
-        self.use_def_maps[ScopeId::from(0)] = file_use_def_map;
-
         let symbol_tables = self.symbol_tables.into_iter().map(|b| b.build()).collect();
+        let use_def_maps = self.use_def_maps.into_iter().map(|b| b.finish()).collect();
+
         SemanticIndex::new(
             self.scopes,
             symbol_tables,
             self.definitions,
             self.uses,
-            self.use_def_maps,
+            use_def_maps,
         )
     }
 }

crates/oak_index/src/use_def_map.rs

@@ -139,12 +139,6 @@ pub struct UseDefMap {
 }
 
 impl UseDefMap {
-    pub(crate) fn empty() -> Self {
-        Self {
-            bindings_by_use: IndexVec::new(),
-        }
-    }
-
     pub fn bindings_at_use(&self, use_id: UseId) -> &Bindings {
         &self.bindings_by_use[use_id]
     }