pr 04.c

crates/prometeu-compiler/src/analysis/symbols/mod.rs

@@ -52,6 +52,11 @@ pub struct DefIndex {
     symbols: HashMap<DefKey, SymbolId>,
 }
 
+#[derive(Debug, Default, Clone)]
+pub struct RefIndex {
+    refs: Vec<Vec<Span>>,
+}
+
 impl SymbolArena {
     pub fn new() -> Self {
         Self { symbols: Vec::new() }
@@ -94,6 +99,32 @@ impl DefIndex {
     }
 }
 
+impl RefIndex {
+    pub fn new() -> Self {
+        Self { refs: Vec::new() }
+    }
+
+    pub fn ensure_symbol(&mut self, symbol_id: SymbolId) {
+        let index = symbol_id.0 as usize;
+        if index >= self.refs.len() {
+            self.refs.resize_with(index + 1, Vec::new);
+        }
+    }
+
+    pub fn record_ref(&mut self, symbol_id: SymbolId, span: Span) {
+        self.ensure_symbol(symbol_id);
+        self.refs[symbol_id.0 as usize].push(span);
+    }
+
+    pub fn refs_of(&self, symbol_id: SymbolId) -> &[Span] {
+        const EMPTY: [Span; 0] = [];
+        self.refs
+            .get(symbol_id.0 as usize)
+            .map(|refs| refs.as_slice())
+            .unwrap_or(&EMPTY)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -165,4 +196,20 @@ mod tests {
         assert_eq!(index.get(type_key), Some(SymbolId(0)));
         assert_eq!(index.get(value_key), Some(SymbolId(1)));
     }
+
+    #[test]
+    fn ref_index_records_refs_per_symbol() {
+        let mut index = RefIndex::new();
+        let span_a1 = Span::new(0, 1, 2);
+        let span_a2 = Span::new(0, 3, 4);
+        let span_b1 = Span::new(1, 10, 12);
+
+        index.record_ref(SymbolId(2), span_a1);
+        index.record_ref(SymbolId(2), span_a2);
+        index.record_ref(SymbolId(5), span_b1);
+
+        assert_eq!(index.refs_of(SymbolId(2)), &[span_a1, span_a2]);
+        assert_eq!(index.refs_of(SymbolId(5)), &[span_b1]);
+        assert!(index.refs_of(SymbolId(9)).is_empty());
+    }
 }