perf(rowan): Speedup next_token and prev_token implementations

Hand write the traversal of `next_token` and `prev_token` to speedup the traversal. The new comment formatting will make heavy use of these two methods which is why it's important that they are fast. ``` ❯ critcmp main fused group PR main ----- ----- ---- formatter/checker.ts 1.00 239.2±1.22ms 10.9 MB/sec 1.01 240.8±3.49ms 10.8 MB/sec formatter/compiler.js 1.00 139.6±3.00ms 7.5 MB/sec 1.04 145.4±1.99ms 7.2 MB/sec formatter/d3.min.js 1.00 104.7±1.04ms 2.5 MB/sec 1.04 108.9±2.69ms 2.4 MB/sec formatter/dojo.js 1.00 7.8±0.11ms 8.8 MB/sec 1.02 7.9±0.01ms 8.6 MB/sec formatter/ios.d.ts 1.00 189.0±3.27ms 9.9 MB/sec 1.01 190.0±1.41ms 9.8 MB/sec formatter/jquery.min.js 1.00 29.5±0.43ms 2.8 MB/sec 1.00 29.4±0.23ms 2.8 MB/sec formatter/math.js 1.00 225.6±0.95ms 2.9 MB/sec 1.04 235.4±4.01ms 2.8 MB/sec formatter/parser.ts 1.02 5.6±0.01ms 8.7 MB/sec 1.00 5.5±0.09ms 8.8 MB/sec formatter/pixi.min.js 1.00 120.9±1.96ms 3.6 MB/sec 1.12 135.3±19.62ms 3.2 MB/sec formatter/react-dom.production.min.js 1.00 35.4±0.54ms 3.2 MB/sec 1.05 37.3±0.39ms 3.1 MB/sec formatter/react.production.min.js 1.00 1815.8±3.00µs 3.4 MB/sec 1.01 1833.6±29.27µs 3.4 MB/sec formatter/router.ts 1.01 4.1±0.01ms 14.8 MB/sec 1.00 4.1±0.08ms 15.0 MB/sec formatter/tex-chtml-full.js 1.00 285.8±1.28ms 3.2 MB/sec 1.04 297.2±1.77ms 3.1 MB/sec formatter/three.min.js 1.00 142.0±1.02ms 4.1 MB/sec 1.05 148.7±4.29ms 3.9 MB/sec formatter/typescript.js 1.00 957.2±94.68ms 9.9 MB/sec 1.01 967.2±20.72ms 9.8 MB/sec formatter/vue.global.prod.js 1.00 47.7±0.19ms 2.5 MB/sec 1.02 48.7±0.39ms 2.5 MB/sec ```
rome · May 25, 2022 · f9fce38 · f9fce38
1 parent a32588d
commit f9fce38
Show file tree

Hide file tree

Showing 3 changed files with 74 additions and 24 deletions.
diff --git a/crates/rome_rowan/src/cursor/node.rs b/crates/rome_rowan/src/cursor/node.rs
@@ -461,6 +461,8 @@ impl Iterator for SyntaxNodeChildren {
     }
 }
 
+impl FusedIterator for SyntaxNodeChildren {}
+
 #[derive(Clone, Debug, Default)]
 pub(crate) struct SyntaxElementChildren {
     next: Option<SyntaxElement>,
@@ -484,6 +486,8 @@ impl Iterator for SyntaxElementChildren {
     }
 }
 
+impl FusedIterator for SyntaxElementChildren {}
+
 pub(crate) struct Preorder {
     start: SyntaxNode,
     next: Option<WalkEvent<SyntaxNode>>,
@@ -543,6 +547,8 @@ impl Iterator for Preorder {
     }
 }
 
+impl FusedIterator for Preorder {}
+
 pub(crate) struct PreorderWithTokens {
     start: SyntaxElement,
     next: Option<WalkEvent<SyntaxElement>>,
@@ -626,6 +632,8 @@ impl Iterator for PreorderWithTokens {
     }
 }
 
+impl FusedIterator for PreorderWithTokens {}
+
 /// Represents a cursor to a green node slot. A slot either contains an element or is empty
 /// if the child isn't present in the source.
 #[derive(Debug, Clone)]
@@ -681,7 +689,7 @@ impl SyntaxSlots {
     }
 }
 
