feat(rome_formatter): Assert formatting of comments

This PR adds a debug assertion to verify if all comments in a program have been formatted to prevent accidental "dropping" of comments.

crates/rome_formatter/Cargo.toml

@@ -13,7 +13,6 @@ rome_rowan = { path = "../rome_rowan" }
 tracing = { version = "0.1.31", default-features = false, features = ["std"] }
 serde = { version = "1.0.136", features = ["derive"], optional = true }
 cfg-if = "1.0.0"
-indexmap = "1.8.2"
 schemars = { version = "0.8.10", optional = true }
 rustc-hash = "1.1.0"
 countme = "3.0.1"

crates/rome_formatter/src/comments.rs

@@ -84,7 +84,7 @@ use rome_rowan::syntax::SyntaxElementKey;
 use rome_rowan::{Language, SyntaxNode, SyntaxToken, SyntaxTriviaPieceComments};
 use rustc_hash::FxHashSet;
 #[cfg(debug_assertions)]
-use std::cell::RefCell;
+use std::cell::{Cell, RefCell};
 use std::rc::Rc;
 
 #[derive(Copy, Clone, Eq, PartialEq, Debug)]
@@ -164,6 +164,10 @@ pub struct SourceComment<L: Language> {
 
     /// The kind of the comment.
     pub(crate) kind: CommentKind,
+
+    /// Whether the comment has been formatted or not.
+    #[cfg(debug_assertions)]
+    pub(crate) formatted: Cell<bool>,
 }
 
 impl<L: Language> SourceComment<L> {
@@ -227,6 +231,16 @@ impl<L: Language> SourceComment<L> {
     pub fn kind(&self) -> CommentKind {
         self.kind
     }
+
+    #[cfg(not(debug_assertions))]
+    #[inline(always)]
+    pub fn mark_formatted(&self) {}
+
+    /// Marks the comment as formatted
+    #[cfg(debug_assertions)]
+    pub fn mark_formatted(&self) {
+        self.formatted.set(true)
+    }
 }
 
 /// A comment decorated with additional information about its surrounding context in the source document.
@@ -446,6 +460,8 @@ impl<L: Language> From<DecoratedComment<L>> for SourceComment<L> {
             lines_after: decorated.lines_after,
             piece: decorated.comment,
             kind: decorated.kind,
+            #[cfg(debug_assertions)]
+            formatted: Cell::new(false),
         }
     }
 }
@@ -912,53 +928,102 @@ impl<L: Language> Comments<L> {
             .any(|comment| is_suppression(comment.piece().text()))
     }
 
+    #[cfg(not(debug_assertions))]
+    #[inline(always)]
+    pub fn mark_suppression_checked(&self, _: &SyntaxNode<L>) {}
+
     /// Marks that it isn't necessary for the given node to check if it has been suppressed or not.
-    #[inline]
+    #[cfg(debug_assertions)]
     pub fn mark_suppression_checked(&self, node: &SyntaxNode<L>) {
-        cfg_if::cfg_if! {
-            if #[cfg(debug_assertions)] {
-                let mut checked_nodes = self.data.checked_suppressions.borrow_mut();
-                checked_nodes.insert(node.clone());
-            } else {
-                let _ = node;
-            }
-        }
+        let mut checked_nodes = self.data.checked_suppressions.borrow_mut();
+        checked_nodes.insert(node.clone());
     }
 
+    #[cfg(not(debug_assertions))]
+    #[inline(always)]
+    pub(crate) fn assert_checked_all_suppressions(&self, _: &SyntaxNode<L>) {}
+
     /// Verifies that [NodeSuppressions::is_suppressed] has been called for every node of `root`.
     /// This is a no-op in builds that have the feature `debug_assertions` disabled.
     ///
     /// # Panics
     /// If theres any node for which the formatting didn't very if it has a suppression comment.
