Enforce and make use of span ordering

2 files changed, 42 insertions, 17 deletions

diff --git a/src/syntax/mod.rs b/src/syntax/mod.rs
index 6cecfd2a..0791a761 100644
--- a/src/syntax/mod.rs
+++ b/src/syntax/mod.rs
@@ -62,7 +62,9 @@ impl SyntaxNode {
     pub fn range(&self, span: Span, offset: usize) -> Option<Range<usize>> {
         match self {
             SyntaxNode::Inner(inner) => inner.range(span, offset),
-            SyntaxNode::Leaf(leaf) => leaf.range(span, offset),
+            SyntaxNode::Leaf(leaf) => {
+                (span == leaf.span).then(|| offset .. offset + self.len())
+            }
         }
     }
 
@@ -242,13 +244,30 @@ impl InnerNode {
 
     /// If the span points into this node, convert it to a byte range.
     pub fn range(&self, span: Span, mut offset: usize) -> Option<Range<usize>> {
-        if let Some(range) = self.data.range(span, offset) {
-            return Some(range);
+        // Check whether we found it.
+        if self.data.span == span {
+            return Some(offset .. offset + self.len());
+        }
+
+        // The parent of a subtree has a smaller span number than all of its
+        // descendants. Therefore, we can bail out early if the target span's
+        // number is smaller than our number.
+        if span.number() < self.span().number() {
+            return None;
         }
 
-        for child in &self.children {
-            if let Some(range) = child.range(span, offset) {
-                return Some(range);
+        let mut children = self.children.iter().peekable();
+        while let Some(child) = children.next() {
+            // Every node in this child's subtree has a smaller span number than
+            // the next sibling. Therefore we only need to recurse if the next
+            // sibling's span number is larger than the target span's number.
+            if children
+                .peek()
+                .map_or(true, |next| next.span().number() > span.number())
+            {
+                if let Some(range) = child.range(span, offset) {
+                    return Some(range);
+                }
             }
 
             offset += child.len();
@@ -388,11 +407,6 @@ impl NodeData {
         self.span
     }
 
-    /// If the span points into this node, convert it to a byte range.
-    pub fn range(&self, span: Span, offset: usize) -> Option<Range<usize>> {
-        (self.span == span).then(|| offset .. offset + self.len())
-    }
-
     /// Assign spans to each node.
     pub fn number(&mut self, id: SourceId, from: u64) {
         self.span = Span::new(id, from);
diff --git a/src/syntax/span.rs b/src/syntax/span.rs
index 3f6e6824..496d699c 100644
--- a/src/syntax/span.rs
+++ b/src/syntax/span.rs
@@ -58,26 +58,37 @@ impl<T: Debug> Debug for Spanned<T> {
 pub struct Span(NonZeroU64);
 
 impl Span {
+    // Number of bits for and minimum and maximum numbers assigned to nodes.
+    const BITS: usize = 48;
+    const DETACHED: u64 = 1;
+    pub(crate) const MIN_NUMBER: u64 = 2;
+    pub(crate) const MAX_NUMBER: u64 = (1 << Self::BITS) - 1;
+
     /// Create a new span from a source id and a unique number.
     pub const fn new(id: SourceId, number: u64) -> Self {
-        assert!(number > 0 && number < (1 << 48));
-        let bits = ((id.into_raw() as u64) << 48) | number;
-        Self(nonzero(bits))
+        assert!(number >= Self::MIN_NUMBER && number <= Self::MAX_NUMBER);
+        let bits = ((id.into_raw() as u64) << Self::BITS) | number;
+        Self(convert(bits))
     }
 
     /// A node that does not belong to any source file.
     pub const fn detached() -> Self {
-        Self(nonzero(1))
+        Self(convert(Self::DETACHED))
     }
 
     /// The id of the source file the span points into.
     pub const fn source(self) -> SourceId {
-        SourceId::from_raw((self.0.get() >> 48) as u16)
+        SourceId::from_raw((self.0.get() >> Self::BITS) as u16)
+    }
+
+    /// The unique number of the span within the source file.
+    pub const fn number(self) -> u64 {
+        self.0.get() & Self::MAX_NUMBER
     }
 }
 
 /// Convert to a non zero u64.
-const fn nonzero(v: u64) -> NonZeroU64 {
+const fn convert(v: u64) -> NonZeroU64 {
     match NonZeroU64::new(v) {
         Some(v) => v,
         None => unreachable!(),