Add tt_from_syntax

DropDemBits · DropDemBits · commit 5fc5f63d09e5 · 2024-06-02T11:10:07.000-04:00
Used for inserting syntax nodes into existing token trees
diff --git a/src/tools/rust-analyzer/crates/ide-assists/src/utils.rs b/src/tools/rust-analyzer/crates/ide-assists/src/utils.rs
@@ -14,9 +14,9 @@ use syntax::{
         edit_in_place::{AttrsOwnerEdit, Indent, Removable},
         make, HasArgList, HasAttrs, HasGenericParams, HasName, HasTypeBounds, Whitespace,
     },
-    ted, AstNode, AstToken, Direction, SourceFile,
+    ted, AstNode, AstToken, Direction, NodeOrToken, SourceFile,
     SyntaxKind::*,
-    SyntaxNode, TextRange, TextSize, T,
+    SyntaxNode, SyntaxToken, TextRange, TextSize, T,
 };
 
 use crate::assist_context::{AssistContext, SourceChangeBuilder};
@@ -916,3 +916,46 @@ pub(crate) fn replace_record_field_expr(
         edit.replace(file_range.range, initializer.syntax().text());
     }
 }
+
+/// Creates a token tree list from a syntax node, creating the needed delimited sub token trees.
+/// Assumes that the input syntax node is a valid syntax tree.
+pub(crate) fn tt_from_syntax(node: SyntaxNode) -> Vec<NodeOrToken<ast::TokenTree, SyntaxToken>> {
+    let mut tt_stack = vec![(None, vec![])];
+
+    for element in node.descendants_with_tokens() {
+        let NodeOrToken::Token(token) = element else { continue };
+
+        match token.kind() {
+            T!['('] | T!['{'] | T!['['] => {
+                // Found an opening delimeter, start a new sub token tree
+                tt_stack.push((Some(token.kind()), vec![]));
+            }
+            T![')'] | T!['}'] | T![']'] => {
+                // Closing a subtree
+                let (delimiter, tt) = tt_stack.pop().expect("unbalanced delimeters");
+                let (_, parent_tt) = tt_stack
+                    .last_mut()
+                    .expect("parent token tree was closed before it was completed");
+                let closing_delimiter = delimiter.map(|it| match it {
+                    T!['('] => T![')'],
+                    T!['{'] => T!['}'],
+                    T!['['] => T![']'],
+                    _ => unreachable!(),
+                });
+                stdx::always!(
+                    closing_delimiter == Some(token.kind()),
+                    "mismatched opening and closing delimiters"
+                );
+
+                let sub_tt = make::token_tree(delimiter.expect("unbalanced delimiters"), tt);
+                parent_tt.push(NodeOrToken::Node(sub_tt));
+            }
+            _ => {
+                let (_, current_tt) = tt_stack.last_mut().expect("unmatched delimiters");
+                current_tt.push(NodeOrToken::Token(token))
+            }
+        }
+    }
+
+    tt_stack.pop().expect("parent token tree was closed before it was completed").1
+}
