Support selectors with and/or followed by before/after (#1883)

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

4 files changed, 173 insertions, 14 deletions

diff --git a/crates/typst/src/model/introspect.rs b/crates/typst/src/model/introspect.rs
index 42c1a9e1..2b2693d9 100644
--- a/crates/typst/src/model/introspect.rs
+++ b/crates/typst/src/model/introspect.rs
@@ -1,5 +1,5 @@
 use std::cell::RefCell;
-use std::collections::HashMap;
+use std::collections::{BTreeSet, HashMap};
 use std::fmt::{self, Debug, Formatter};
 use std::hash::Hash;
 use std::num::NonZeroUsize;
@@ -237,6 +237,15 @@ impl Introspector {
             .get_index_of(&elem.location().unwrap())
             .unwrap_or(usize::MAX)
     }
+
+    /// Perform a binary search for `elem` among the `list`.
+    fn binary_search(
+        &self,
+        list: &[Prehashed<Content>],
+        elem: &Content,
+    ) -> Result<usize, usize> {
+        list.binary_search_by_key(&self.index(elem), |elem| self.index(elem))
+    }
 }
 
 #[comemo::track]
@@ -252,12 +261,9 @@ impl Introspector {
             Selector::Elem(..)
             | Selector::Label(_)
             | Selector::Regex(_)
-            | Selector::Can(_)
-            | Selector::Or(_)
-            | Selector::And(_) => {
+            | Selector::Can(_) => {
                 self.all().filter(|elem| selector.matches(elem)).cloned().collect()
             }
-
             Selector::Location(location) => {
                 self.get(location).cloned().into_iter().collect()
             }
@@ -265,9 +271,7 @@ impl Introspector {
                 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))
-                    {
+                    let split = match self.binary_search(&list, &end) {
                         // Element itself is contained.
                         Ok(i) => i + *inclusive as usize,
                         // Element itself is not contained.
@@ -281,10 +285,7 @@ impl Introspector {
                 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)
-                        }) {
+                    let split = match self.binary_search(&list, &start) {
                         // Element itself is contained.
                         Ok(i) => i + !*inclusive as usize,
                         // Element itself is not contained.
@@ -294,6 +295,37 @@ impl Introspector {
                 }
                 list
             }
+            Selector::And(selectors) => {
+                let mut results: Vec<_> =
+                    selectors.iter().map(|sel| self.query(sel)).collect();
+
+                // Extract the smallest result list and then keep only those
+                // elements in the smallest list that are also in all other
+                // lists.
+                results
+                    .iter()
+                    .enumerate()
+                    .min_by_key(|(_, vec)| vec.len())
+                    .map(|(i, _)| i)
+                    .map(|i| results.swap_remove(i))
+                    .iter()
+                    .flatten()
+                    .filter(|candidate| {
+                        results
+                            .iter()
+                            .all(|other| self.binary_search(other, candidate).is_ok())
+                    })
+                    .cloned()
+                    .collect()
+            }
+            Selector::Or(selectors) => selectors
+                .iter()
+                .flat_map(|sel| self.query(sel))
+                .map(|elem| self.index(&elem))
+                .collect::<BTreeSet<usize>>()
+                .into_iter()
+                .map(|index| self.elems[index].0.clone())
+                .collect(),
         };
 
         self.queries.borrow_mut().insert(hash, output.clone());
diff --git a/crates/typst/src/model/selector.rs b/crates/typst/src/model/selector.rs
index 12f7fa0e..5e4f257b 100644
--- a/crates/typst/src/model/selector.rs
+++ b/crates/typst/src/model/selector.rs
@@ -64,13 +64,13 @@ impl Selector {
     }
 
     /// Transforms this selector and an iterator of other selectors into a
-    /// [`Selector::Or`] selector.
+    /// [`Selector::And`] selector.
     pub fn and(self, others: impl IntoIterator<Item = Self>) -> Self {
         Self::And(others.into_iter().chain(Some(self)).collect())
     }
 
     /// Transforms this selector and an iterator of other selectors into a
-    /// [`Selector::And`] selector.
+    /// [`Selector::Or`] selector.
     pub fn or(self, others: impl IntoIterator<Item = Self>) -> Self {
         Self::Or(others.into_iter().chain(Some(self)).collect())
     }
diff --git a/tests/ref/compiler/selector-logical.png b/tests/ref/compiler/selector-logical.png
new file mode 100644
index 00000000..eafa93c8
--- /dev/null
+++ b/tests/ref/compiler/selector-logical.png
diff --git a/tests/typ/compiler/selector-logical.typ b/tests/typ/compiler/selector-logical.typ
new file mode 100644
index 00000000..64f97384
--- /dev/null
+++ b/tests/typ/compiler/selector-logical.typ
@@ -0,0 +1,127 @@
+//Tests for logical (and/or) selectors
+
+---
+= A
+== B
+#figure([Cat], kind: "cat", supplement: [Other])
+=== D
+= E <first>
+#figure([Frog], kind: "frog", supplement: none)
+#figure([Giraffe], kind: "giraffe", supplement: none) <second>
+#figure([GiraffeCat], kind: "cat", supplement: [Other]) <second>
+= H
+#figure([Iguana], kind: "iguana", supplement: none)
+== I
+
+#let test-selector(selector, ref) = locate(loc => {
+  let elems = query(selector, loc)
+  test(elems.map(e => e.body), ref)
+})
+
+// Test `or`.
+#test-selector(
+  heading.where(level: 1).or(heading.where(level: 3)),
+  ([A], [D], [E], [H]),
+)
+
+#test-selector(
+  heading.where(level: 1).or(
+    heading.where(level: 3),
+    figure.where(kind: "frog"),
+  ),
+  ([A], [D], [E], [Frog], [H]),
+)
+
+#test-selector(
+  heading.where(level: 1).or(
+    heading.where(level: 2),
+    figure.where(kind: "frog"),
+    figure.where(kind: "cat"),
+  ),
+  ([A], [B], [Cat], [E], [Frog], [GiraffeCat], [H], [I]),
+)
+
+#test-selector(
+  figure.where(kind: "dog").or(heading.where(level: 3)),
+  ([D],),
+)
+
+#test-selector(
+  figure.where(kind: "dog").or(figure.where(kind: "fish")),
+  (),
+)
+
+// Test `or` duplicates removal.
+#test-selector(
+  heading.where(level: 1).or(heading.where(level: 1)),
+  ([A], [E], [H]),
+)
+
+// Test `and`.
+#test-selector(
+  figure.where(kind: "cat").and(figure.where(kind: "frog")),
+  (),
+)
+
+// Test `or` with `before`/`after`
+#test-selector(
+  selector(heading)
+    .before(<first>)
+    .or(selector(figure).before(<first>)),
+  ([A], [B], [Cat], [D], [E]),
+)
+
+#test-selector(
+  heading.where(level: 2)
+    .after(<first>)
+    .or(selector(figure).after(<first>)),
+  ([Frog], [Giraffe], [GiraffeCat], [Iguana], [I]),
+)
+
+// Test `and` with `after`
+#test-selector(
+   figure.where(kind: "cat")
+    .and(figure.where(supplement: [Other]))
+    .after(<first>),
+   ([GiraffeCat],),
+)
+
+// Test `and` (with nested `or`)
+#test-selector(
+  heading.where(level: 2)
+    .or(heading.where(level: 3))
+    .and(heading.where(level: 2).or(heading.where(level: 1))),
+  ([B], [I]),
+)
+
+#test-selector(
+  heading.where(level: 2)
+    .or(heading.where(level: 3), heading.where(level:1))
+    .and(
+      heading.where(level: 2).or(heading.where(level: 1)),
+      heading.where(level: 3).or(heading.where(level: 1)),
+    ),
+  ([A], [E], [H]),
+)
+
+// Test `and` with `or` and `before`/`after`
+#test-selector(
+  heading.where(level: 1).before(<first>)
+    .or(heading.where(level: 3).before(<first>))
+    .and(
+      heading.where(level: 1).before(<first>)
+        .or(heading.where(level: 2).before(<first>))
+    ),
+  ([A], [E]),
+)
+
+#test-selector(
+  heading.where(level: 1).before(<first>, inclusive: false)
+    .or(selector(figure).after(<first>))
+    .and(figure.where(kind: "iguana").or(
+      figure.where(kind: "frog"),
+      figure.where(kind: "cat"),
+      heading.where(level: 1).after(<first>),
+    )),
+  ([Frog], [GiraffeCat], [Iguana])
+)