feat: add def use analysis inside of module (#19)

* dev: add snapshot map

* feat: add def use analysis inside of module

crates/tinymist-query/src/adt/mod.rs

@@ -0,0 +1 @@
+pub mod snapshot_map;

crates/tinymist-query/src/adt/snapshot_map.rs

@@ -0,0 +1,162 @@
+//! Upstream [rustc_data_structures::snapshot_map].
+//! Last checked commit: f4bb4500ddb4
+//! Last checked time:   2023-12-28
+//!
+//! [rustc_data_structures::snapshot_map]: https://github.com/rust-lang/rust/blob/master/compiler/rustc_data_structures/src/snapshot_map/mod.rs
+
+#![allow(missing_docs)]
+#![allow(unused)]
+
+use ena::undo_log::{Rollback, Snapshots, UndoLogs, VecLog};
+use std::borrow::{Borrow, BorrowMut};
+use std::hash::Hash;
+use std::marker::PhantomData;
+use std::ops;
+
+pub use ena::undo_log::Snapshot;
+
+type FxHashMap<K, V> = fxhash::FxHashMap<K, V>;
+
+pub type SnapshotMapStorage<K, V> = SnapshotMap<K, V, FxHashMap<K, V>, ()>;
+pub type SnapshotMapRef<'a, K, V, L> = SnapshotMap<K, V, &'a mut FxHashMap<K, V>, &'a mut L>;
+
+#[derive(Clone)]
+pub struct SnapshotMap<K, V, M = FxHashMap<K, V>, L = VecLog<UndoLog<K, V>>> {
+    map: M,
+    undo_log: L,
+    _marker: PhantomData<(K, V)>,
+}
+
+// HACK(eddyb) manual impl avoids `Default` bounds on `K` and `V`.
+impl<K, V, M, L> Default for SnapshotMap<K, V, M, L>
+where
+    M: Default,
+    L: Default,
+{
+    fn default() -> Self {
+        SnapshotMap {
+            map: Default::default(),
+            undo_log: Default::default(),
+            _marker: PhantomData,
+        }
+    }
+}
+
+#[derive(Clone)]
+pub enum UndoLog<K, V> {
+    Inserted(K),
+    Overwrite(K, V),
+    Purged,
+}
+
+impl<K, V, M, L> SnapshotMap<K, V, M, L> {
+    #[inline]
+    pub fn with_log<L2>(&mut self, undo_log: L2) -> SnapshotMap<K, V, &mut M, L2> {
+        SnapshotMap {
+            map: &mut self.map,
+            undo_log,
+            _marker: PhantomData,
+        }
+    }
+}
+
+impl<K, V, M, L> SnapshotMap<K, V, M, L>
+where
+    K: Hash + Clone + Eq,
+    M: BorrowMut<FxHashMap<K, V>> + Borrow<FxHashMap<K, V>>,
+    L: UndoLogs<UndoLog<K, V>>,
+{
+    pub fn clear(&mut self) {
+        self.map.borrow_mut().clear();
+        self.undo_log.clear();
+    }
+
+    pub fn insert(&mut self, key: K, value: V) -> bool {
+        match self.map.borrow_mut().insert(key.clone(), value) {
+            None => {
+                self.undo_log.push(UndoLog::Inserted(key));
+                true
+            }
+            Some(old_value) => {
+                self.undo_log.push(UndoLog::Overwrite(key, old_value));
+                false
+            }
+        }
+    }
+
+    pub fn remove(&mut self, key: K) -> bool {
+        match self.map.borrow_mut().remove(&key) {
+            Some(old_value) => {
+                self.undo_log.push(UndoLog::Overwrite(key, old_value));
+                true
+            }
+            None => false,
+        }
+    }
+
+    pub fn get<Q: ?Sized>(&self, k: &Q) -> Option<&V>
+    where
+        K: Borrow<Q>,
+        Q: Hash + Eq,
+    {
+        self.map.borrow().get(k)
+    }
+}
+
+impl<K, V> SnapshotMap<K, V>
+where
+    K: Hash + Clone + Eq,
+{
+    pub fn snapshot(&mut self) -> Snapshot {
+        self.undo_log.start_snapshot()
+    }
+
+    pub fn commit(&mut self, snapshot: Snapshot) {
+        self.undo_log.commit(snapshot)
+    }
+
+    pub fn rollback_to(&mut self, snapshot: Snapshot) {
+        let map = &mut self.map;
+        self.undo_log.rollback_to(|| map, snapshot)
+    }
+}
+
+impl<'k, K, V, M, L> ops::Index<&'k K> for SnapshotMap<K, V, M, L>
+where
+    K: Hash + Clone + Eq,
+    M: Borrow<FxHashMap<K, V>>,
+{
+    type Output = V;
+    fn index(&self, key: &'k K) -> &V {
+        &self.map.borrow()[key]
+    }
+}
+
+impl<K, V, M, L> Rollback<UndoLog<K, V>> for SnapshotMap<K, V, M, L>
+where
+    K: Eq + Hash,
+    M: Rollback<UndoLog<K, V>>,
+{
+    fn reverse(&mut self, undo: UndoLog<K, V>) {
+        self.map.reverse(undo)
+    }
+}
+
+impl<K, V> Rollback<UndoLog<K, V>> for FxHashMap<K, V>
+where
+    K: Eq + Hash,
+{
+    fn reverse(&mut self, undo: UndoLog<K, V>) {
+        match undo {
+            UndoLog::Inserted(key) => {
+                self.remove(&key);
+            }
+
+            UndoLog::Overwrite(key, old_value) => {
+                self.insert(key, old_value);
+            }
+
+            UndoLog::Purged => {}
+        }
+    }
+}

crates/tinymist-query/src/analysis.rs

@@ -13,6 +13,9 @@ pub use def_use::*;
 
 #[cfg(test)]
 mod lexical_hierarchy_tests {
+    use def_use::DefUseSnapshot;
+
+    use crate::analysis::def_use;
     use crate::analysis::lexical_hierarchy;
     use crate::prelude::*;
     use crate::tests::*;
@@ -30,4 +33,16 @@ mod lexical_hierarchy_tests {
             assert_snapshot!(JsonRepr::new_redacted(result, &REDACT_LOC));
         });
     }
