Selector rework (#640)

1 files changed, 47 insertions, 28 deletions

diff --git a/src/model/introspect.rs b/src/model/introspect.rs
index 9697f850..031f2d5e 100644
--- a/src/model/introspect.rs
+++ b/src/model/introspect.rs
@@ -84,9 +84,11 @@ impl StabilityProvider {
 
 /// Can be queried for elements and their positions.
 pub struct Introspector {
+    /// The number of pages in the document.
     pages: usize,
-    elems: IndexMap<Option<Location>, (Content, Position)>,
-    // Indexed by page number.
+    /// All introspectable elements.
+    elems: IndexMap<Location, (Content, Position)>,
+    /// The page numberings, indexed by page number minus 1.
     page_numberings: Vec<Value>,
 }
 
@@ -106,8 +108,8 @@ impl Introspector {
     }
 
     /// Iterate over all elements.
-    pub fn all(&self) -> impl Iterator<Item = &Content> {
-        self.elems.values().map(|(elem, _)| elem)
+    pub fn all(&self) -> impl Iterator<Item = Content> + '_ {
+        self.elems.values().map(|(c, _)| c).cloned()
     }
 
     /// Extract metadata from a frame.
@@ -121,11 +123,11 @@ impl Introspector {
                     self.extract(&group.frame, page, ts);
                 }
                 FrameItem::Meta(Meta::Elem(content), _)
-                    if !self.elems.contains_key(&content.location()) =>
+                    if !self.elems.contains_key(&content.location().unwrap()) =>
                 {
                     let pos = pos.transform(ts);
                     let ret = self.elems.insert(
-                        content.location(),
+                        content.location().unwrap(),
                         (content.clone(), Position { page, point: pos }),
                     );
                     assert!(ret.is_none(), "duplicate locations");
@@ -146,32 +148,33 @@ impl Introspector {
         self.pages > 0
     }
 
-    /// Query for all matching elements.
-    pub fn query(&self, selector: Selector) -> Vec<Content> {
-        self.all().filter(|elem| selector.matches(elem)).cloned().collect()
+    /// 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 element up to the given location.
-    pub fn query_before(&self, selector: Selector, location: Location) -> Vec<Content> {
-        let mut matches = vec![];
-        for elem in self.all() {
-            if selector.matches(elem) {
-                matches.push(elem.clone());
-            }
-            if elem.location() == Some(location) {
-                break;
-            }
+    /// Query for all matching elements.
+    pub fn query<'a>(&'a self, selector: &'a Selector) -> Vec<Content> {
+        match selector {
+            Selector::Location(location) => self
+                .elems
+                .get(location)
+                .map(|(content, _)| content)
+                .cloned()
+                .into_iter()
+                .collect(),
+            _ => selector.match_iter(self).collect(),
         }
-        matches
     }
 
-    /// Query for all matching elements starting from the given location.
-    pub fn query_after(&self, selector: Selector, location: Location) -> Vec<Content> {
-        self.all()
-            .skip_while(|elem| elem.location() != Some(location))
-            .filter(|elem| selector.matches(elem))
-            .cloned()
-            .collect()
+    /// Query for the first matching element.
+    pub fn query_first<'a>(&'a self, selector: &'a Selector) -> Option<Content> {
+        match selector {
+            Selector::Location(location) => {
+                self.elems.get(location).map(|(content, _)| content).cloned()
+            }
+            _ => selector.match_iter(self).next(),
+        }
     }
 
     /// Query for a unique element with the label.
@@ -205,8 +208,24 @@ impl Introspector {
     /// Find the position for the given location.
     pub fn position(&self, location: Location) -> Position {
         self.elems
-            .get(&Some(location))
+            .get(&location)
             .map(|(_, loc)| *loc)
             .unwrap_or(Position { page: NonZeroUsize::ONE, point: Point::zero() })
     }
+
+    /// Checks whether `a` is before `b` in the document.
+    pub fn is_before(&self, a: Location, b: Location, inclusive: bool) -> bool {
+        let a = self.elems.get_index_of(&a).unwrap();
+        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)
+    }
 }