More efficient introspection

Switches from a mutable locator to one based on tracked chains and optimizes query performance.

library/src/layout/mod.rs

@@ -75,11 +75,12 @@ impl LayoutRoot for Content {
             content: &Content,
             world: Tracked<dyn World + '_>,
             tracer: TrackedMut<Tracer>,
-            provider: TrackedMut<StabilityProvider>,
+            locator: Tracked<Locator>,
             introspector: Tracked<Introspector>,
             styles: StyleChain,
         ) -> SourceResult<Document> {
-            let mut vt = Vt { world, tracer, provider, introspector };
+            let mut locator = Locator::chained(locator);
+            let mut vt = Vt { world, tracer, locator: &mut locator, introspector };
             let scratch = Scratch::default();
             let (realized, styles) = realize_root(&mut vt, &scratch, content, styles)?;
             realized
@@ -94,7 +95,7 @@ impl LayoutRoot for Content {
             self,
             vt.world,
             TrackedMut::reborrow_mut(&mut vt.tracer),
-            TrackedMut::reborrow_mut(&mut vt.provider),
+            vt.locator.track(),
             vt.introspector,
             styles,
         )
@@ -121,10 +122,14 @@ pub trait Layout {
         styles: StyleChain,
         regions: Regions,
     ) -> SourceResult<Fragment> {
-        vt.provider.save();
+        let mut locator = Locator::chained(vt.locator.track());
-        let result = self.layout(vt, styles, regions);
+        let mut vt = Vt {
-        vt.provider.restore();
+            world: vt.world,
-        result
+            tracer: TrackedMut::reborrow_mut(&mut vt.tracer),
+            locator: &mut locator,
+            introspector: vt.introspector,
+        };
+        self.layout(&mut vt, styles, regions)
     }
 }
 
@@ -141,12 +146,13 @@ impl Layout for Content {
             content: &Content,
             world: Tracked<dyn World + '_>,
             tracer: TrackedMut<Tracer>,
-            provider: TrackedMut<StabilityProvider>,
+            locator: Tracked<Locator>,
             introspector: Tracked<Introspector>,
             styles: StyleChain,
             regions: Regions,
         ) -> SourceResult<Fragment> {
-            let mut vt = Vt { world, tracer, provider, introspector };
+            let mut locator = Locator::chained(locator);
+            let mut vt = Vt { world, tracer, locator: &mut locator, introspector };
             let scratch = Scratch::default();
             let (realized, styles) = realize_block(&mut vt, &scratch, content, styles)?;
             realized
@@ -157,15 +163,18 @@ impl Layout for Content {
 
         tracing::info!("Layouting `Content`");
 
-        cached(
+        let fragment = cached(
             self,
             vt.world,
             TrackedMut::reborrow_mut(&mut vt.tracer),
-            TrackedMut::reborrow_mut(&mut vt.provider),
+            vt.locator.track(),
             vt.introspector,
             styles,
             regions,
-        )
+        )?;
+
+        vt.locator.visit_frames(&fragment);
+        Ok(fragment)
     }
 }
 

library/src/layout/par.rs

@@ -141,14 +141,15 @@ impl ParElem {
             par: &ParElem,
             world: Tracked<dyn World + '_>,
             tracer: TrackedMut<Tracer>,
-            provider: TrackedMut<StabilityProvider>,
+            locator: Tracked<Locator>,
             introspector: Tracked<Introspector>,
             styles: StyleChain,
             consecutive: bool,
             region: Size,
             expand: bool,
         ) -> SourceResult<Fragment> {
-            let mut vt = Vt { world, tracer, provider, introspector };
+            let mut locator = Locator::chained(locator);
+            let mut vt = Vt { world, tracer, locator: &mut locator, introspector };
             let children = par.children();
 
             // Collect all text into one string for BiDi analysis.
@@ -166,17 +167,20 @@ impl ParElem {
             finalize(&mut vt, &p, &lines, region, expand)
         }
 
-        cached(
+        let fragment = cached(
             self,
             vt.world,
             TrackedMut::reborrow_mut(&mut vt.tracer),
-            TrackedMut::reborrow_mut(&mut vt.provider),
+            vt.locator.track(),
             vt.introspector,
             styles,
             consecutive,
             region,
             expand,
-        )
+        )?;
+
+        vt.locator.visit_frames(&fragment);
+        Ok(fragment)
     }
 }
 

library/src/meta/counter.rs

@@ -396,7 +396,7 @@ impl Counter {
         self.sequence_impl(
             vt.world,
             TrackedMut::reborrow_mut(&mut vt.tracer),
-            TrackedMut::reborrow_mut(&mut vt.provider),
+            vt.locator.track(),
             vt.introspector,
         )
     }