+
+    #[test]
+    fn def_use() {
+        snapshot_testing("lexical_hierarchy", &|world, path| {
+            let source = get_suitable_source_in_workspace(world, &path).unwrap();
+
+            let result = def_use::get_def_use(source);
+            let result = result.as_ref().map(DefUseSnapshot);
+
+            assert_snapshot!(JsonRepr::new_redacted(result, &REDACT_LOC));
+        });
+    }
 }

crates/tinymist-query/src/analysis/def_use.rs

@@ -1,7 +1,192 @@
+use core::fmt;
+use std::{collections::HashMap, ops::Range};
+
+use serde::Serialize;
 use typst::syntax::Source;
 
-use super::{get_lexical_hierarchy, LexicalScopeKind};
+use crate::adt::snapshot_map::SnapshotMap;
 
-pub fn get_def_use(source: Source) {
-    let _ = get_lexical_hierarchy(source, LexicalScopeKind::DefUse);
+use super::{get_lexical_hierarchy, LexicalHierarchy, LexicalKind, LexicalScopeKind};
+
+pub use typst_ts_core::vector::ir::DefId;
+
+enum Ns {
+    Label,
+    Value,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct IdentRef {
+    name: String,
+    range: Range<usize>,
+}
+
+impl fmt::Display for IdentRef {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}@{:?}", self.name, self.range)
+    }
+}
+
+impl Serialize for IdentRef {
+    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
+        let s = self.to_string();
+        serializer.serialize_str(&s)
+    }
+}
+
+#[derive(Serialize)]
+pub struct IdentDef {
+    name: String,
+    kind: LexicalKind,
+    range: Range<usize>,
+}
+
+#[derive(Default)]
+pub struct DefUseInfo {
+    ident_defs: indexmap::IndexMap<IdentRef, IdentDef>,
+    ident_refs: HashMap<IdentRef, DefId>,
+    undefined_refs: Vec<IdentRef>,
+}
+
+pub struct DefUseSnapshot<'a>(pub &'a DefUseInfo);
+
+impl<'a> Serialize for DefUseSnapshot<'a> {
+    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
+        use serde::ser::SerializeMap;
+        // HashMap<IdentRef, DefId>
+        let references: HashMap<DefId, Vec<IdentRef>> = {
+            let mut map = HashMap::new();
+            for (k, v) in &self.0.ident_refs {
+                map.entry(*v).or_insert_with(Vec::new).push(k.clone());
+            }
+            map
+        };
+
+        #[derive(Serialize)]
+        struct DefUseEntry<'a> {
+            def: &'a IdentDef,
+            refs: &'a Vec<IdentRef>,
+        }
+
+        let mut state = serializer.serialize_map(None)?;
+        for (k, (ident_ref, ident_def)) in self.0.ident_defs.as_slice().iter().enumerate() {
+            let id = DefId(k as u64);
+
+            let empty_ref = Vec::new();
+            let entry = DefUseEntry {
+                def: ident_def,
+                refs: references.get(&id).unwrap_or(&empty_ref),
+            };
+
+            state.serialize_entry(&ident_ref.to_string(), &entry)?;
+        }
+
+        if !self.0.undefined_refs.is_empty() {
+            let entry = DefUseEntry {
+                def: &IdentDef {
+                    name: "<nil>".to_string(),
+                    kind: LexicalKind::Block,
+                    range: 0..0,
+                },
+                refs: &self.0.undefined_refs,
+            };
+            state.serialize_entry("<nil>", &entry)?;
+        }
+
+        state.end()
+    }
+}
+
+pub fn get_def_use(source: Source) -> Option<DefUseInfo> {
+    let e = get_lexical_hierarchy(source, LexicalScopeKind::DefUse)?;
+
+    let mut collector = DefUseCollector {
+        info: DefUseInfo::default(),
+        id_scope: SnapshotMap::default(),
+        label_scope: SnapshotMap::default(),
+    };
+
+    collector.scan(&e);
+    Some(collector.info)
+}
+
+struct DefUseCollector {
+    info: DefUseInfo,
+    label_scope: SnapshotMap<String, DefId>,
+    id_scope: SnapshotMap<String, DefId>,
+}
+
+impl DefUseCollector {
+    fn enter<T>(&mut self, f: impl FnOnce(&mut Self) -> T) -> T {
+        let id_snap = self.id_scope.snapshot();
+        let res = f(self);
+        self.id_scope.rollback_to(id_snap);
+        res
+    }
+
+    fn scan(&mut self, e: &[LexicalHierarchy]) -> Option<()> {
+        for e in e {
+            match e.info.kind {
+                LexicalKind::Heading(..) => unreachable!(),
+                LexicalKind::Label => self.insert(Ns::Label, e),
+                LexicalKind::LabelRef => self.insert_ref(Ns::Label, e),
+                LexicalKind::Function | LexicalKind::Variable => self.insert(Ns::Value, e),
+                LexicalKind::ValRef => self.insert_ref(Ns::Value, e),
+                LexicalKind::Block => {
+                    if let Some(e) = &e.children {
+                        self.enter(|this| this.scan(e.as_slice()))?;
+                    }
+                }
+            }
+        }
+
+        Some(())
+    }
+
+    fn insert(&mut self, label: Ns, e: &LexicalHierarchy) {
+        let snap = match label {
+            Ns::Label => &mut self.label_scope,
+            Ns::Value => &mut self.id_scope,
+        };
+
+        let id_ref = IdentRef {
+            name: e.info.name.clone(),
+            range: e.info.range.clone(),
+        };
+        let (id, old_def) = self.info.ident_defs.insert_full(
+            id_ref.clone(),
+            IdentDef {
+                name: e.info.name.clone(),
+                kind: e.info.kind,
+                range: e.info.range.clone(),
+            },
+        );
+        if let Some(old_def) = old_def {
+            assert_eq!(old_def.kind, e.info.kind);
+        }
+
+        let id = DefId(id as u64);
+        snap.insert(e.info.name.clone(), id);
+    }
+
+    fn insert_ref(&mut self, label: Ns, e: &LexicalHierarchy) {
+        let snap = match label {
+            Ns::Label => &mut self.label_scope,
+            Ns::Value => &mut self.id_scope,
+        };
+
+        let id_ref = IdentRef {
+            name: e.info.name.clone(),
+            range: e.info.range.clone(),
+        };
+
+        match snap.get(&e.info.name) {
+            Some(id) => {
+                self.info.ident_refs.insert(id_ref, *id);
+            }
+            None => {
+                self.info.undefined_refs.push(id_ref);
+            }
+        }
+    }
 }

