From 73340f70a6f36bc46514a81b220373112fc7399e Mon Sep 17 00:00:00 2001
From: Johann Birnick <johann.birnick@hotmail.de>
Date: Sun, 19 Jan 2025 04:42:47 -0800
Subject: [PATCH] initial draft

---
 .../typst-library/src/introspection/count.rs  | 92 +++++++++++++++++++
 .../src/introspection/introspector.rs         | 54 +++++++++++
 crates/typst-library/src/introspection/mod.rs |  4 +
 3 files changed, 150 insertions(+)
 create mode 100644 crates/typst-library/src/introspection/count.rs

diff --git a/crates/typst-library/src/introspection/count.rs b/crates/typst-library/src/introspection/count.rs
new file mode 100644
index 000000000..43af61ad7
--- /dev/null
+++ b/crates/typst-library/src/introspection/count.rs
@@ -0,0 +1,92 @@
+use comemo::Tracked;
+
+use crate::diag::HintedStrResult;
+use crate::engine::Engine;
+use crate::foundations::{func, Array, Context, IntoValue, LocatableSelector, Selector};
+
+/// Counts element in the document.
+///
+/// The `count` function lets you count elements in your document of a
+/// particular type or with a particular label. It also allows to count them
+/// in a hierachical way. To use it, you first need to ensure that [context]
+/// is available.
+///
+/// Always returns an array of integers, even if called with just a single
+/// target.
+
+/// # Counting elements - simple
+///
+/// To just count elements of a single type/label up to the current location,
+/// pass the selector you want to count:
+/// ```example
+/// = Heading
+///
+/// = Another Heading
+///
+/// #context count(heading)
+///
+/// = Third Heading
+/// ```
+///
+/// Note that it will not return an integer, but an array with a single integer
+/// entry.
+
+/// # Counting elements - hierarchical
+///
+/// If you pass multiple targets, then it starts by counting the first target.
+/// Then the second target is counted _starting only from the last counted
+/// element of the first target_.
+///
+/// ```example
+/// = Some Heading
+///
+/// == Some Subheading
+///
+/// = Another Heading
+///
+/// == Some Subheading
+///
+/// #context count(heading.where(level: 1), heading.where(level: 2))
+///
+/// == Another Subheading
+/// ```
+#[func(contextual)]
+pub fn count(
+    /// The engine.
+    engine: &mut Engine,
+    /// The callsite context.
+    context: Tracked<Context>,
+    /// Each target can be
+    /// - an element function like a `heading` or `figure`,
+    /// - a `{<label>}`,
+    /// - a more complex selector like `{heading.where(level: 1)}`,
+    /// - or `{selector(heading).before(here())}`.
+    ///
+    /// Only [locatable]($location/#locatable) element functions are supported.
+    #[variadic]
+    targets: Vec<LocatableSelector>,
+    /// When passing this argument, it will count everything only from a
+    /// certain location on.
+    /// The selector must match exactly one element in the document. The most
+    /// useful kinds of selectors for this are [labels]($label) and
+    /// [locations]($location).
+    // TODO remove Option as soon as there is a special `start` location
+    #[named]
+    after: Option<LocatableSelector>,
+) -> HintedStrResult<Array> {
+    // NOTE this could be made more efficient
+    // one could directly get a slice &[Selector] from Vec<LocatableSelector>
+    // by using #[repr(transparent)] on LocatableSelector
+    let selectors: Vec<Selector> = targets.into_iter().map(|sel| sel.0).collect();
+    // TODO add argument "at" with default value "here"
+    // and compute `before` accordingly
+    let before = Some(context.location()?);
+    let after = match after {
+        Some(selector) => Some(selector.resolve_unique(engine.introspector, context)?),
+        None => None,
+    };
+
+    let nums = engine.introspector.count(&selectors, after, before);
+
+    Ok(nums.into_iter().map(IntoValue::into_value).collect())
+}
diff --git a/crates/typst-library/src/introspection/introspector.rs b/crates/typst-library/src/introspection/introspector.rs
index 8cbaea891..f175a63db 100644
--- a/crates/typst-library/src/introspection/introspector.rs
+++ b/crates/typst-library/src/introspection/introspector.rs
@@ -251,6 +251,60 @@ impl Introspector {
         }
     }
 
+    /// Count hierarchically.
+    pub fn count(
+        &self,
+        selectors: &[Selector],
+        // TODO remove Option as soon as there is a special `start` location
+        after: Option<Location>,
+        // TODO remove Option as soon as there is a special `end` location
+        before: Option<Location>,
+    ) -> SmallVec<[usize; 3]> {
+        let mut nums = SmallVec::with_capacity(selectors.len());
+
+        // can't write `mut` directly into function parameters because not supported by comemo
+        let mut after = after;
+
+        for selector in selectors.iter() {
+            let list = self.query(selector);
+
+            // count how many elements in `list` are between after and end
+            let index_start = if let Some(after) = after {
+                if let Some(after) = self.get_by_loc(&after) {
+                    match self.binary_search(&list, after) {
+                        Ok(i) => i, // TODO if after is exclusive then must add 1 here
+                        Err(i) => i,
+                    }
+                } else {
+                    0
+                }
+            } else {
+                0
+            };
+            let index_end = if let Some(before) = before {
+                if let Some(before) = self.get_by_loc(&before) {
+                    match self.binary_search(&list, before) {
+                        Ok(i) => i + 1,
+                        Err(i) => i,
+                    }
+                } else {
+                    list.len()
+                }
+            } else {
+                list.len()
+            };
+            nums.push(index_end - index_start);
+
+            // the next level should be counted only from
+            // the last element of the previous level on
+            if index_end != 0 {
+                after = Some(list[index_end - 1].location().unwrap());
+            }
+        }
+
+        nums
+    }
+
     /// The total number pages.
     pub fn pages(&self) -> NonZeroUsize {
         NonZeroUsize::new(self.pages).unwrap_or(NonZeroUsize::ONE)
diff --git a/crates/typst-library/src/introspection/mod.rs b/crates/typst-library/src/introspection/mod.rs
index b1ff2e080..ada1da633 100644
--- a/crates/typst-library/src/introspection/mod.rs
+++ b/crates/typst-library/src/introspection/mod.rs
@@ -1,5 +1,7 @@
 //! Interaction between document parts.
 
+#[path = "count.rs"]
+mod count_;
 mod counter;
 #[path = "here.rs"]
 mod here_;
@@ -14,6 +16,7 @@ mod query_;
 mod state;
 mod tag;
 
+pub use self::count_::*;
 pub use self::counter::*;
 pub use self::here_::*;
 pub use self::introspector::*;
@@ -49,5 +52,6 @@ pub fn define(global: &mut Scope) {
     global.define_elem::<MetadataElem>();
     global.define_func::<here>();
     global.define_func::<query>();
+    global.define_func::<count>();
     global.define_func::<locate>();
 }