-    #[inline]
+    #[cfg(debug_assertions)]
     pub(crate) fn assert_checked_all_suppressions(&self, root: &SyntaxNode<L>) {
-        cfg_if::cfg_if! {
-            if #[cfg(debug_assertions)] {
-                use rome_rowan::SyntaxKind;
-
-                let checked_nodes = self.data.checked_suppressions.borrow();
-                for node in root.descendants() {
-                    if node.kind().is_list() || node.kind().is_root() {
-                        continue;
-                    }
-
-                    if !checked_nodes.contains(&node) {
-                        panic!(r#"
+        use rome_rowan::SyntaxKind;
+
+        let checked_nodes = self.data.checked_suppressions.borrow();
+        for node in root.descendants() {
+            if node.kind().is_list() || node.kind().is_root() {
+                continue;
+            }
+
+            if !checked_nodes.contains(&node) {
+                panic!(
+                    r#"
 The following node has been formatted without checking if it has suppression comments.
 Ensure that the formatter calls into the node's formatting rule by using `node.format()` or
 manually test if the node has a suppression comment using `f.context().comments().is_suppressed(node.syntax())`
 if using the node's format rule isn't an option."
 
 Node:
 {node:#?}"#
-                        );
-                    }
-                }
-            } else {
-                let _ = root;
+                );
             }
         }
     }
+
+    #[inline(always)]
+    #[cfg(not(debug_assertions))]
+    pub(crate) fn assert_formatted_all_comments(&self) {}
+
+    #[cfg(debug_assertions)]
+    pub(crate) fn assert_formatted_all_comments(&self) {
+        let has_unformatted_comments = self
+            .data
+            .comments
+            .all_parts()
+            .any(|comment| !comment.formatted.get());
+
+        if has_unformatted_comments {
+            let mut unformatted_comments = Vec::new();
+
+            for node in self
+                .data
+                .root
+                .as_ref()
+                .expect("Expected root for comments with data")
+                .descendants()
+            {
+                unformatted_comments.extend(self.leading_comments(&node).iter().filter_map(
+                    |comment| {
+                        (!comment.formatted.get()).then_some(DebugComment::Leading {
+                            node: node.clone(),
+                            comment,
+                        })
+                    },
+                ));
+                unformatted_comments.extend(self.dangling_comments(&node).iter().filter_map(
+                    |comment| {
+                        (!comment.formatted.get()).then_some(DebugComment::Dangling {
+                            node: node.clone(),
+                            comment,
+                        })
+                    },
+                ));
+                unformatted_comments.extend(self.trailing_comments(&node).iter().filter_map(
+                    |comment| {
+                        (!comment.formatted.get()).then_some(DebugComment::Trailing {
+                            node: node.clone(),
+                            comment,
+                        })
+                    },
+                ));
+            }
+
+            panic!("The following comments have not been formatted.\n{unformatted_comments:#?}")
+        }
+    }
 }
 
 struct CommentsData<L: Language> {

crates/rome_formatter/src/comments/map.rs

@@ -243,14 +243,17 @@ impl<K: std::hash::Hash + Eq, V> CommentsMap<K, V> {
     pub fn parts(&self, key: &K) -> PartsIterator<V> {
         match self.index.get(key) {
             None => PartsIterator::Slice([].iter()),
-            Some(Entry::OutOfOrder(entry)) => PartsIterator::Leading {
-                leading: self.out_of_order[entry.leading_index()].iter(),
-                dangling: &self.out_of_order[entry.dangling_index()],
-                trailing: &self.out_of_order[entry.trailing_index()],
-            },
-            Some(Entry::InOrder(entry)) => PartsIterator::Slice(self.parts[entry.range()].iter()),
+            Some(entry) => PartsIterator::from_entry(entry, self),
         }
     }
+
+    /// Returns an iterator over the parts of all keys.
+    #[allow(unused)]
+    pub fn all_parts(&self) -> impl Iterator<Item = &V> {
+        self.index
+            .values()
+            .flat_map(|entry| PartsIterator::from_entry(entry, self))
+    }
 }
 
 impl<K: std::hash::Hash + Eq, V> Default for CommentsMap<K, V> {
@@ -297,6 +300,19 @@ pub(super) enum PartsIterator<'a, V> {
     },
 }
 
+impl<'a, V> PartsIterator<'a, V> {
+    fn from_entry<K>(entry: &Entry, map: &'a CommentsMap<K, V>) -> Self {
+        match entry {
+            Entry::OutOfOrder(entry) => PartsIterator::Leading {
+                leading: map.out_of_order[entry.leading_index()].iter(),
+                dangling: &map.out_of_order[entry.dangling_index()],
+                trailing: &map.out_of_order[entry.trailing_index()],
+            },
+            Entry::InOrder(entry) => PartsIterator::Slice(map.parts[entry.range()].iter()),
+        }
+    }
+}
+
 impl<'a, V> Iterator for PartsIterator<'a, V> {
     type Item = &'a V;
 

crates/rome_formatter/src/format_element.rs

@@ -6,10 +6,10 @@ use crate::{
     FormatResult, Formatter, GroupId, IndentStyle, LineWidth, PrinterOptions, TextSize,
     TransformSourceMap,
 };
-use indexmap::IndexSet;
 #[cfg(target_pointer_width = "64")]
 use rome_rowan::static_assert;
 use rome_rowan::SyntaxTokenText;
+use rustc_hash::FxHashMap;
 #[cfg(debug_assertions)]
 use std::any::type_name;
 use std::any::TypeId;
@@ -757,7 +757,7 @@ impl From<ConditionalGroupContent> for FormatElement {
 #[derive(Clone, Default, Debug)]
 struct IrFormatContext {
     /// The interned elements that have been printed to this point
-    printed_interned_elements: IndexSet<Interned>,
+    printed_interned_elements: FxHashMap<Interned, usize>,
 }
 
 impl FormatContext for IrFormatContext {
@@ -964,10 +964,16 @@ impl Format<IrFormatContext> for FormatElement {
                 )
             }
             FormatElement::Interned(interned) => {
-                let (index, inserted) = f
-                    .context_mut()
-                    .printed_interned_elements
-                    .insert_full(interned.clone());
+                let interned_elements = &mut f.context_mut().printed_interned_elements;
+
+                let (index, inserted) = match interned_elements.get(interned) {
+                    None => {
+                        let index = interned_elements.len();
+                        interned_elements.insert(interned.clone(), index);
+                        (index, true)
+                    }
+                    Some(index) => (*index, false),
+                };
 
                 if inserted {
                     write!(

crates/rome_formatter/src/lib.rs

@@ -882,10 +882,12 @@ pub fn format_node<L: FormatLanguage>(
         let document = buffer.into_element();
 
         state.assert_formatted_all_tokens(&root);
-        state
-            .context()
-            .comments()
-            .assert_checked_all_suppressions(&root);
+        {
+            let comments = state.context().comments();
+
+            comments.assert_checked_all_suppressions(&root);
+            comments.assert_formatted_all_comments();
+        }
 
         Ok(Formatted::new(document, state.into_context()))
     })

crates/rome_formatter/src/trivia.rs

@@ -7,6 +7,8 @@ use crate::{
     TextRange, VecBuffer,
 };
 use rome_rowan::{Language, SyntaxNode, SyntaxToken};
+#[cfg(debug_assertions)]
+use std::cell::Cell;
 
 /// Formats the leading comments of `node`
 pub const fn format_leading_comments<L: Language>(
@@ -57,6 +59,8 @@ where
                     _ => write!(f, [empty_line()])?,
                 },
             }
+
+            comment.mark_formatted()
         }
 
         Ok(())
@@ -129,6 +133,8 @@ where
                     write!(f, [content])?;
                 }
             }
+
+            comment.mark_formatted();
         }
 
         Ok(())
@@ -247,6 +253,8 @@ where
                 let format_comment =
                     FormatRefWithRule::new(comment, Context::CommentRule::default());
                 join.entry(&format_comment);
+
+                comment.mark_formatted();
             }
 
             join.finish()?;
@@ -484,6 +492,8 @@ impl<L: Language> FormatSkippedTokenTrivia<'_, L> {
                     lines_before: lines,
                     lines_after: 0,
                     piece: comment,
+                    #[cfg(debug_assertions)]
+                    formatted: Cell::new(true),
                 };
 
                 dangling_comments.push(source_comment);