API to binary search child by range

Author: matklad

Cargo.toml

@@ -10,7 +10,7 @@ edition = "2018"
 [dependencies]
 rustc-hash = "1.0.1"
 smol_str = "0.1.10"
-text-size = "1.0.0"
+text-size = "1.1.0"
 triomphe = "0.1.1"
 
 serde = { version = "1.0.89", optional = true, default-features = false }

src/api.rs

@@ -274,6 +274,10 @@ impl<L: Language> SyntaxNode<L> {
     pub fn covering_element(&self, range: TextRange) -> SyntaxElement<L> {
         NodeOrToken::from(self.raw.covering_element(range))
     }
+
+    pub fn child_or_token_at_range(&self, range: TextRange) -> Option<SyntaxElement<L>> {
+        self.raw.child_or_token_at_range(range).map(SyntaxElement::from)
+    }
 }
 
 impl<L: Language> SyntaxToken<L> {

src/cursor.rs

@@ -470,18 +470,29 @@ impl SyntaxNode {
             );
             res = match &res {
                 NodeOrToken::Token(_) => return res,
-                NodeOrToken::Node(node) => {
-                    match node
-                        .children_with_tokens()
-                        .find(|child| child.text_range().contains_range(range))
-                    {
-                        Some(child) => child,
-                        None => return res,
-                    }
-                }
+                NodeOrToken::Node(node) => match node.child_or_token_at_range(range) {
+                    Some(it) => it,
+                    None => return res,
+                },
             };
         }
     }
+
+    pub fn child_or_token_at_range(&self, range: TextRange) -> Option<SyntaxElement> {
+        let start_offset = self.text_range().start();
+        let (index, offset, child) = self.green().child_at_range(range - start_offset)?;
+        let index = index as u32;
+        let offset = offset + start_offset;
+        let res: SyntaxElement = match child {
+            NodeOrToken::Node(node) => {
+                let data =
+                    NodeData::new(Kind::Child { parent: self.clone(), index, offset }, node.into());
+                SyntaxNode::new(data).into()
+            }
+            NodeOrToken::Token(_token) => SyntaxToken::new(self.clone(), index, offset).into(),
+        };
+        Some(res)
+    }
 }
 
 impl SyntaxToken {

src/green/node.rs

@@ -4,7 +4,7 @@ use triomphe::{Arc, ThinArc};
 
 use crate::{
     green::{GreenElement, GreenElementRef, SyntaxKind},
-    GreenToken, NodeOrToken, TextSize,
+    GreenToken, NodeOrToken, TextRange, TextSize,
 };
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -43,6 +43,16 @@ impl GreenChild {
             GreenChild::Token { token, .. } => NodeOrToken::Token(token),
         }
     }
+    fn offset_in_parent(&self) -> TextSize {
+        match self {
+            GreenChild::Node { offset_in_parent, .. }
+            | GreenChild::Token { offset_in_parent, .. } => *offset_in_parent,
+        }
+    }
+    fn range_in_parent(&self) -> TextRange {
+        let len = self.as_ref().text_len();
+        TextRange::at(self.offset_in_parent(), len)
+    }
 }
 
 #[cfg(target_pointer_width = "64")]
@@ -101,6 +111,24 @@ impl GreenNode {
         Children { inner: self.data.slice.iter() }
     }
 
+    pub(crate) fn child_at_range(
+        &self,
+        range: TextRange,
+    ) -> Option<(usize, TextSize, GreenElementRef<'_>)> {
+        let idx = self
+            .data
+            .slice
+            .binary_search_by(|it| {
+                let child_range = it.range_in_parent();
+                TextRange::ordering(child_range, range)
+            })
+            // XXX: this handles empty ranges
+            .unwrap_or_else(|it| it.saturating_sub(1));
+        let child =
+            &self.data.slice.get(idx).filter(|it| it.range_in_parent().contains_range(range))?;
+        Some((idx, child.offset_in_parent(), child.as_ref()))
+    }
+
     pub fn ptr(&self) -> *const c_void {
         self.data.heap_ptr()
     }