Finish call args

crates/rome_formatter/Cargo.toml

@@ -16,6 +16,7 @@ cfg-if = "1.0.0"
 schemars = { version = "0.8.10", optional = true }
 rustc-hash = "1.1.0"
 countme = "3.0.1"
+indexmap = "1.9.1"
 
 [dev-dependencies]
 rome_js_parser = { path = "../rome_js_parser"}

crates/rome_formatter/src/lib.rs

@@ -1494,11 +1494,22 @@ impl<Context> FormatState<Context> {
 
     #[cfg(not(debug_assertions))]
     #[inline]
-    pub fn remove_tracked_token<L: Language>(&self, _: SyntaxToken<L>) {}
+    pub fn set_token_tracking_enabled(&mut self, _: bool) {}
 
     #[cfg(debug_assertions)]
-    pub fn remove_tracked_token<L: Language>(&mut self, token: &SyntaxToken<L>) {
-        self.printed_tokens.remove_tracked_token(token);
+    pub fn set_token_tracking_enabled(&mut self, enabled: bool) {
+        self.printed_tokens.set_enabled(enabled)
+    }
+
+    #[cfg(not(debug_assertions))]
+    #[inline]
+    pub fn is_token_tracking_enabled(&self) -> bool {
+        false
+    }
+
+    #[cfg(debug_assertions)]
+    pub fn is_token_tracking_enabled(&self) -> bool {
+        self.printed_tokens.is_enabled()
     }
 
     /// Asserts in debug builds that all tokens have been printed.
@@ -1522,7 +1533,7 @@ where
     pub fn snapshot(&self) -> FormatStateSnapshot {
         FormatStateSnapshot {
             #[cfg(debug_assertions)]
-            printed_tokens: self.printed_tokens.clone(),
+            printed_tokens: self.printed_tokens.snapshot(),
         }
     }
 
@@ -1534,13 +1545,13 @@ where
 
         cfg_if::cfg_if! {
             if #[cfg(debug_assertions)] {
-                self.printed_tokens = printed_tokens;
+                self.printed_tokens.restore(printed_tokens);
             }
         }
     }
 }
 
 pub struct FormatStateSnapshot {
     #[cfg(debug_assertions)]
-    printed_tokens: PrintedTokens,
+    printed_tokens: printed_tokens::PrintedTokensSnapshot,
 }

crates/rome_formatter/src/printed_tokens.rs

@@ -1,5 +1,5 @@
+use indexmap::IndexSet;
 use rome_rowan::{Direction, Language, SyntaxNode, SyntaxToken, TextSize};
-use std::collections::BTreeSet;
 
 /// Tracks the ranges of the formatted (including replaced or tokens formatted as verbatim) tokens.
 ///
