Add side parameter to leaf_at (#3767)

2 files changed, 55 insertions, 8 deletions

diff --git a/crates/typst-syntax/src/lib.rs b/crates/typst-syntax/src/lib.rs
index 0ddb1460..6a6063d2 100644
--- a/crates/typst-syntax/src/lib.rs
+++ b/crates/typst-syntax/src/lib.rs
@@ -21,7 +21,7 @@ pub use self::kind::SyntaxKind;
 pub use self::lexer::{
     is_id_continue, is_id_start, is_ident, is_newline, link_prefix, split_newlines,
 };
-pub use self::node::{LinkedChildren, LinkedNode, SyntaxError, SyntaxNode};
+pub use self::node::{LinkedChildren, LinkedNode, Side, SyntaxError, SyntaxNode};
 pub use self::parser::{parse, parse_code, parse_math};
 pub use self::path::VirtualPath;
 pub use self::source::Source;
diff --git a/crates/typst-syntax/src/node.rs b/crates/typst-syntax/src/node.rs
index 3c93cd84..d047e751 100644
--- a/crates/typst-syntax/src/node.rs
+++ b/crates/typst-syntax/src/node.rs
@@ -811,6 +811,13 @@ impl<'a> LinkedNode<'a> {
     }
 }
 
+/// Indicates whether the cursor is before the related byte index, or after.
+#[derive(Debug, Clone)]
+pub enum Side {
+    Before,
+    After,
+}
+
 /// Access to leafs.
 impl<'a> LinkedNode<'a> {
     /// Get the rightmost non-trivia leaf before this node.
@@ -840,8 +847,8 @@ impl<'a> LinkedNode<'a> {
         None
     }
 
-    /// Get the leaf at the specified byte offset.
-    pub fn leaf_at(&self, cursor: usize) -> Option<Self> {
+    /// Get the leaf immediately before the specified byte offset.
+    fn leaf_before(&self, cursor: usize) -> Option<Self> {
         if self.node.children().len() == 0 && cursor <= self.offset + self.len() {
             return Some(self.clone());
         }
@@ -853,7 +860,7 @@ impl<'a> LinkedNode<'a> {
             if (offset < cursor && cursor <= offset + len)
                 || (offset == cursor && i + 1 == count)
             {
-                return child.leaf_at(cursor);
+                return child.leaf_before(cursor);
             }
             offset += len;
         }
@@ -861,6 +868,32 @@ impl<'a> LinkedNode<'a> {
         None
     }
 
+    /// Get the leaf after the specified byte offset.
+    fn leaf_after(&self, cursor: usize) -> Option<Self> {
+        if self.node.children().len() == 0 && cursor < self.offset + self.len() {
+            return Some(self.clone());
+        }
+
+        let mut offset = self.offset;
+        for child in self.children() {
+            let len = child.len();
+            if offset <= cursor && cursor < offset + len {
+                return child.leaf_after(cursor);
+            }
+            offset += len;
+        }
+
+        None
+    }
+
+    /// Get the leaf at the specified byte offset.
+    pub fn leaf_at(&self, cursor: usize, side: Side) -> Option<Self> {
+        match side {
+            Side::Before => self.leaf_before(cursor),
+            Side::After => self.leaf_after(cursor),
+        }
+    }
+
     /// Find the rightmost contained non-trivia leaf.
     pub fn rightmost_leaf(&self) -> Option<Self> {
         if self.is_leaf() && !self.kind().is_trivia() {
@@ -974,8 +1007,13 @@ mod tests {
     fn test_linked_node() {
         let source = Source::detached("#set text(12pt, red)");
 
-        // Find "text".
-        let node = LinkedNode::new(source.root()).leaf_at(7).unwrap();
+        // Find "text" with Before.
+        let node = LinkedNode::new(source.root()).leaf_at(7, Side::Before).unwrap();
+        assert_eq!(node.offset(), 5);
+        assert_eq!(node.text(), "text");
+
+        // Find "text" with After.
+        let node = LinkedNode::new(source.root()).leaf_at(7, Side::After).unwrap();
         assert_eq!(node.offset(), 5);
         assert_eq!(node.text(), "text");
 
@@ -988,17 +1026,26 @@ mod tests {
     #[test]
     fn test_linked_node_non_trivia_leaf() {
         let source = Source::detached("#set fun(12pt, red)");
-        let leaf = LinkedNode::new(source.root()).leaf_at(6).unwrap();
+        let leaf = LinkedNode::new(source.root()).leaf_at(6, Side::Before).unwrap();
         let prev = leaf.prev_leaf().unwrap();
         assert_eq!(leaf.text(), "fun");
         assert_eq!(prev.text(), "set");
 
+        // Check position 9 with Before.
         let source = Source::detached("#let x = 10");
-        let leaf = LinkedNode::new(source.root()).leaf_at(9).unwrap();
+        let leaf = LinkedNode::new(source.root()).leaf_at(9, Side::Before).unwrap();
         let prev = leaf.prev_leaf().unwrap();
         let next = leaf.next_leaf().unwrap();
         assert_eq!(prev.text(), "=");
         assert_eq!(leaf.text(), " ");
         assert_eq!(next.text(), "10");
+
+        // Check position 9 with After.
+        let source = Source::detached("#let x = 10");
+        let leaf = LinkedNode::new(source.root()).leaf_at(9, Side::After).unwrap();
+        let prev = leaf.prev_leaf().unwrap();
+        assert!(leaf.next_leaf().is_none());
+        assert_eq!(prev.text(), "=");
+        assert_eq!(leaf.text(), "10");
     }
 }