@@ -407,10 +407,11 @@ impl Counter {
         &self,
         world: Tracked<dyn World + '_>,
         tracer: TrackedMut<Tracer>,
-        provider: TrackedMut<StabilityProvider>,
+        locator: Tracked<Locator>,
         introspector: Tracked<Introspector>,
     ) -> SourceResult<EcoVec<(CounterState, NonZeroUsize)>> {
-        let mut vt = Vt { world, tracer, provider, introspector };
+        let mut locator = Locator::chained(locator);
+        let mut vt = Vt { world, tracer, locator: &mut locator, introspector };
         let mut state = CounterState(match &self.0 {
             // special case, because pages always start at one.
             CounterKey::Page => smallvec![1],

library/src/meta/query.rs

@@ -122,7 +122,12 @@ pub fn query(
     location: Location,
 ) -> Value {
     let _ = location;
-    vm.vt.introspector.query(&target.0).into()
+    let vec = vm.vt.introspector.query(&target.0);
+    Value::Array(
+        vec.into_iter()
+            .map(|elem| Value::Content(elem.into_inner()))
+            .collect(),
+    )
 }
 
 /// Turns a value into a selector. The following values are accepted:

library/src/meta/reference.rs

@@ -130,7 +130,7 @@ impl Synthesize for RefElem {
         let target = self.target();
         if vt.introspector.init() && !BibliographyElem::has(vt, &target.0) {
             if let Ok(elem) = vt.introspector.query_label(&target) {
-                self.push_element(Some(elem));
+                self.push_element(Some(elem.into_inner()));
                 return Ok(());
             }
         }
@@ -173,7 +173,7 @@ impl Show for RefElem {
         let supplement = match self.supplement(styles) {
             Smart::Auto | Smart::Custom(None) => None,
             Smart::Custom(Some(supplement)) => {
-                Some(supplement.resolve(vt, [elem.clone().into()])?)
+                Some(supplement.resolve(vt, [(*elem).clone().into()])?)
             }
         };
 

library/src/meta/state.rs

@@ -311,7 +311,7 @@ impl State {
         self.sequence_impl(
             vt.world,
             TrackedMut::reborrow_mut(&mut vt.tracer),
-            TrackedMut::reborrow_mut(&mut vt.provider),
+            vt.locator.track(),
             vt.introspector,
         )
     }
@@ -322,10 +322,11 @@ impl State {
         &self,
         world: Tracked<dyn World + '_>,
         tracer: TrackedMut<Tracer>,
-        provider: TrackedMut<StabilityProvider>,
+        locator: Tracked<Locator>,
         introspector: Tracked<Introspector>,
     ) -> SourceResult<EcoVec<Value>> {
-        let mut vt = Vt { world, tracer, provider, introspector };
+        let mut locator = Locator::chained(locator);
+        let mut vt = Vt { world, tracer, locator: &mut locator, introspector };
         let mut state = self.init.clone();
         let mut stops = eco_vec![state.clone()];
 

library/src/prelude.rs

@@ -23,9 +23,9 @@ pub use typst::geom::*;
 #[doc(no_inline)]
 pub use typst::model::{
     element, Behave, Behaviour, Construct, Content, ElemFunc, Element, Finalize, Fold,
-    Introspector, Label, Locatable, LocatableSelector, Location, MetaElem, PlainText,
+    Introspector, Label, Locatable, LocatableSelector, Location, Locator, MetaElem,
-    Resolve, Selector, Set, Show, StabilityProvider, StyleChain, StyleVec, Styles,
+    PlainText, Resolve, Selector, Set, Show, StyleChain, StyleVec, Styles, Synthesize,
-    Synthesize, Unlabellable, Vt,
+    Unlabellable, Vt,
 };
 #[doc(no_inline)]
 pub use typst::syntax::{Span, Spanned};

src/eval/func.rs

@@ -12,7 +12,7 @@ use super::{
     cast_to_value, Args, CastInfo, Eval, Flow, Route, Scope, Scopes, Tracer, Value, Vm,
 };
 use crate::diag::{bail, SourceResult, StrResult};
-use crate::model::{ElemFunc, Introspector, StabilityProvider, Vt};
+use crate::model::{ElemFunc, Introspector, Locator, Vt};
 use crate::syntax::ast::{self, AstNode, Expr, Ident};
 use crate::syntax::{SourceId, Span, SyntaxNode};
 use crate::World;
@@ -104,7 +104,7 @@ impl Func {
                     vm.world(),
                     route,
                     TrackedMut::reborrow_mut(&mut vm.vt.tracer),
-                    TrackedMut::reborrow_mut(&mut vm.vt.provider),
+                    vm.vt.locator.track(),
                     vm.vt.introspector,
                     vm.depth + 1,
                     args,
@@ -127,7 +127,14 @@ impl Func {
         let route = Route::default();
         let id = SourceId::detached();
         let scopes = Scopes::new(None);
-        let mut vm = Vm::new(vt.reborrow_mut(), route.track(), id, scopes);
+        let mut locator = Locator::chained(vt.locator.track());
+        let vt = Vt {
+            world: vt.world,
+            tracer: TrackedMut::reborrow_mut(&mut vt.tracer),
+            locator: &mut locator,
+            introspector: vt.introspector,
+        };
+        let mut vm = Vm::new(vt, route.track(), id, scopes);
         let args = Args::new(self.span(), args);
         self.call_vm(&mut vm, args)
     }
@@ -320,7 +327,7 @@ impl Closure {
         world: Tracked<dyn World + '_>,
         route: Tracked<Route>,
         tracer: TrackedMut<Tracer>,
-        provider: TrackedMut<StabilityProvider>,
+        locator: Tracked<Locator>,
         introspector: Tracked<Introspector>,
         depth: usize,
         mut args: Args,
@@ -335,8 +342,11 @@ impl Closure {
         let mut scopes = Scopes::new(None);
         scopes.top = closure.captured.clone();
 
-        // Evaluate the body.
+        // Prepare VT.
-        let vt = Vt { world, tracer, provider, introspector };
+        let mut locator = Locator::chained(locator);
+        let vt = Vt { world, tracer, locator: &mut locator, introspector };
+
+        // Prepare VM.
         let mut vm = Vm::new(vt, route, closure.location, scopes);
         vm.depth = depth;
 

src/eval/mod.rs

@@ -48,9 +48,8 @@ use unicode_segmentation::UnicodeSegmentation;
 use crate::diag::{
     bail, error, At, SourceError, SourceResult, StrResult, Trace, Tracepoint,
 };
-use crate::model::ShowableSelector;
 use crate::model::{
-    Content, Introspector, Label, Recipe, StabilityProvider, Styles, Transform,
+    Content, Introspector, Label, Locator, Recipe, ShowableSelector, Styles, Transform,
     Unlabellable, Vt,
 };
 use crate::syntax::ast::AstNode;
@@ -83,23 +82,26 @@ pub fn eval(
     let library = world.library();
     set_lang_items(library.items.clone());
 
-    // Evaluate the module.
+    // Prepare VT.
-    let route = Route::insert(route, id);
+    let mut locator = Locator::default();
-    let scopes = Scopes::new(Some(library));
+    let introspector = Introspector::default();
-    let mut provider = StabilityProvider::new();
-    let introspector = Introspector::new(&[]);
     let vt = Vt {
         world,
         tracer,
-        provider: provider.track_mut(),
+        locator: &mut locator,
         introspector: introspector.track(),
     };
+
+    // Prepare VM.
+    let route = Route::insert(route, id);
+    let scopes = Scopes::new(Some(library));
     let mut vm = Vm::new(vt, route.track(), id, scopes);
     let root = match source.root().cast::<ast::Markup>() {
         Some(markup) if vm.traced.is_some() => markup,
         _ => source.ast()?,
     };
 
+    // Evaluate the module.
     let result = root.eval(&mut vm);
 
     // Handle control flow.
@@ -129,20 +131,24 @@ pub fn eval_string(
         return Err(Box::new(errors));
     }
 
-    let id = SourceId::detached();
+    // Prepare VT.
-    let library = world.library();
-    let scopes = Scopes::new(Some(library));
-    let route = Route::default();
     let mut tracer = Tracer::default();
-    let mut provider = StabilityProvider::new();
+    let mut locator = Locator::default();
-    let introspector = Introspector::new(&[]);
+    let introspector = Introspector::default();
     let vt = Vt {
         world,
         tracer: tracer.track_mut(),
-        provider: provider.track_mut(),
+        locator: &mut locator,
         introspector: introspector.track(),
     };
+
+    // Prepare VM.
+    let route = Route::default();
+    let id = SourceId::detached();
+    let scopes = Scopes::new(Some(world.library()));
     let mut vm = Vm::new(vt, route.track(), id, scopes);
+
+    // Evaluate the code.
     let code = root.cast::<ast::Code>().unwrap();
     let result = code.eval(&mut vm);
 

src/export/pdf/outline.rs

@@ -12,7 +12,7 @@ use crate::util::NonZeroExt;
 pub fn write_outline(ctx: &mut PdfContext) -> Option<Ref> {
     let mut tree: Vec<HeadingNode> = vec![];
     for heading in ctx.introspector.query(&item!(heading_func).select()) {
-        let leaf = HeadingNode::leaf(heading);
+        let leaf = HeadingNode::leaf((*heading).clone());
         if let Some(last) = tree.last_mut() {
             if last.try_insert(leaf.clone(), NonZeroUsize::ONE) {
                 continue;

src/ide/analyze.rs

@@ -103,6 +103,7 @@ pub fn analyze_labels(
                 Value::Content(content) => Some(content),
                 _ => None,
             })
+            .as_ref()
             .unwrap_or(elem)
             .plain_text();
         output.push((label, Some(details)));

src/lib.rs

@@ -54,7 +54,7 @@ pub mod syntax;
 
 use std::path::Path;
 
-use comemo::{Prehashed, Track};
+use comemo::{Prehashed, Track, TrackedMut};
 
 use crate::diag::{FileResult, SourceResult};
 use crate::doc::Document;
@@ -66,16 +66,24 @@ use crate::util::Buffer;
 /// Compile a source file into a fully layouted document.
 #[tracing::instrument(skip(world))]
 pub fn compile(world: &dyn World) -> SourceResult<Document> {
-    // Evaluate the source file into a module.
     let route = Route::default();
     let mut tracer = Tracer::default();
-    let module =
-        eval::eval(world.track(), route.track(), tracer.track_mut(), world.main())?;
 
-    tracing::info!("Evaluation successful");
+    // Call `track` just once to keep comemo's ID stable.
+    let world = world.track();
+    let mut tracer = tracer.track_mut();
+
+    // Evaluate the source file into a module.
+    tracing::info!("Starting evaluation");
+    let module = eval::eval(
+        world,
+        route.track(),
+        TrackedMut::reborrow_mut(&mut tracer),
+        world.main(),
+    )?;
 
     // Typeset the module's contents.
-    model::typeset(world.track(), tracer.track_mut(), &module.content())
+    model::typeset(world, tracer, &module.content())
 }
 
 /// The environment in which typesetting occurs.

src/model/introspect.rs

@@ -1,7 +1,11 @@
+use std::cell::RefCell;
+use std::collections::HashMap;
 use std::fmt::{self, Debug, Formatter};
 use std::hash::Hash;
 use std::num::NonZeroUsize;
 
+use comemo::{Prehashed, Track, Tracked, Validate};
+use ecow::EcoVec;
 use indexmap::IndexMap;
 
 use super::{Content, Selector};
@@ -12,15 +16,15 @@ use crate::geom::{Point, Transform};
 use crate::model::Label;
 use crate::util::NonZeroExt;
 
-/// Stably identifies an element in the document across multiple layout passes.
+/// Uniquely identifies an element in the document across multiple layout passes.
 ///
-/// This struct is created by [`StabilityProvider::locate`].
+/// This struct is created by [`Locator::locate`].
 #[derive(Copy, Clone, Eq, PartialEq, Hash)]
 pub struct Location {
     /// The hash of the element.
     hash: u128,
     /// An unique number among elements with the same hash. This is the reason
-    /// we need a mutable `StabilityProvider` everywhere.
+    /// we need a `Locator` everywhere.
     disambiguator: usize,
     /// A synthetic location created from another one. This is used for example
     /// in bibliography management to create individual linkable locations for
@@ -46,39 +50,113 @@ cast_from_value! {
     Location: "location",
 }
 
-/// Provides stable identities to elements.
+/// Provides locations for elements in the document.
-#[derive(Clone, Default)]
+///
-pub struct StabilityProvider {
+/// A [`Location`] consists of an element's hash plus a disambiguator. Just the
-    hashes: Vec<u128>,
+/// hash is not enough because we can have multiple equal elements with the same
-    checkpoints: Vec<usize>,
+/// hash (not a hash collision, just equal elements!). Between these, we
+/// disambiguate with an increasing number. In principle, the disambiguator
+/// could just be counted up. However, counting is an impure operation and as
+/// such we can't count across a memoization boundary. [^1]
+///
+/// Instead, we only mutate within a single "layout run" and combine the results
+/// with disambiguators from an outer tracked locator. Thus, the locators form a
+/// "tracked chain". When a layout run ends, its mutations are discarded and, on
+/// the other side of the memoization boundary, we
+/// [reconstruct](Self::visit_frame) them from the resulting [frames](Frame).
+///
+/// [^1]: Well, we could with [`TrackedMut`](comemo::TrackedMut), but the
+/// overhead is quite high, especially since we need to save & undo the counting
+/// when only measuring.
+#[derive(Default)]
+pub struct Locator<'a> {
+    /// Maps from a hash to the maximum number we've seen for this hash. This
+    /// number becomes the `disambiguator`.
+    hashes: RefCell<HashMap<u128, usize>>,
+    /// An outer `Locator`, from which we can get disambiguator for hashes
+    /// outside of the current "layout run".
+    ///
+    /// We need to override the constraint's lifetime here so that `Tracked` is
+    /// covariant over the constraint. If it becomes invariant, we're in for a
+    /// world of lifetime pain.
+    outer: Option<Tracked<'a, Self, <Locator<'static> as Validate>::Constraint>>,
 }
 
-impl StabilityProvider {
+impl<'a> Locator<'a> {
-    /// Create a new stability provider.
+    /// Create a new locator.
     pub fn new() -> Self {
         Self::default()
     }
+
+    /// Create a new chained locator.
+    pub fn chained(outer: Tracked<'a, Self>) -> Self {
+        Self { outer: Some(outer), ..Default::default() }
+    }
+
+    /// Start tracking this locator.
+    ///
+    /// In comparison to [`Track::track`], this method skips this chain link
+    /// if it does not contribute anything.
+    pub fn track(&self) -> Tracked<'_, Self> {
+        match self.outer {
+            Some(outer) if self.hashes.borrow().is_empty() => outer,
+            _ => Track::track(self),
+        }
+    }
+
+    /// Produce a stable identifier for this call site.
+    pub fn locate(&mut self, hash: u128) -> Location {
+        // Get the current disambiguator for this hash.
+        let disambiguator = self.disambiguator_impl(hash);
+
+        // Bump the next disambiguator up by one.
+        self.hashes.borrow_mut().insert(hash, disambiguator + 1);
+
+        // Create the location in its default variant.
+        Location { hash, disambiguator, variant: 0 }
+    }
+
+    /// Advance past a frame.
+    pub fn visit_frame(&mut self, frame: &Frame) {
+        for (_, item) in frame.items() {
+            match item {
+                FrameItem::Group(group) => self.visit_frame(&group.frame),
+                FrameItem::Meta(Meta::Elem(elem), _) => {
+                    let mut hashes = self.hashes.borrow_mut();
+                    let loc = elem.location().unwrap();
+                    let entry = hashes.entry(loc.hash).or_default();
+
+                    // Next disambiguator needs to be at least one larger than
+                    // the maximum we've seen so far.
+                    *entry = (*entry).max(loc.disambiguator + 1);
+                }
+                _ => {}
+            }
+        }
+    }
+
+    /// Advance past a number of frames.
+    pub fn visit_frames<'b>(&mut self, frames: impl IntoIterator<Item = &'b Frame>) {
+        for frame in frames {
+            self.visit_frame(frame);
+        }
+    }
+
+    /// The current disambiguator for the given hash.
+    fn disambiguator_impl(&self, hash: u128) -> usize {
+        *self
+            .hashes
+            .borrow_mut()
+            .entry(hash)
+            .or_insert_with(|| self.outer.map_or(0, |outer| outer.disambiguator(hash)))
+    }
 }
 
 #[comemo::track]
-impl StabilityProvider {
+impl<'a> Locator<'a> {
-    /// Produce a stable identifier for this call site.
+    /// The current disambiguator for the hash.
-    pub fn locate(&mut self, hash: u128) -> Location {
+    fn disambiguator(&self, hash: u128) -> usize {
-        let disambiguator = self.hashes.iter().filter(|&&prev| prev == hash).count();
+        self.disambiguator_impl(hash)
-        self.hashes.push(hash);
-        Location { hash, disambiguator, variant: 0 }
-    }
-
-    /// Create a checkpoint of the state that can be restored.
-    pub fn save(&mut self) {
-        self.checkpoints.push(self.hashes.len());
-    }
-
-    /// Restore the last checkpoint.
-    pub fn restore(&mut self) {
-        if let Some(checkpoint) = self.checkpoints.pop() {
-            self.hashes.truncate(checkpoint);
-        }
     }
 }
 
@@ -87,9 +165,14 @@ pub struct Introspector {
     /// The number of pages in the document.
     pages: usize,
     /// All introspectable elements.
-    elems: IndexMap<Location, (Content, Position)>,
+    elems: IndexMap<Location, (Prehashed<Content>, Position)>,
     /// The page numberings, indexed by page number minus 1.
     page_numberings: Vec<Value>,
+    /// 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>>>>,
 }
 
 impl Introspector {
@@ -100,6 +183,7 @@ impl Introspector {
             pages: frames.len(),
             elems: IndexMap::new(),
             page_numberings: vec![],
+            queries: RefCell::default(),
         };
         for (i, frame) in frames.iter().enumerate() {
             let page = NonZeroUsize::new(1 + i).unwrap();
@@ -108,11 +192,6 @@ impl Introspector {
         introspector
     }
 
-    /// Iterate over all elements.
-    pub fn all(&self) -> impl Iterator<Item = Content> + '_ {
-        self.elems.values().map(|(c, _)| c).cloned()
-    }
-
     /// Extract metadata from a frame.
     #[tracing::instrument(skip_all)]
     fn extract(&mut self, frame: &Frame, page: NonZeroUsize, ts: Transform) {
@@ -130,7 +209,7 @@ impl Introspector {
                     let pos = pos.transform(ts);
                     let ret = self.elems.insert(
                         content.location().unwrap(),
-                        (content.clone(), Position { page, point: pos }),
+                        (Prehashed::new(content.clone()), Position { page, point: pos }),
                     );
                     assert!(ret.is_none(), "duplicate locations");
                 }
@@ -141,6 +220,23 @@ impl Introspector {
             }
         }
     }
+
+    /// Iterate over all locatable elements.
+    pub fn all(&self) -> impl Iterator<Item = &Prehashed<Content>> + '_ {
+        self.elems.values().map(|(c, _)| c)
+    }
+
+    /// Get an element by its location.
+    fn get(&self, location: &Location) -> Option<&Prehashed<Content>> {
+        self.elems.get(location).map(|(elem, _)| elem)
+    }
+
+    /// Get the index of this element among all.
+    fn index(&self, elem: &Content) -> usize {
+        self.elems
+            .get_index_of(&elem.location().unwrap())
+            .unwrap_or(usize::MAX)
+    }
 }
 
 #[comemo::track]
@@ -150,40 +246,75 @@ impl Introspector {
         self.pages > 0
     }
 
-    /// Get an element from the position cache.
-    pub fn location(&self, location: &Location) -> Option<Content> {
-        self.elems.get(location).map(|(c, _)| c).cloned()
-    }
-
     /// Query for all matching elements.
-    #[tracing::instrument(skip_all)]
+    pub fn query(&self, selector: &Selector) -> EcoVec<Prehashed<Content>> {
-    pub fn query<'a>(&'a self, selector: &'a Selector) -> Vec<Content> {
+        let hash = crate::util::hash128(selector);
-        match selector {
+        if let Some(output) = self.queries.borrow().get(&hash) {
-            Selector::Location(location) => self
+            return output.clone();
-                .elems
-                .get(location)
-                .map(|(content, _)| content)
-                .cloned()
-                .into_iter()
-                .collect(),
-            _ => selector.match_iter(self).collect(),
         }
+
+        let output = match selector {
+            Selector::Elem(..)
+            | Selector::Label(_)
+            | Selector::Regex(_)
+            | Selector::Can(_)
+            | Selector::Or(_)
+            | Selector::And(_) => {
+                self.all().filter(|elem| selector.matches(elem)).cloned().collect()
+            }
+
+            Selector::Location(location) => {
+                self.get(location).cloned().into_iter().collect()
+            }
+            Selector::Before { selector, end, inclusive } => {
+                let mut list = self.query(selector);
+                if let Some(end) = self.query_first(end) {
+                    // Determine which elements are before `end`.
+                    let split = match list
+                        .binary_search_by_key(&self.index(&end), |elem| self.index(elem))
+                    {
+                        // Element itself is contained.
+                        Ok(i) => i + *inclusive as usize,
+                        // Element itself is not contained.
+                        Err(i) => i,
+                    };
+                    list = list[..split].into();
+                }
+                list
+            }
+            Selector::After { selector, start, inclusive } => {
+                let mut list = self.query(selector);
+                if let Some(start) = self.query_first(start) {
+                    // Determine which elements are after `start`.
+                    let split = match list
+                        .binary_search_by_key(&self.index(&start), |elem| {
+                            self.index(elem)
+                        }) {
+                        // Element itself is contained.
+                        Ok(i) => i + !*inclusive as usize,
+                        // Element itself is not contained.
+                        Err(i) => i,
+                    };
+                    list = list[split..].into();
+                }
+                list
+            }
+        };
+
+        self.queries.borrow_mut().insert(hash, output.clone());
+        output
     }
 
-    /// Query for the first matching element.
+    /// Query for the first element that matches the selector.
-    #[tracing::instrument(skip_all)]
+    pub fn query_first(&self, selector: &Selector) -> Option<Prehashed<Content>> {
-    pub fn query_first<'a>(&'a self, selector: &'a Selector) -> Option<Content> {
         match selector {
-            Selector::Location(location) => {
+            Selector::Location(location) => self.get(location).cloned(),
-                self.elems.get(location).map(|(content, _)| content).cloned()
+            _ => self.query(selector).first().cloned(),
-            }
-            _ => selector.match_iter(self).next(),
         }
     }
 
     /// Query for a unique element with the label.
-    #[tracing::instrument(skip(self))]
+    pub fn query_label(&self, label: &Label) -> StrResult<Prehashed<Content>> {
-    pub fn query_label(&self, label: &Label) -> StrResult<Content> {
         let mut found = None;
         for elem in self.all().filter(|elem| elem.label() == Some(label)) {
             if found.is_some() {
@@ -199,40 +330,28 @@ impl Introspector {
         NonZeroUsize::new(self.pages).unwrap_or(NonZeroUsize::ONE)
     }
 
-    /// Find the page number for the given location.
-    pub fn page(&self, location: Location) -> NonZeroUsize {
-        self.position(location).page
-    }
-
     /// Gets the page numbering for the given location, if any.
-    #[tracing::instrument(skip(self))]
     pub fn page_numbering(&self, location: Location) -> Value {
         let page = self.page(location);
         self.page_numberings.get(page.get() - 1).cloned().unwrap_or_default()
     }
 
+    /// Find the page number for the given location.
+    pub fn page(&self, location: Location) -> NonZeroUsize {
+        self.position(location).page
+    }
+
     /// Find the position for the given location.
-    #[tracing::instrument(skip(self))]
     pub fn position(&self, location: Location) -> Position {
         self.elems
             .get(&location)
             .map(|(_, loc)| *loc)
             .unwrap_or(Position { page: NonZeroUsize::ONE, point: Point::zero() })
     }
+}
 
-    /// Checks whether `a` is before `b` in the document.
+impl Default for Introspector {
-    pub fn is_before(&self, a: Location, b: Location, inclusive: bool) -> bool {
+    fn default() -> Self {
-        let a = self.elems.get_index_of(&a).unwrap();
+        Self::new(&[])
-        let b = self.elems.get_index_of(&b).unwrap();
-        if inclusive {
-            a <= b
-        } else {
-            a < b
-        }
-    }
-
-    /// Checks whether `a` is after `b` in the document.
-    pub fn is_after(&self, a: Location, b: Location, inclusive: bool) -> bool {
-        !self.is_before(a, b, !inclusive)
     }
 }

src/model/mod.rs

@@ -6,13 +6,15 @@ mod introspect;
 mod realize;
 mod styles;
 
+pub use typst_macros::element;
+
 pub use self::content::*;
 pub use self::element::*;
 pub use self::introspect::*;
 pub use self::realize::*;
 pub use self::styles::*;
 
-pub use typst_macros::element;
+use std::mem::ManuallyDrop;
 
 use comemo::{Track, Tracked, TrackedMut, Validate};
 
@@ -29,38 +31,51 @@ pub fn typeset(
     mut tracer: TrackedMut<Tracer>,
     content: &Content,
 ) -> SourceResult<Document> {
-    tracing::info!("Starting layout");
+    tracing::info!("Starting typesetting");
+
     let library = world.library();
     let styles = StyleChain::new(&library.styles);
 
     let mut document;
     let mut iter = 0;
-    let mut introspector = Introspector::new(&[]);
+
+    // We need `ManuallyDrop` until this lands in stable:
+    // https://github.com/rust-lang/rust/issues/70919
+    let mut introspector = ManuallyDrop::new(Introspector::new(&[]));
 
     // Relayout until all introspections stabilize.
     // If that doesn't happen within five attempts, we give up.
     loop {
         tracing::info!("Layout iteration {iter}");
 
-        let mut provider = StabilityProvider::new();
         let constraint = <Introspector as Validate>::Constraint::new();
+        let mut locator = Locator::new();
         let mut vt = Vt {
             world,
             tracer: TrackedMut::reborrow_mut(&mut tracer),
-            provider: provider.track_mut(),
+            locator: &mut locator,
             introspector: introspector.track_with(&constraint),
         };
 
-        document = (library.items.layout)(&mut vt, content, styles)?;
+        // Layout!
-        iter += 1;
+        let result = (library.items.layout)(&mut vt, content, styles)?;
 
-        introspector = Introspector::new(&document.pages);
+        // Drop the old introspector.
+        ManuallyDrop::into_inner(introspector);
+
+        // Only now assign the document and construct the new introspector.
+        document = result;
+        introspector = ManuallyDrop::new(Introspector::new(&document.pages));
+        iter += 1;
 
         if iter >= 5 || introspector.validate(&constraint) {
             break;
         }
     }
 
+    // Drop the introspector.
+    ManuallyDrop::into_inner(introspector);
+
     Ok(document)
 }
 
@@ -73,19 +88,7 @@ pub struct Vt<'a> {
     /// The tracer for inspection of the values an expression produces.
     pub tracer: TrackedMut<'a, Tracer>,
     /// Provides stable identities to elements.
-    pub provider: TrackedMut<'a, StabilityProvider>,
+    pub locator: &'a mut Locator<'a>,
     /// Provides access to information about the document.
     pub introspector: Tracked<'a, Introspector>,
 }
-
-impl Vt<'_> {
-    /// Mutably reborrow with a shorter lifetime.
-    pub fn reborrow_mut(&mut self) -> Vt<'_> {
-        Vt {
-            world: self.world,
-            tracer: TrackedMut::reborrow_mut(&mut self.tracer),
-            provider: TrackedMut::reborrow_mut(&mut self.provider),
-            introspector: self.introspector,
-        }
-    }
-}

src/model/realize.rs

@@ -37,7 +37,7 @@ pub fn realize(
     if target.needs_preparation() {
         let mut elem = target.clone();
         if target.can::<dyn Locatable>() || target.label().is_some() {
-            let location = vt.provider.locate(hash128(target));
+            let location = vt.locator.locate(hash128(target));
             elem.set_location(location);
         }
 

src/model/styles.rs

@@ -8,7 +8,7 @@ use std::sync::Arc;
 use comemo::Prehashed;
 use ecow::{eco_format, eco_vec, EcoString, EcoVec};
 
-use super::{Content, ElemFunc, Element, Introspector, Label, Location, Vt};
+use super::{Content, ElemFunc, Element, Label, Location, Vt};
 use crate::diag::{SourceResult, StrResult, Trace, Tracepoint};
 use crate::eval::{cast_from_value, Args, Cast, CastInfo, Dict, Func, Regex, Value, Vm};
 use crate::model::Locatable;
@@ -50,8 +50,7 @@ impl Styles {
         *self = outer;
     }
 
-    /// Apply one outer styles. Like [`chain_one`](StyleChain::chain_one), but
+    /// Apply one outer styles.
-    /// in-place.
     pub fn apply_one(&mut self, outer: Style) {
         self.0.insert(0, Prehashed::new(outer));
     }
@@ -322,77 +321,6 @@ impl Selector {
         }
     }
 
-    /// Matches the selector for an introspector.
-    pub fn match_iter<'a>(
-        &'a self,
-        introspector: &'a Introspector,
-    ) -> Box<dyn Iterator<Item = Content> + 'a> {
-        self.match_iter_inner(introspector, introspector.all())
-    }
-
-    /// Match the selector against the given list of elements. Returns an
-    /// iterator over the matching elements.
-    fn match_iter_inner<'a>(
-        &'a self,
-        introspector: &'a Introspector,
-        parent: impl Iterator<Item = Content> + 'a,
-    ) -> Box<dyn Iterator<Item = Content> + 'a> {
-        match self {
-            Self::Location(location) => {
-                Box::new(introspector.location(location).into_iter())
-            }
-            Self::Or(selectors) => Box::new(parent.filter(|element| {
-                selectors.iter().any(|selector| {
-                    selector
-                        .match_iter_inner(introspector, std::iter::once(element.clone()))
-                        .next()
-                        .is_some()
-                })
-            })),
-            Self::And(selectors) => Box::new(parent.filter(|element| {
-                selectors.iter().all(|selector| {
-                    selector
-                        .match_iter_inner(introspector, std::iter::once(element.clone()))
-                        .next()
-                        .is_some()
-                })
-            })),
-            Self::Before { selector, end: location, inclusive } => {
-                if let Some(content) = introspector.query_first(location) {
-                    let loc = content.location().unwrap();
-                    Box::new(selector.match_iter_inner(introspector, parent).take_while(
-                        move |elem| {
-                            introspector.is_before(
-                                elem.location().unwrap(),
-                                loc,
-                                *inclusive,
-                            )
-                        },
-                    ))
-                } else {
-                    Box::new(selector.match_iter_inner(introspector, parent))
-                }
-            }
-            Self::After { selector, start: location, inclusive } => {
-                if let Some(content) = introspector.query_first(location) {
-                    let loc = content.location().unwrap();
-                    Box::new(selector.match_iter_inner(introspector, parent).skip_while(
-                        move |elem| {
-                            introspector.is_before(
-                                elem.location().unwrap(),
-                                loc,
-                                !*inclusive,
-                            )
-                        },
-                    ))
-                } else {
-                    Box::new(std::iter::empty())
-                }
-            }
-            other => Box::new(parent.filter(move |content| other.matches(content))),
-        }
-    }
-
     /// Whether the selector matches for the target.
     pub fn matches(&self, target: &Content) -> bool {
         match self {
@@ -412,6 +340,7 @@ impl Selector {
             Self::Or(selectors) => selectors.iter().any(move |sel| sel.matches(target)),
             Self::And(selectors) => selectors.iter().all(move |sel| sel.matches(target)),
             Self::Location(location) => target.location() == Some(*location),
+            // Not supported here.
             Self::Before { .. } | Self::After { .. } => false,
         }
     }

tests/typ/meta/query-before-after.typ

@@ -34,6 +34,7 @@
 === Subsubsection 1
 === Subsubsection 2
 == Subsection 3
+
 = Section 2
 == Subsection 1
 == Subsection 2