From 715f9fb0a5b5242c405e2c9277596fad7509e0db Mon Sep 17 00:00:00 2001
From: Laurenz <laurmaedje@gmail.com>
Date: Tue, 4 Apr 2023 23:27:51 +0200
Subject: [PATCH] Locatable selectors

---
 library/src/meta/counter.rs | 23 ++++++++---------
 library/src/meta/query.rs   | 21 +++-------------
 library/src/prelude.rs      |  5 ++--
 src/eval/cast.rs            | 16 +++++++++---
 src/model/styles.rs         | 50 +++++++++++++++++++++++++++++++++++--
 5 files changed, 76 insertions(+), 39 deletions(-)

diff --git a/library/src/meta/counter.rs b/library/src/meta/counter.rs
index 5ad9b48d4..d771168d4 100644
--- a/library/src/meta/counter.rs
+++ b/library/src/meta/counter.rs
@@ -276,10 +276,11 @@ use crate::prelude::*;
 pub fn counter(
     /// The key that identifies this counter.
     ///
-    /// - If this is the [`page`]($func/page) function, counts through pages.
-    /// - If this is any other element function, counts through its elements.
     /// - If it is a string, creates a custom counter that is only affected by
-    ///   manual updates.
+    ///   manual updates,
+    /// - If this is a `{<label>}`, counts through all elements with that label,
+    /// - If this is an element function or selector, counts through its elements,
+    /// - If this is the [`page`]($func/page) function, counts through pages.
     key: CounterKey,
 ) -> Value {
     Value::dynamic(Counter::new(key))
@@ -485,18 +486,14 @@ cast_from_value! {
     CounterKey,
     v: Str => Self::Str(v),
     label: Label => Self::Selector(Selector::Label(label)),
-    element: ElemFunc => {
-        if element == PageElem::func() {
-            return Ok(Self::Page);
+    v: ElemFunc => {
+        if v == PageElem::func() {
+            Self::Page
+        } else {
+            Self::Selector(LocatableSelector::cast(Value::from(v))?.0)
         }
-
-        if !Content::new(element).can::<dyn Locatable>() {
-            Err(eco_format!("cannot count through {}s", element.name()))?;
-        }
-
-        Self::Selector(Selector::Elem(element, None))
     },
-    selector: Selector => Self::Selector(selector),
+    selector: LocatableSelector => Self::Selector(selector.0),
 }
 
 impl Debug for CounterKey {
diff --git a/library/src/meta/query.rs b/library/src/meta/query.rs
index d4fab6353..4241c19d7 100644
--- a/library/src/meta/query.rs
+++ b/library/src/meta/query.rs
@@ -94,8 +94,8 @@ use crate::prelude::*;
 /// Returns: content
 #[func]
 pub fn query(
-    /// Can be an element function like a `heading` or `figure` or a
-    /// `{<label>}`.
+    /// Can be an element function like a `heading` or `figure`, a `{<label>}`
+    /// or a more complex selector like `{heading.where(level: 1)}`.
     ///
     /// Currently, only a subset of element functions is supported. Aside from
     /// headings and figures, this includes equations, references and all
@@ -104,7 +104,7 @@ pub fn query(
     /// have an explicit label attached to them. This limitation will be
     /// resolved
     /// in the future.
-    target: Target,
+    target: LocatableSelector,
 
     /// Can be any location. Why is it required then? As noted before, Typst has
     /// to evaluate parts of your code multiple times to determine the values of
@@ -151,18 +151,3 @@ pub fn query(
     };
     elements.into()
 }
-
-/// A query target.
-struct Target(Selector);
-
-cast_from_value! {
-    Target,
-    label: Label => Self(Selector::Label(label)),
-    element: ElemFunc => {
-        if !Content::new(element).can::<dyn Locatable>() {
-            Err(eco_format!("cannot query for {}s", element.name()))?;
-        }
-
-        Self(Selector::Elem(element, None))
-    }
-}
diff --git a/library/src/prelude.rs b/library/src/prelude.rs
index 4c83cf312..fc71e2c2b 100644
--- a/library/src/prelude.rs
+++ b/library/src/prelude.rs
@@ -23,8 +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, Location, MetaElem, Resolve, Selector, Set, Show,
-    StabilityProvider, StyleChain, StyleVec, Styles, Synthesize, Unlabellable, Vt,
+    Introspector, Label, Locatable, LocatableSelector, Location, MetaElem, Resolve,
+    Selector, Set, Show, StabilityProvider, StyleChain, StyleVec, Styles, Synthesize,
+    Unlabellable, Vt,
 };
 #[doc(no_inline)]
 pub use typst::syntax::{Span, Spanned};
diff --git a/src/eval/cast.rs b/src/eval/cast.rs
index 7b507dbb0..92e4212f8 100644
--- a/src/eval/cast.rs
+++ b/src/eval/cast.rs
@@ -237,7 +237,7 @@ impl<T> Variadics for Vec<T> {
 }
 
 /// Describes a possible value for a cast.
-#[derive(Debug, Clone, Hash)]
+#[derive(Debug, Clone, Hash, PartialEq, PartialOrd)]
 pub enum CastInfo {
     /// Any value is okay.
     Any,
@@ -302,15 +302,23 @@ impl Add for CastInfo {
     fn add(self, rhs: Self) -> Self {
         Self::Union(match (self, rhs) {
             (Self::Union(mut lhs), Self::Union(rhs)) => {
-                lhs.extend(rhs);
+                for cast in rhs {
+                    if !lhs.contains(&cast) {
+                        lhs.push(cast);
+                    }
+                }
                 lhs
             }
             (Self::Union(mut lhs), rhs) => {
-                lhs.push(rhs);
+                if !lhs.contains(&rhs) {
+                    lhs.push(rhs);
+                }
                 lhs
             }
             (lhs, Self::Union(mut rhs)) => {
-                rhs.insert(0, lhs);
+                if !rhs.contains(&lhs) {
+                    rhs.insert(0, lhs);
+                }
                 rhs
             }
             (lhs, rhs) => vec![lhs, rhs],
diff --git a/src/model/styles.rs b/src/model/styles.rs
index 4b309ed04..097c5138c 100644
--- a/src/model/styles.rs
+++ b/src/model/styles.rs
@@ -6,8 +6,9 @@ use std::mem;
 use ecow::{eco_format, eco_vec, EcoString, EcoVec};
 
 use super::{Content, ElemFunc, Element, Label, Vt};
-use crate::diag::{SourceResult, Trace, Tracepoint};
-use crate::eval::{cast_from_value, Args, Cast, Dict, Func, Regex, Value, Vm};
+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;
 use crate::syntax::Span;
 use crate::util::pretty_array_like;
 
@@ -337,6 +338,51 @@ cast_from_value! {
     regex: Regex => Self::Regex(regex),
 }
 
+/// A selector that can be used with `query`. Hopefully, this is made obsolote
+/// by a more powerful query mechanism in the future.
+#[derive(Clone, PartialEq, Hash)]
+pub struct LocatableSelector(pub Selector);
+
+impl Cast for LocatableSelector {
+    fn is(value: &Value) -> bool {
+        matches!(value, Value::Label(_) | Value::Func(_))
+            || value.type_name() == "selector"
+    }
+
+    fn cast(value: Value) -> StrResult<Self> {
+        fn validate(selector: &Selector) -> StrResult<()> {
+            match &selector {
+                Selector::Elem(elem, _) if !elem.can::<dyn Locatable>() => {
+                    Err(eco_format!("{} is not locatable", elem.name()))?
+                }
+                Selector::Regex(_) => Err("text is not locatable")?,
+                Selector::Any(list) | Selector::All(list) => {
+                    for selector in list {
+                        validate(selector)?;
+                    }
+                }
+                _ => {}
+            }
+            Ok(())
+        }
+
+        if !Self::is(&value) {
+            return <Self as Cast>::error(value);
+        }
+
+        let selector = Selector::cast(value)?;
+        validate(&selector)?;
+        Ok(Self(selector))
+    }
+
+    fn describe() -> CastInfo {
+        CastInfo::Union(vec![
+            CastInfo::Type("label"),
+            CastInfo::Type("function"),
+            CastInfo::Type("selector"),
+        ])
+    }
+}
 /// A show rule transformation that can be applied to a match.
 #[derive(Clone, PartialEq, Hash)]
 pub enum Transform {