crates/tinymist-query/src/analysis/lexical_hierarchy.rs

@@ -24,7 +24,7 @@ pub(crate) fn get_lexical_hierarchy(
     worker.stack.push((
         LexicalInfo {
             name: "deadbeef".to_string(),
-            kind: LexicalKind::Namespace(-1),
+            kind: LexicalKind::Heading(-1),
             range: 0..0,
         },
         eco_vec![],
@@ -40,9 +40,9 @@ pub(crate) fn get_lexical_hierarchy(
     res.map(|_| worker.stack.pop().unwrap().1)
 }
 
-#[derive(Debug, Clone, Copy, Hash, Serialize, Deserialize)]
+#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, Serialize, Deserialize)]
 pub(crate) enum LexicalKind {
-    Namespace(i16),
+    Heading(i16),
     ValRef,
     LabelRef,
     Variable,
@@ -56,7 +56,7 @@ impl TryFrom<LexicalKind> for SymbolKind {
 
     fn try_from(value: LexicalKind) -> Result<Self, Self::Error> {
         match value {
-            LexicalKind::Namespace(..) => Ok(SymbolKind::NAMESPACE),
+            LexicalKind::Heading(..) => Ok(SymbolKind::NAMESPACE),
             LexicalKind::Variable => Ok(SymbolKind::VARIABLE),
             LexicalKind::Function => Ok(SymbolKind::FUNCTION),
             LexicalKind::Label => Ok(SymbolKind::CONSTANT),
@@ -218,10 +218,10 @@ impl LexicalHierarchyWorker {
         let checkpoint = self.enter_symbol_context(&node)?;
 
         if let Some(symbol) = own_symbol {
-            if let LexicalKind::Namespace(level) = symbol.kind {
+            if let LexicalKind::Heading(level) = symbol.kind {
                 'heading_break: while let Some((w, _)) = self.stack.last() {
                     match w.kind {
-                        LexicalKind::Namespace(l) if l < level => break 'heading_break,
+                        LexicalKind::Heading(l) if l < level => break 'heading_break,
                         LexicalKind::Block => break 'heading_break,
                         _ if self.stack.len() <= 1 => break 'heading_break,
                         _ => {}
@@ -230,7 +230,7 @@ impl LexicalHierarchyWorker {
                     self.symbreak();
                 }
             }
-            let is_heading = matches!(symbol.kind, LexicalKind::Namespace(..));
+            let is_heading = matches!(symbol.kind, LexicalKind::Heading(..));
 
             self.stack.push((symbol, eco_vec![]));
             let stack_height = self.stack.len();
@@ -416,7 +416,7 @@ impl LexicalHierarchyWorker {
                     return Ok(None);
                 };
                 let kind = match parent.kind() {
-                    SyntaxKind::Heading if self.g.affect_heading() => LexicalKind::Namespace(
+                    SyntaxKind::Heading if self.g.affect_heading() => LexicalKind::Heading(
                         parent.cast::<ast::Heading>().unwrap().depth().get() as i16,
                     ),
                     _ => return Ok(None),

crates/tinymist-query/src/fixtures/lexical_hierarchy/def_use@base.typ.snap

@@ -0,0 +1,17 @@
+---
+source: crates/tinymist-query/src/analysis.rs
+expression: "JsonRepr::new_redacted(result, &REDACT_LOC)"
+input_file: crates/tinymist-query/src/fixtures/lexical_hierarchy/base.typ
+---
+{
+ "x@5..6": {
+  "def": {
+   "kind": "Variable",
+   "name": "x",
+   "range": "5:6"
+  },
+  "refs": [
+   "x@14..15"
+  ]
+ }
+}

crates/tinymist-query/src/fixtures/lexical_hierarchy/def_use@destructing.typ.snap

@@ -0,0 +1,43 @@
+---
+source: crates/tinymist-query/src/analysis.rs
+expression: "JsonRepr::new_redacted(result, &REDACT_LOC)"
+input_file: crates/tinymist-query/src/fixtures/lexical_hierarchy/destructing.typ
+---
+{
+ "a@29..30": {
+  "def": {
+   "kind": "Variable",
+   "name": "a",
+   "range": "29:30"
+  },
+  "refs": []
+ },
+ "a@6..7": {
+  "def": {
+   "kind": "Variable",
+   "name": "a",
+   "range": "6:7"
+  },
+  "refs": [
+   "a@41..42"
+  ]
+ },
+ "b@32..33": {
+  "def": {
+   "kind": "Variable",
+   "name": "b",
+   "range": "32:33"
+  },
+  "refs": []
+ },
+ "b@9..10": {
+  "def": {
+   "kind": "Variable",
+   "name": "b",
+   "range": "9:10"
+  },
+  "refs": [
+   "b@38..39"
+  ]
+ }
+}

crates/tinymist-query/src/fixtures/lexical_hierarchy/def_use@dict.typ.snap

@@ -0,0 +1,26 @@
+---
+source: crates/tinymist-query/src/analysis.rs
+expression: "JsonRepr::new_redacted(result, &REDACT_LOC)"
+input_file: crates/tinymist-query/src/fixtures/lexical_hierarchy/dict.typ
+---
+{
+ "x@18..19": {
+  "def": {
+   "kind": "Variable",
+   "name": "x",
+   "range": "18:19"
+  },
+  "refs": []
+ },
+ "z@5..6": {
+  "def": {
+   "kind": "Variable",
+   "name": "z",
+   "range": "5:6"
+  },
+  "refs": [
+   "z@43..44",
+   "z@30..31"
+  ]
+ }
+}

crates/tinymist-query/src/fixtures/lexical_hierarchy/def_use@func.typ.snap

@@ -0,0 +1,33 @@
+---
+source: crates/tinymist-query/src/analysis.rs
+expression: "JsonRepr::new_redacted(result, &REDACT_LOC)"
+input_file: crates/tinymist-query/src/fixtures/lexical_hierarchy/func.typ
+---
+{
+ "a@20..21": {
+  "def": {
+   "kind": "Variable",
+   "name": "a",
+   "range": "20:21"
+  },
+  "refs": [
+   "a@25..26"
+  ]
+ },
+ "f@18..19": {
+  "def": {
+   "kind": "Function",
+   "name": "f",
+   "range": "18:19"
+  },
+  "refs": []
+ },
+ "x@5..6": {
+  "def": {
+   "kind": "Variable",
+   "name": "x",
+   "range": "5:6"
+  },
+  "refs": []
+ }
+}

crates/tinymist-query/src/fixtures/lexical_hierarchy/def_use@redefine.typ.snap

@@ -0,0 +1,25 @@
+---
+source: crates/tinymist-query/src/analysis.rs
+expression: "JsonRepr::new_redacted(result, &REDACT_LOC)"
+input_file: crates/tinymist-query/src/fixtures/lexical_hierarchy/redefine.typ
+---
+{
+ "x@18..19": {
+  "def": {
+   "kind": "Variable",
+   "name": "x",
+   "range": "18:19"
+  },
+  "refs": []
+ },
+ "x@5..6": {
+  "def": {
+   "kind": "Variable",
+   "name": "x",
+   "range": "5:6"
+  },
+  "refs": [
+   "x@22..23"
+  ]
+ }
+}

crates/tinymist-query/src/folding_range.rs

@@ -80,7 +80,7 @@ fn calc_folding_range(
             last_loc
         };
 
-        if matches!(e.info.kind, LexicalKind::Namespace(..)) {
+        if matches!(e.info.kind, LexicalKind::Heading(..)) {
             range.end_line = range.end_line.max(if is_not_last_range {
                 next_start.0.saturating_sub(1)
             } else {

crates/tinymist-query/src/lib.rs

@@ -1,3 +1,4 @@
+mod adt;
 pub mod analysis;
 
 pub(crate) mod diagnostics;