Optimized labels & introspector (#2801)

Co-authored-by: Laurenz <laurmaedje@gmail.com>

14 files changed, 109 insertions, 81 deletions

diff --git a/Cargo.lock b/Cargo.lock
index dfd01d8b..2dcef771 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2644,6 +2644,7 @@ dependencies = [
  "ciborium",
  "comemo",
  "csv",
+ "dashmap",
  "ecow",
  "fontdb",
  "hayagriva",
diff --git a/Cargo.toml b/Cargo.toml
index 5c1d21a1..2aca6f18 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -38,6 +38,7 @@ clap_mangen = "0.2.10"
 codespan-reporting = "0.11"
 comemo = "0.3.1"
 csv = "1"
+dashmap = "5.5"
 dirs = "5"
 ecow = { version = "0.2", features = ["serde"] }
 env_proxy = "0.4"
diff --git a/crates/typst-cli/src/query.rs b/crates/typst-cli/src/query.rs
index a84cef79..dadcd978 100644
--- a/crates/typst-cli/src/query.rs
+++ b/crates/typst-cli/src/query.rs
@@ -4,7 +4,6 @@ use serde::Serialize;
 use typst::diag::{bail, StrResult};
 use typst::eval::{eval_string, EvalMode, Tracer};
 use typst::foundations::{Content, IntoValue, LocatableSelector, Scope};
-use typst::introspection::Introspector;
 use typst::model::Document;
 use typst::syntax::Span;
 use typst::World;
@@ -76,7 +75,8 @@ fn retrieve(
     })?
     .cast::<LocatableSelector>()?;
 
-    Ok(Introspector::new(&document.pages)
+    Ok(document
+        .introspector
         .query(&selector.0)
         .into_iter()
         .map(|x| x.into_inner())
diff --git a/crates/typst-ide/src/analyze.rs b/crates/typst-ide/src/analyze.rs
index 8c5117f7..e78f0c12 100644
--- a/crates/typst-ide/src/analyze.rs
+++ b/crates/typst-ide/src/analyze.rs
@@ -4,8 +4,7 @@ use typst::engine::{Engine, Route};
 use typst::eval::{Tracer, Vm};
 use typst::foundations::{Label, Scopes, Value};
 use typst::introspection::{Introspector, Locator};
-use typst::layout::Frame;
-use typst::model::BibliographyElem;
+use typst::model::{BibliographyElem, Document};
 use typst::syntax::{ast, LinkedNode, Span, SyntaxKind};
 use typst::World;
 
@@ -75,12 +74,11 @@ pub fn analyze_import(world: &dyn World, source: &LinkedNode) -> Option<Value> {
 /// - All labels and descriptions for them, if available
 /// - A split offset: All labels before this offset belong to nodes, all after
 ///   belong to a bibliography.
-pub fn analyze_labels(frames: &[Frame]) -> (Vec<(Label, Option<EcoString>)>, usize) {
+pub fn analyze_labels(document: &Document) -> (Vec<(Label, Option<EcoString>)>, usize) {
     let mut output = vec![];
-    let introspector = Introspector::new(frames);
 
     // Labels in the document.
-    for elem in introspector.all() {
+    for elem in document.introspector.all() {
         let Some(label) = elem.label() else { continue };
         let details = elem
             .get_by_name("caption")
@@ -98,7 +96,7 @@ pub fn analyze_labels(frames: &[Frame]) -> (Vec<(Label, Option<EcoString>)>, usi
     let split = output.len();
 
     // Bibliography keys.
-    for (key, detail) in BibliographyElem::keys(introspector.track()) {
+    for (key, detail) in BibliographyElem::keys(document.introspector.track()) {
         output.push((Label::new(&key), detail));
     }
 
diff --git a/crates/typst-ide/src/complete.rs b/crates/typst-ide/src/complete.rs
index 6e94ee6a..8acf1540 100644
--- a/crates/typst-ide/src/complete.rs
+++ b/crates/typst-ide/src/complete.rs
@@ -8,7 +8,7 @@ use typst::foundations::{
     fields_on, format_str, mutable_methods_on, repr, AutoValue, CastInfo, Func, Label,
     NoneValue, Repr, Scope, Type, Value,
 };
-use typst::layout::Frame;
+use typst::model::Document;
 use typst::syntax::{
     ast, is_id_continue, is_id_start, is_ident, LinkedNode, Source, SyntaxKind,
 };
@@ -27,14 +27,18 @@ use crate::{plain_docs_sentence, summarize_font_family};
 ///
 /// When `explicit` is `true`, the user requested the completion by pressing
 /// control and space or something similar.
+///
+/// Passing a `document` (from a previous compilation) is optional, but enhances
+/// the autocompletions. Label completions, for instance, are only generated
+/// when the document is available.
 pub fn autocomplete(
     world: &dyn World,
-    frames: &[Frame],
+    document: Option<&Document>,
     source: &Source,
     cursor: usize,
     explicit: bool,
 ) -> Option<(usize, Vec<Completion>)> {
-    let mut ctx = CompletionContext::new(world, frames, source, cursor, explicit)?;
+    let mut ctx = CompletionContext::new(world, document, source, cursor, explicit)?;
 
     let _ = complete_comments(&mut ctx)
         || complete_field_accesses(&mut ctx)
@@ -966,7 +970,7 @@ fn code_completions(ctx: &mut CompletionContext, hash: bool) {
 /// Context for autocompletion.
 struct CompletionContext<'a> {
     world: &'a (dyn World + 'a),
-    frames: &'a [Frame],
+    document: Option<&'a Document>,
     global: &'a Scope,
     math: &'a Scope,
     text: &'a str,
@@ -984,7 +988,7 @@ impl<'a> CompletionContext<'a> {
     /// Create a new autocompletion context.
     fn new(
         world: &'a (dyn World + 'a),
-        frames: &'a [Frame],
+        document: Option<&'a Document>,
         source: &'a Source,
         cursor: usize,
         explicit: bool,
@@ -994,7 +998,7 @@ impl<'a> CompletionContext<'a> {
         let leaf = LinkedNode::new(source.root()).leaf_at(cursor)?;
         Some(Self {
             world,
-            frames,
+            document,
             global: library.global.scope(),
             math: library.math.scope(),
             text,
@@ -1094,7 +1098,8 @@ impl<'a> CompletionContext<'a> {
 
     /// Add completions for labels and references.
     fn label_completions(&mut self) {
-        let (labels, split) = analyze_labels(self.frames);
+        let Some(document) = self.document else { return };
+        let (labels, split) = analyze_labels(document);
 
         let head = &self.text[..self.from];
         let at = head.ends_with('@');
diff --git a/crates/typst-ide/src/jump.rs b/crates/typst-ide/src/jump.rs
index 700f475f..8612ffc3 100644
--- a/crates/typst-ide/src/jump.rs
+++ b/crates/typst-ide/src/jump.rs
@@ -1,9 +1,9 @@
 use std::num::NonZeroUsize;
 
 use ecow::EcoString;
-use typst::introspection::{Introspector, Meta};
+use typst::introspection::Meta;
 use typst::layout::{Frame, FrameItem, Point, Position, Size};
-use typst::model::Destination;
+use typst::model::{Destination, Document};
 use typst::syntax::{FileId, LinkedNode, Source, Span, SyntaxKind};
 use typst::visualize::Geometry;
 use typst::World;
@@ -31,12 +31,10 @@ impl Jump {
 /// Determine where to jump to based on a click in a frame.
 pub fn jump_from_click(
     world: &dyn World,
-    frames: &[Frame],
+    document: &Document,
     frame: &Frame,
     click: Point,
 ) -> Option<Jump> {
-    let mut introspector = None;
-
     // Try to find a link first.
     for (pos, item) in frame.items() {
         if let FrameItem::Meta(Meta::Link(dest), size) = item {
@@ -44,11 +42,9 @@ pub fn jump_from_click(
                 return Some(match dest {
                     Destination::Url(url) => Jump::Url(url.clone()),
                     Destination::Position(pos) => Jump::Position(*pos),
-                    Destination::Location(loc) => Jump::Position(
-                        introspector
-                            .get_or_insert_with(|| Introspector::new(frames))
-                            .position(*loc),
-                    ),
+                    Destination::Location(loc) => {
+                        Jump::Position(document.introspector.position(*loc))
+                    }
                 });
             }
         }
@@ -60,7 +56,7 @@ pub fn jump_from_click(
             FrameItem::Group(group) => {
                 // TODO: Handle transformation.
                 if let Some(span) =
-                    jump_from_click(world, frames, &group.frame, click - pos)
+                    jump_from_click(world, document, &group.frame, click - pos)
                 {
                     return Some(span);
                 }
diff --git a/crates/typst-ide/src/tooltip.rs b/crates/typst-ide/src/tooltip.rs
index 4f079166..67614b9e 100644
--- a/crates/typst-ide/src/tooltip.rs
+++ b/crates/typst-ide/src/tooltip.rs
@@ -4,7 +4,8 @@ use ecow::{eco_format, EcoString};
 use if_chain::if_chain;
 use typst::eval::{CapturesVisitor, Tracer};
 use typst::foundations::{repr, CastInfo, Repr, Value};
-use typst::layout::{Frame, Length};
+use typst::layout::Length;
+use typst::model::Document;
 use typst::syntax::{ast, LinkedNode, Source, SyntaxKind};
 use typst::util::{round_2, Numeric};
 use typst::World;
@@ -13,9 +14,13 @@ use crate::analyze::{analyze_expr, analyze_labels};
 use crate::{plain_docs_sentence, summarize_font_family};
 
 /// Describe the item under the cursor.
+///
+/// Passing a `document` (from a previous compilation) is optional, but enhances
+/// the autocompletions. Label completions, for instance, are only generated
+/// when the document is available.
 pub fn tooltip(
     world: &dyn World,
-    frames: &[Frame],
+    document: Option<&Document>,
     source: &Source,
     cursor: usize,
 ) -> Option<Tooltip> {
@@ -26,7 +31,7 @@ pub fn tooltip(
 
     named_param_tooltip(world, &leaf)
         .or_else(|| font_tooltip(world, &leaf))
-        .or_else(|| label_tooltip(frames, &leaf))
+        .or_else(|| document.and_then(|doc| label_tooltip(doc, &leaf)))
         .or_else(|| expr_tooltip(world, &leaf))
         .or_else(|| closure_tooltip(&leaf))
 }
@@ -145,14 +150,14 @@ fn length_tooltip(length: Length) -> Option<Tooltip> {
 }
 
 /// Tooltip for a hovered reference or label.
-fn label_tooltip(frames: &[Frame], leaf: &LinkedNode) -> Option<Tooltip> {
+fn label_tooltip(document: &Document, leaf: &LinkedNode) -> Option<Tooltip> {
     let target = match leaf.kind() {
         SyntaxKind::RefMarker => leaf.text().trim_start_matches('@'),
         SyntaxKind::Label => leaf.text().trim_start_matches('<').trim_end_matches('>'),
         _ => return None,
     };
 
-    for (label, detail) in analyze_labels(frames).0 {
+    for (label, detail) in analyze_labels(document).0 {
         if label.as_str() == target {
             return Some(Tooltip::Text(detail?));
         }
diff --git a/crates/typst-pdf/src/lib.rs b/crates/typst-pdf/src/lib.rs
index 0c82cd90..3982d156 100644
--- a/crates/typst-pdf/src/lib.rs
+++ b/crates/typst-pdf/src/lib.rs
@@ -19,7 +19,6 @@ use ecow::{eco_format, EcoString};
 use pdf_writer::types::Direction;
 use pdf_writer::{Finish, Name, Pdf, Ref, TextStr};
 use typst::foundations::Datetime;
-use typst::introspection::Introspector;
 use typst::layout::{Abs, Dir, Em, Transform};
 use typst::model::Document;
 use typst::text::{Font, Lang};
@@ -70,10 +69,6 @@ pub fn pdf(
 struct PdfContext<'a> {
     /// The document that we're currently exporting.
     document: &'a Document,
-    /// An introspector for the document, used to resolve locations links and
-    /// the document outline.
-    introspector: Introspector,
-
     /// The writer we are writing the PDF into.
     pdf: Pdf,
     /// Content of exported pages.
@@ -128,7 +123,6 @@ impl<'a> PdfContext<'a> {
         let page_tree_ref = alloc.bump();
         Self {
             document,
-            introspector: Introspector::new(&document.pages),
             pdf: Pdf::new(),
             pages: vec![],
             glyph_sets: HashMap::new(),
diff --git a/crates/typst-pdf/src/outline.rs b/crates/typst-pdf/src/outline.rs
index aafaa5b9..9c9ef413 100644
--- a/crates/typst-pdf/src/outline.rs
+++ b/crates/typst-pdf/src/outline.rs
@@ -18,7 +18,7 @@ pub(crate) fn write_outline(ctx: &mut PdfContext) -> Option<Ref> {
     // Therefore, its next descendant must be added at its level, which is
     // enforced in the manner shown below.
     let mut last_skipped_level = None;
-    for heading in ctx.introspector.query(&HeadingElem::elem().select()).iter() {
+    for heading in ctx.document.introspector.query(&HeadingElem::elem().select()).iter() {
         let leaf = HeadingNode::leaf((**heading).clone());
 
         if leaf.bookmarked {
@@ -163,7 +163,7 @@ fn write_outline_item(
     outline.title(TextStr(body.plain_text().trim()));
 
     let loc = node.element.location().unwrap();
-    let pos = ctx.introspector.position(loc);
+    let pos = ctx.document.introspector.position(loc);
     let index = pos.page.get() - 1;
     if let Some(page) = ctx.pages.get(index) {
         let y = (pos.point.y - Abs::pt(10.0)).max(Abs::zero());
diff --git a/crates/typst-pdf/src/page.rs b/crates/typst-pdf/src/page.rs
index 1bbef7af..e4322f5f 100644
--- a/crates/typst-pdf/src/page.rs
+++ b/crates/typst-pdf/src/page.rs
@@ -180,7 +180,7 @@ fn write_page(ctx: &mut PdfContext, i: usize) {
                 continue;
             }
             Destination::Position(pos) => *pos,
-            Destination::Location(loc) => ctx.introspector.position(*loc),
+            Destination::Location(loc) => ctx.document.introspector.position(*loc),
         };
 
         let index = pos.page.get() - 1;
diff --git a/crates/typst/Cargo.toml b/crates/typst/Cargo.toml
index 6334b67e..93cdfd8a 100644
--- a/crates/typst/Cargo.toml
+++ b/crates/typst/Cargo.toml
@@ -24,6 +24,7 @@ chinese-number = { workspace = true }
 ciborium = { workspace = true }
 comemo = { workspace = true }
 csv = { workspace = true }
+dashmap = { workspace = true }
 ecow = { workspace = true}
 fontdb = { workspace = true }
 hayagriva = { workspace = true }
diff --git a/crates/typst/src/introspection/introspector.rs b/crates/typst/src/introspection/introspector.rs
index df454b84..ff9d6c64 100644
--- a/crates/typst/src/introspection/introspector.rs
+++ b/crates/typst/src/introspection/introspector.rs
@@ -1,11 +1,13 @@
-use std::cell::RefCell;
 use std::collections::{BTreeSet, HashMap};
+use std::fmt::{self, Debug, Formatter};
 use std::hash::Hash;
 use std::num::NonZeroUsize;
 
 use comemo::Prehashed;
+use dashmap::DashMap;
 use ecow::{eco_format, EcoVec};
 use indexmap::IndexMap;
+use smallvec::SmallVec;
 
 use crate::diag::{bail, StrResult};
 use crate::foundations::{Content, Label, Repr, Selector};
@@ -15,35 +17,38 @@ use crate::model::Numbering;
 use crate::util::NonZeroExt;
 
 /// Can be queried for elements and their positions.
+#[derive(Clone)]
 pub struct Introspector {
     /// The number of pages in the document.
     pages: usize,
     /// All introspectable elements.
     elems: IndexMap<Location, (Prehashed<Content>, Position)>,
+    /// Maps labels to their indices in the element list. We use a smallvec such
+    /// that if the label is unique, we don't need to allocate.
+    labels: HashMap<Label, SmallVec<[usize; 1]>>,
     /// The page numberings, indexed by page number minus 1.
     page_numberings: Vec<Option<Numbering>>,
     /// Caches queries done on the introspector. This is important because
     /// even if all top-level queries are distinct, they often have shared
     /// subqueries. Example: Individual counter queries with `before` that
     /// all depend on a global counter query.
-    queries: RefCell<HashMap<u128, EcoVec<Prehashed<Content>>>>,
+    queries: DashMap<u128, EcoVec<Prehashed<Content>>>,
 }
 
 impl Introspector {
-    /// Create a new introspector.
-    #[tracing::instrument(skip(frames))]
-    pub fn new(frames: &[Frame]) -> Self {
-        let mut introspector = Self {
-            pages: frames.len(),
-            elems: IndexMap::new(),
-            page_numberings: vec![],
-            queries: RefCell::default(),
-        };
+    /// Applies new frames in-place, reusing the existing allocations.
+    #[tracing::instrument(skip_all)]
+    pub fn rebuild(&mut self, frames: &[Frame]) {
+        self.pages = frames.len();
+        self.elems.clear();
+        self.labels.clear();
+        self.page_numberings.clear();
+        self.queries.clear();
+
         for (i, frame) in frames.iter().enumerate() {
             let page = NonZeroUsize::new(1 + i).unwrap();
-            introspector.extract(frame, page, Transform::identity());
+            self.extract(frame, page, Transform::identity());
         }
-        introspector
     }
 
     /// Extract metadata from a frame.
@@ -61,11 +66,17 @@ impl Introspector {
                     if !self.elems.contains_key(&content.location().unwrap()) =>
                 {
                     let pos = pos.transform(ts);
+                    let content = Prehashed::new(content.clone());
                     let ret = self.elems.insert(
                         content.location().unwrap(),
-                        (Prehashed::new(content.clone()), Position { page, point: pos }),
+                        (content.clone(), Position { page, point: pos }),
                     );
                     assert!(ret.is_none(), "duplicate locations");
+
+                    // Build the label cache.
+                    if let Some(label) = content.label() {
+                        self.labels.entry(label).or_default().push(self.elems.len() - 1);
+                    }
                 }
                 FrameItem::Meta(Meta::PageNumbering(numbering), _) => {
                     self.page_numberings.push(numbering.clone());
@@ -107,15 +118,19 @@ impl Introspector {
     /// Query for all matching elements.
     pub fn query(&self, selector: &Selector) -> EcoVec<Prehashed<Content>> {
         let hash = crate::util::hash128(selector);
-        if let Some(output) = self.queries.borrow().get(&hash) {
+        if let Some(output) = self.queries.get(&hash) {
             return output.clone();
         }
 
         let output = match selector {
-            Selector::Elem(..)
-            | Selector::Label(_)
-            | Selector::Regex(_)
-            | Selector::Can(_) => {
+            Selector::Label(label) => self
+                .labels
+                .get(label)
+                .map(|indices| {
+                    indices.iter().map(|&index| self.elems[index].0.clone()).collect()
+                })
+                .unwrap_or_default(),
+            Selector::Elem(..) | Selector::Regex(_) | Selector::Can(_) => {
                 self.all().filter(|elem| selector.matches(elem)).cloned().collect()
             }
             Selector::Location(location) => {
@@ -182,7 +197,7 @@ impl Introspector {
                 .collect(),
         };
 
-        self.queries.borrow_mut().insert(hash, output.clone());
+        self.queries.insert(hash, output.clone());
         output
     }
 
@@ -196,16 +211,15 @@ impl Introspector {
 
     /// Query for a unique element with the label.
     pub fn query_label(&self, label: Label) -> StrResult<&Prehashed<Content>> {
-        let mut found = None;
-        for elem in self.all().filter(|elem| elem.label() == Some(label)) {
-            if found.is_some() {
-                bail!("label `{}` occurs multiple times in the document", label.repr());
-            }
-            found = Some(elem);
-        }
-        found.ok_or_else(|| {
+        let indices = self.labels.get(&label).ok_or_else(|| {
             eco_format!("label `{}` does not exist in the document", label.repr())
-        })
+        })?;
+
+        if indices.len() > 1 {
+            bail!("label `{}` occurs multiple times in the document", label.repr());
+        }
+
+        Ok(&self.elems[indices[0]].0)
     }
 
     /// The total number pages.
@@ -237,6 +251,18 @@ impl Introspector {
 
 impl Default for Introspector {
     fn default() -> Self {
-        Self::new(&[])
+        Self {
+            pages: 0,
+            elems: IndexMap::new(),
+            labels: HashMap::new(),
+            page_numberings: vec![],
+            queries: DashMap::new(),
+        }
+    }
+}
+
+impl Debug for Introspector {
+    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+        f.pad("Introspector(..)")
     }
 }
diff --git a/crates/typst/src/lib.rs b/crates/typst/src/lib.rs
index 33e688b7..303d5183 100644
--- a/crates/typst/src/lib.rs
+++ b/crates/typst/src/lib.rs
@@ -111,8 +111,7 @@ fn typeset(
     let styles = StyleChain::new(&library.styles);
 
     let mut iter = 0;
-    let mut document;
-    let mut introspector = Introspector::new(&[]);
+    let mut document = Document::default();
 
     // Relayout until all introspections stabilize.
     // If that doesn't happen within five attempts, we give up.
@@ -129,16 +128,15 @@ fn typeset(
             route: Route::default(),
             tracer: tracer.track_mut(),
             locator: &mut locator,
-            introspector: introspector.track_with(&constraint),
+            introspector: document.introspector.track_with(&constraint),
         };
 
         // Layout!
         document = content.layout_root(&mut engine, styles)?;
-
-        introspector = Introspector::new(&document.pages);
+        document.introspector.rebuild(&document.pages);
         iter += 1;
 
-        if introspector.validate(&constraint) {
+        if document.introspector.validate(&constraint) {
             break;
         }
 
diff --git a/crates/typst/src/model/document.rs b/crates/typst/src/model/document.rs
index 1dc241c0..39fb2ac3 100644
--- a/crates/typst/src/model/document.rs
+++ b/crates/typst/src/model/document.rs
@@ -5,7 +5,7 @@ use crate::engine::Engine;
 use crate::foundations::{
     cast, elem, Args, Array, Construct, Content, Datetime, Smart, StyleChain, Value,
 };
-use crate::introspection::ManualPageCounter;
+use crate::introspection::{Introspector, ManualPageCounter};
 use crate::layout::{Frame, LayoutRoot, PageElem};
 
 /// The root element of a document and its metadata.
@@ -110,6 +110,7 @@ impl LayoutRoot for DocumentElem {
             author: self.author(styles).0,
             keywords: self.keywords(styles).0,
             date: self.date(styles),
+            introspector: Introspector::default(),
         })
     }
 }
@@ -137,7 +138,7 @@ cast! {
 }
 
 /// A finished document with metadata and page frames.
-#[derive(Debug, Default, Clone, Hash)]
+#[derive(Debug, Default, Clone)]
 pub struct Document {
     /// The page frames.
     pub pages: Vec<Frame>,
@@ -149,6 +150,8 @@ pub struct Document {
     pub keywords: Vec<EcoString>,
     /// The document's creation date.
     pub date: Smart<Option<Datetime>>,
+    /// Provides the ability to execute queries on the document.
+    pub introspector: Introspector,
 }
 
 #[cfg(test)]
@@ -156,8 +159,8 @@ mod tests {
     use super::*;
 
     #[test]
-    fn test_document_is_send() {
-        fn ensure_send<T: Send>() {}
-        ensure_send::<Document>();
+    fn test_document_is_send_and_sync() {
+        fn ensure_send_and_sync<T: Send + Sync>() {}
+        ensure_send_and_sync::<Document>();
     }
 }