@@ -8,7 +8,14 @@ use std::collections::BTreeSet;
 #[derive(Debug, Clone, Default)]
 pub struct PrintedTokens {
     /// Key: Start of a token's range
-    offsets: BTreeSet<TextSize>,
+    offsets: IndexSet<TextSize>,
+    disabled: bool,
+}
+
+#[derive(Copy, Clone)]
+pub struct PrintedTokensSnapshot {
+    len: usize,
+    disabled: bool,
 }
 
 impl PrintedTokens {
@@ -17,48 +24,58 @@ impl PrintedTokens {
     /// ## Panics
     /// If this token has been formatted before.
     pub fn track_token<L: Language>(&mut self, token: &SyntaxToken<L>) {
+        if self.disabled {
+            return;
+        }
+
         let range = token.text_trimmed_range();
 
         if !self.offsets.insert(range.start()) {
             panic!("You tried to print the token '{token:?}' twice, and this is not valid.");
         }
     }
 
-    /// Removes a token if it has been tracked before. No-op otherwise.
-    pub(crate) fn remove_tracked_token<L: Language>(&mut self, token: &SyntaxToken<L>) {
-        self.offsets.remove(&token.text_trimmed_range().start());
+    /// Enables or disables the assertion tracking
+    pub(crate) fn set_enabled(&mut self, enabled: bool) {
+        self.disabled = !enabled;
+    }
+
+    pub(crate) fn is_enabled(&self) -> bool {
+        !self.disabled
+    }
+
+    pub(crate) fn snapshot(&self) -> PrintedTokensSnapshot {
+        PrintedTokensSnapshot {
+            len: self.offsets.len(),
+            disabled: self.disabled,
+        }
+    }
+
+    pub(crate) fn restore(&mut self, snapshot: PrintedTokensSnapshot) {
+        let PrintedTokensSnapshot { len, disabled } = snapshot;
+
+        self.offsets.truncate(len);
+        self.disabled = disabled
     }
 
     /// Asserts that all tokens of the passed in node have been tracked
     ///
     /// ## Panics
     /// If any descendant token of `root` hasn't been tracked
     pub fn assert_all_tracked<L: Language>(&self, root: &SyntaxNode<L>) {
-        let mut descendants = root.descendants_tokens(Direction::Next);
-        let mut offsets = self.offsets.iter();
+        let mut offsets = self.offsets.clone();
 
-        loop {
-            match (descendants.next(), offsets.next()) {
-                (Some(descendant), Some(offset)) => match descendant.text_trimmed_range().start() {
-                    descendant_offset if descendant_offset < *offset => {
-                        panic!("token has not been seen by the formatter: {descendant:#?}.\
+        for token in root.descendants_tokens(Direction::Next) {
+            if !offsets.remove(&token.text_trimmed_range().start()) {
+                panic!("token has not been seen by the formatter: {token:#?}.\
                         \nUse `format_replaced` if you want to replace a token from the formatted output.\
                         \nUse `format_removed` if you want to remove a token from the formatted output.\n\
-                        parent: {:#?}", descendant.parent())
-                    }
-                    descendant_offset if descendant_offset > *offset => {
-                        panic!("tracked offset {offset:?} doesn't match any token of {root:#?}. Have you passed a token from another tree?");
-                    }
-                    _ => {}
-                },
-                (Some(descendant), None) => {
-                    panic!("token has not been seen by the formatter: {descendant:#?}.\n Use `formatter.format_replaced` if you intentionally remove or replace a token from the formatted output.")
-                }
-                (None, Some(offset)) => {
-                    panic!("tracked offset {offset:?} doesn't match any token of {root:#?}. Have you passed a token from another tree?");
-                }
-                (None, None) => break,
-            };
+                        parent: {:#?}", token.parent())
+            }
+        }
+
+        for offset in offsets {
+            panic!("tracked offset {offset:?} doesn't match any token of {root:#?}. Have you passed a token from another tree?");
         }
     }
 }

crates/rome_js_formatter/src/js/declarations/function_declaration.rs

@@ -137,54 +137,49 @@ impl FormatFunction {
         let parameters = self.parameters()?;
         let return_type_annotation = self.return_type_annotation();
 
-        if expand {
-            write!(
-                f,
-                [
-                    group(&format_with(|f| {
-                        let mut buffer = RemoveSoftLinesBuffer::new(f);
-
-                        let mut recording = buffer.start_recording();
-                        write!(recording, [type_parameters.format(), parameters.format()])?;
-                        let recorded = recording.stop();
-
-                        if recorded.will_break() {
-                            return Err(FormatError::PoorLayout);
-                        } else {
-                            Ok(())
-                        }
-                    })),
-                    return_type_annotation.format()
-                ]
-            )?
-        } else {
-            write!(f, [type_parameters.format()])?;
-
-            write!(
-                f,
-                [group(&format_with(|f| {
-                    let mut format_return_type_annotation =
-                        return_type_annotation.format().memoized();
-                    let group_parameters = should_group_function_parameters(
-                        type_parameters.as_ref(),
-                        parameters.items().len(),
-                        return_type_annotation
-                            .as_ref()
-                            .map(|annotation| annotation.ty()),
-                        &mut format_return_type_annotation,
-                        f,
-                    )?;
-
-                    if group_parameters {
-                        write!(f, [group(&parameters.format())])?;
-                    } else {
-                        write!(f, [parameters.format()])?;
-                    }
-
-                    write![f, [format_return_type_annotation]]
-                }))]
-            )?;
-        }
+        write!(f, [type_parameters.format()])?;
+
+        let format_parameters = format_with(|f| {
+            if expand {
+                let mut buffer = RemoveSoftLinesBuffer::new(f);
+
+                let mut recording = buffer.start_recording();
+                write!(recording, [parameters.format()])?;
+                let recorded = recording.stop();
+
+                if recorded.will_break() {
+                    return Err(FormatError::PoorLayout);
+                } else {
+                    Ok(())
+                }
+            } else {
+                parameters.format().fmt(f)
+            }
+        });
+
+        write!(
+            f,
+            [group(&format_with(|f| {
+                let mut format_return_type_annotation = return_type_annotation.format().memoized();
+                let group_parameters = should_group_function_parameters(
+                    type_parameters.as_ref(),
+                    parameters.items().len(),
+                    return_type_annotation
+                        .as_ref()
+                        .map(|annotation| annotation.ty()),
+                    &mut format_return_type_annotation,
+                    f,
+                )?;
+
+                if group_parameters {
+                    write!(f, [group(&format_parameters)])?;
+                } else {
+                    write!(f, [format_parameters])?;
+                }
+
+                write!(f, [format_return_type_annotation])
+            }))]
+        )?;
 
         if let Some(body) = self.body()? {
             write!(f, [space(), body.format()])?;