-impl<'a> Iterator for SyntaxSlots {
+impl Iterator for SyntaxSlots {
     type Item = SyntaxSlot;
 
     fn next(&mut self) -> Option<Self::Item> {
@@ -786,6 +794,8 @@ impl Iterator for SlotsPreorder {
     }
 }
 
+impl FusedIterator for SlotsPreorder {}
+
 #[derive(Debug, Clone)]
 pub(crate) struct Siblings<'a> {
     parent: &'a GreenNodeData,

diff --git a/crates/rome_rowan/src/cursor/token.rs b/crates/rome_rowan/src/cursor/token.rs
@@ -1,6 +1,8 @@
 use crate::cursor::{NodeData, SyntaxElement, SyntaxNode, SyntaxTrivia};
 use crate::green::GreenElementRef;
-use crate::{green, Direction, GreenToken, GreenTokenData, RawSyntaxKind, SyntaxTokenText};
+use crate::{
+    green, Direction, GreenToken, GreenTokenData, RawSyntaxKind, SyntaxTokenText, WalkEvent,
+};
 use std::hash::{Hash, Hasher};
 use std::rc::Rc;
 use std::{fmt, iter};
@@ -140,30 +142,65 @@ impl SyntaxToken {
     }
 
     pub fn next_token(&self) -> Option<SyntaxToken> {
-        iter::successors(
-            self.next_sibling_or_token(),
-            SyntaxElement::next_sibling_or_token,
-        )
-        .chain(self.ancestors().flat_map(|node| {
-            iter::successors(
-                node.next_sibling_or_token(),
-                SyntaxElement::next_sibling_or_token,
-            )
-        }))
-        .find_map(|element| element.first_token())
+        self.next_token_impl(Direction::Next)
     }
+
     pub fn prev_token(&self) -> Option<SyntaxToken> {
-        iter::successors(
-            self.prev_sibling_or_token(),
-            SyntaxElement::prev_sibling_or_token,
-        )
-        .chain(self.ancestors().flat_map(|node| {
-            iter::successors(
-                node.prev_sibling_or_token(),
-                SyntaxElement::prev_sibling_or_token,
-            )
-        }))
-        .find_map(|element| element.last_token())
+        self.next_token_impl(Direction::Prev)
+    }
+
+    /// Returns the token preceding or following this token depending on the passed `direction`.
+    fn next_token_impl(&self, direction: Direction) -> Option<SyntaxToken> {
+        let mut current: WalkEvent<SyntaxElement> =
+            WalkEvent::Leave(SyntaxElement::Token(self.clone()));
+
+        loop {
+            current = match current {
+                WalkEvent::Enter(element) => match element {
+                    SyntaxElement::Token(token) => break Some(token),
+                    SyntaxElement::Node(node) => {
+                        let first_child = match direction {
+                            Direction::Next => node.first_child_or_token(),
+                            Direction::Prev => node.last_child_or_token(),
+                        };
+
+                        match first_child {
+                            // If node is empty, leave parent
+                            None => WalkEvent::Leave(SyntaxElement::Node(node)),
+                            // Otherwise traverse full sub-tree
+                            Some(child) => WalkEvent::Enter(child),
+                        }
+                    }
+                },
+                WalkEvent::Leave(element) => {
+                    let mut current_element = element;
+
+                    loop {
+                        // Only traverse the left (pref) / right (next) siblings of the parent
+                        // to avoid traversing into the same children again.
+                        let sibling = match direction {
+                            Direction::Next => current_element.next_sibling_or_token(),
+                            Direction::Prev => current_element.prev_sibling_or_token(),
+                        };
+
+                        match sibling {
+                            // Traverse all children of the sibling
+                            Some(sibling) => break WalkEvent::Enter(sibling),
+                            None => {
+                                match current_element.parent() {
+                                    Some(node) => {
+                                        current_element = SyntaxElement::Node(node);
+                                    }
+                                    None => {
+                                        return None; // Reached root, no token found
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        }
     }
 
     #[must_use = "syntax elements are immutable, the result of update methods must be propagated to have any effect"]

diff --git a/crates/rome_rowan/src/cursor/trivia.rs b/crates/rome_rowan/src/cursor/trivia.rs
@@ -2,6 +2,7 @@ use crate::cursor::SyntaxToken;
 use crate::green::GreenTrivia;
 use crate::TriviaPiece;
 use std::fmt;
+use std::iter::FusedIterator;
 use text_size::{TextRange, TextSize};
 
 #[derive(PartialEq, Eq, Clone, Hash)]
@@ -129,6 +130,8 @@ impl Iterator for SyntaxTriviaPiecesIterator {
     }
 }
 
+impl FusedIterator for SyntaxTriviaPiecesIterator {}
+
 impl DoubleEndedIterator for SyntaxTriviaPiecesIterator {
     fn next_back(&mut self) -> Option<Self::Item> {
         if self.end_index == self.next_index {