Optimize counters and state (#4223)

crates/typst/src/introspection/counter.rs

@@ -223,10 +223,7 @@ impl Counter {
         location: Location,
     ) -> SourceResult<CounterState> {
         let sequence = self.sequence(engine)?;
-        let offset = engine
+        let offset = engine.introspector.query_count_before(&self.selector(), location);
-            .introspector
-            .query(&self.selector().before(location.into(), true))
-            .len();
         let (mut at_state, at_page) = sequence[offset].clone();
         let (mut final_state, final_page) = sequence.last().unwrap().clone();
         if self.is_page() {
@@ -245,16 +242,14 @@ impl Counter {
     pub fn at_loc(
         &self,
         engine: &mut Engine,
-        loc: Location,
+        location: Location,
     ) -> SourceResult<CounterState> {
         let sequence = self.sequence(engine)?;
-        let offset = engine
+        let offset = engine.introspector.query_count_before(&self.selector(), location);
-            .introspector
-            .query(&self.selector().before(loc.into(), true))
-            .len();
         let (mut state, page) = sequence[offset].clone();
         if self.is_page() {
-            let delta = engine.introspector.page(loc).get().saturating_sub(page.get());
+            let delta =
+                engine.introspector.page(location).get().saturating_sub(page.get());
             state.step(NonZeroUsize::ONE, delta);
         }
         Ok(state)

crates/typst/src/introspection/introspector.rs

@@ -121,7 +121,7 @@ impl Introspector {
                     indices.iter().map(|&index| self.elems[index].0.clone()).collect()
                 })
                 .unwrap_or_default(),
-            Selector::Elem(..) | Selector::Regex(_) | Selector::Can(_) => self
+            Selector::Elem(..) | Selector::Can(_) => self
                 .all()
                 .filter(|elem| selector.matches(elem, None))
                 .cloned()
@@ -188,6 +188,8 @@ impl Introspector {
                 .into_iter()
                 .map(|index| self.elems[index].0.clone())
                 .collect(),
+            // Not supported here.
+            Selector::Regex(_) => EcoVec::new(),
         };
 
         self.queries.insert(hash, output.clone());
@@ -198,6 +200,11 @@ impl Introspector {
     pub fn query_first(&self, selector: &Selector) -> Option<Content> {
         match selector {
             Selector::Location(location) => self.get(location).cloned(),
+            Selector::Label(label) => self
+                .labels
+                .get(label)
+                .and_then(|indices| indices.first())
+                .map(|&index| self.elems[index].0.clone()),
             _ => self.query(selector).first().cloned(),
         }
     }
@@ -236,6 +243,21 @@ impl Introspector {
         Ok(&self.elems[indices[0]].0)
     }
 
+    /// This is an optimized version of
+    /// `query(selector.before(end, true).len()` used by counters and state.
+    pub fn query_count_before(&self, selector: &Selector, end: Location) -> usize {
+        // See `query()` for details.
+        let list = self.query(selector);
+        if let Some(end) = self.get(&end) {
+            match self.binary_search(&list, end) {
+                Ok(i) => i + 1,
+                Err(i) => i,
+            }
+        } else {
+            list.len()
+        }
+    }
+
     /// The total number pages.
     pub fn pages(&self) -> NonZeroUsize {
         NonZeroUsize::new(self.pages).unwrap_or(NonZeroUsize::ONE)

crates/typst/src/introspection/state.rs

@@ -202,10 +202,7 @@ impl State {
     /// Get the value of the state at the given location.
     pub fn at_loc(&self, engine: &mut Engine, loc: Location) -> SourceResult<Value> {
         let sequence = self.sequence(engine)?;
-        let offset = engine
+        let offset = engine.introspector.query_count_before(&self.selector(), loc);
-            .introspector
-            .query(&self.selector().before(loc.into(), true))
-            .len();
         Ok(sequence[offset].clone())
     }