Add side parameter to leaf_at (#3767)

5 files changed, 69 insertions, 16 deletions

diff --git a/crates/typst-ide/src/complete.rs b/crates/typst-ide/src/complete.rs
index 4e4b8918..be28a431 100644
--- a/crates/typst-ide/src/complete.rs
+++ b/crates/typst-ide/src/complete.rs
@@ -10,7 +10,7 @@ use typst::foundations::{
 };
 use typst::model::Document;
 use typst::syntax::{
-    ast, is_id_continue, is_id_start, is_ident, LinkedNode, Source, SyntaxKind,
+    ast, is_id_continue, is_id_start, is_ident, LinkedNode, Side, Source, SyntaxKind,
 };
 use typst::text::RawElem;
 use typst::visualize::Color;
@@ -1033,7 +1033,7 @@ impl<'a> CompletionContext<'a> {
     ) -> Option<Self> {
         let text = source.text();
         let library = world.library();
-        let leaf = LinkedNode::new(source.root()).leaf_at(cursor)?;
+        let leaf = LinkedNode::new(source.root()).leaf_at(cursor, Side::Before)?;
         Some(Self {
             world,
             document,
diff --git a/crates/typst-ide/src/jump.rs b/crates/typst-ide/src/jump.rs
index c8b7343a..7c3c7569 100644
--- a/crates/typst-ide/src/jump.rs
+++ b/crates/typst-ide/src/jump.rs
@@ -4,7 +4,7 @@ use ecow::EcoString;
 use typst::introspection::Meta;
 use typst::layout::{Frame, FrameItem, Point, Position, Size};
 use typst::model::{Destination, Document};
-use typst::syntax::{FileId, LinkedNode, Source, Span, SyntaxKind};
+use typst::syntax::{FileId, LinkedNode, Side, Source, Span, SyntaxKind};
 use typst::visualize::Geometry;
 use typst::World;
 
@@ -115,11 +115,16 @@ pub fn jump_from_cursor(
     source: &Source,
     cursor: usize,
 ) -> Option<Position> {
-    let node = LinkedNode::new(source.root()).leaf_at(cursor)?;
-    if node.kind() != SyntaxKind::Text {
-        return None;
+    fn is_text(node: &LinkedNode) -> bool {
+        node.get().kind() == SyntaxKind::Text
     }
 
+    let root = LinkedNode::new(source.root());
+    let node = root
+        .leaf_at(cursor, Side::Before)
+        .filter(is_text)
+        .or_else(|| root.leaf_at(cursor, Side::After).filter(is_text))?;
+
     let span = node.span();
     for (i, page) in document.pages.iter().enumerate() {
         if let Some(pos) = find_in_frame(&page.frame, span) {
diff --git a/crates/typst-ide/src/tooltip.rs b/crates/typst-ide/src/tooltip.rs
index 2f04be87..3416e5f8 100644
--- a/crates/typst-ide/src/tooltip.rs
+++ b/crates/typst-ide/src/tooltip.rs
@@ -6,7 +6,7 @@ use typst::eval::{CapturesVisitor, Tracer};
 use typst::foundations::{repr, Capturer, CastInfo, Repr, Value};
 use typst::layout::Length;
 use typst::model::Document;
-use typst::syntax::{ast, LinkedNode, Source, SyntaxKind};
+use typst::syntax::{ast, LinkedNode, Side, Source, SyntaxKind};
 use typst::util::{round_2, Numeric};
 use typst::World;
 
@@ -23,8 +23,9 @@ pub fn tooltip(
     document: Option<&Document>,
     source: &Source,
     cursor: usize,
+    side: Side,
 ) -> Option<Tooltip> {
-    let leaf = LinkedNode::new(source.root()).leaf_at(cursor)?;
+    let leaf = LinkedNode::new(source.root()).leaf_at(cursor, side)?;
     if leaf.kind().is_trivia() {
         return None;
     }
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");
     }
 }