From 2d6dc1cc9778bea161f8521238d04da243847211 Mon Sep 17 00:00:00 2001
From: Micha Reiser <micha@rome.tools>
Date: Thu, 15 Sep 2022 14:01:28 +0200
Subject: [PATCH] Comments infrastructure

---
 Cargo.lock                                    |   2 +
 crates/rome_formatter/Cargo.toml              |   4 +
 crates/rome_formatter/src/comments.rs         | 555 ++++++++--
 crates/rome_formatter/src/comments/builder.rs | 999 ++++++++++++++++++
 crates/rome_formatter/src/lib.rs              |  22 +-
 5 files changed, 1458 insertions(+), 124 deletions(-)
 create mode 100644 crates/rome_formatter/src/comments/builder.rs

diff --git a/Cargo.lock b/Cargo.lock
index d8fd1524a0c..1a4de63b6a9 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1571,6 +1571,8 @@ dependencies = [
  "cfg-if",
  "countme",
  "indexmap",
+ "rome_js_parser",
+ "rome_js_syntax",
  "rome_rowan",
  "rustc-hash",
  "schemars",
diff --git a/crates/rome_formatter/Cargo.toml b/crates/rome_formatter/Cargo.toml
index fde1f730112..5c63d0c7b5b 100644
--- a/crates/rome_formatter/Cargo.toml
+++ b/crates/rome_formatter/Cargo.toml
@@ -18,5 +18,9 @@ schemars = { version = "0.8.10", optional = true }
 rustc-hash = "1.1.0"
 countme = "3.0.1"
 
+[dev-dependencies]
+rome_js_parser = { path = "../rome_js_parser"}
+rome_js_syntax = { path = "../rome_js_syntax" }
+
 [features]
 serde = ["dep:serde", "schemars", "rome_rowan/serde"]
diff --git a/crates/rome_formatter/src/comments.rs b/crates/rome_formatter/src/comments.rs
index 4c53a59b443..f022ae8bbee 100644
--- a/crates/rome_formatter/src/comments.rs
+++ b/crates/rome_formatter/src/comments.rs
@@ -1,12 +1,13 @@
+mod builder;
 mod map;
 
-use rome_rowan::{
-    Direction, Language, SyntaxElement, SyntaxKind, SyntaxNode, SyntaxTriviaPieceComments,
-    WalkEvent,
-};
+use self::{builder::CommentsBuilderVisitor, map::CommentsMap};
+use crate::{TextSize, TransformSourceMap};
+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::collections::HashSet;
 use std::rc::Rc;
 
 #[derive(Copy, Clone, Eq, PartialEq, Debug)]
@@ -47,43 +48,6 @@ pub enum CommentKind {
     Line,
 }
 
-#[derive(Debug, Clone)]
-pub struct SourceComment<L: Language> {
-    /// The number of lines appearing before this comment
-    lines_before: u32,
-
-    /// The comment piece
-    piece: SyntaxTriviaPieceComments<L>,
-}
-
-impl<L: Language> SourceComment<L> {
-    /// Creates a new trailing comment. A trailing comment always has 0 lines before.
-    pub fn trailing(piece: SyntaxTriviaPieceComments<L>) -> Self {
-        Self {
-            lines_before: 0,
-            piece,
-        }
-    }
-
-    /// Creates a leading comment with the specified lines before
-    pub fn leading(piece: SyntaxTriviaPieceComments<L>, lines_before: u32) -> Self {
-        Self {
-            lines_before,
-            piece,
-        }
-    }
-
-    /// Returns the underlining comment trivia piece
-    pub fn piece(&self) -> &SyntaxTriviaPieceComments<L> {
-        &self.piece
-    }
-
-    /// Returns the number of lines before directly before this comment
-    pub fn lines_before(&self) -> u32 {
-        self.lines_before
-    }
-}
-
 impl CommentKind {
     pub const fn is_line(&self) -> bool {
         matches!(self, CommentKind::Line)
@@ -116,24 +80,236 @@ impl CommentKind {
     }
 }
 
+#[derive(Debug, Clone)]
+pub struct SourceComment<L: Language> {
+    /// The number of lines appearing before this comment
+    pub(crate) lines_before: u32,
+
+    pub(crate) lines_after: u32,
+
+    /// The comment piece
+    pub(crate) piece: SyntaxTriviaPieceComments<L>,
+
+    pub(crate) kind: CommentKind,
+}
+
+impl<L: Language> SourceComment<L> {
+    /// Returns the underlining comment trivia piece
+    pub fn piece(&self) -> &SyntaxTriviaPieceComments<L> {
+        &self.piece
+    }
+
+    /// Returns the number of lines before directly before this comment
+    pub fn lines_before(&self) -> u32 {
+        self.lines_before
+    }
+
+    pub fn lines_after(&self) -> u32 {
+        self.lines_after
+    }
+
+    /// The kind of the comment
+    pub fn kind(&self) -> CommentKind {
+        self.kind
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct DecoratedComment<L: Language> {
+    enclosing: SyntaxNode<L>,
+    preceding: Option<SyntaxNode<L>>,
+    following: Option<SyntaxNode<L>>,
+    following_token: SyntaxToken<L>,
+    position: CommentTextPosition,
+    lines_before: u32,
+    lines_after: u32,
+    comment: SyntaxTriviaPieceComments<L>,
+    kind: CommentKind,
+}
+
+impl<L: Language> DecoratedComment<L> {
+    /// The node that fully encloses the comment (the comment's start and end position are fully in the
+    /// node's bounds).
+    pub fn enclosing_node(&self) -> &SyntaxNode<L> {
+        &self.enclosing
+    }
+
+    pub fn piece(&self) -> &SyntaxTriviaPieceComments<L> {
+        &self.comment
+    }
+
+    /// The node directly preceding the comment or [None] if the comment is preceded by a token or is the first
+    /// token in the program.
+    pub fn preceding_node(&self) -> Option<&SyntaxNode<L>> {
+        self.preceding.as_ref()
+    }
+
+    fn take_preceding_node(&mut self) -> Option<SyntaxNode<L>> {
+        self.preceding.take()
+    }
+
+    /// The node directly following the comment or [None] if the comment is followed by a token or is the last token in the program.
+    pub fn following_node(&self) -> Option<&SyntaxNode<L>> {
+        self.following.as_ref()
+    }
+
+    fn take_following_node(&mut self) -> Option<SyntaxNode<L>> {
+        self.following.take()
+    }
+
+    /// The number of lines between this comment and the **previous** token, comment or skipped trivia.
+    pub fn lines_before(&self) -> u32 {
+        self.lines_before
+    }
+
+    pub fn lines_after(&self) -> u32 {
+        self.lines_after
+    }
+
+    /// Returns the [kind](CommentKind) of the comment.
+    pub fn kind(&self) -> CommentKind {
+        self.kind
+    }
+
+    pub fn position(&self) -> CommentTextPosition {
+        self.position
+    }
+
+    pub fn following_token(&self) -> &SyntaxToken<L> {
+        &self.following_token
+    }
+}
+
+impl<L: Language> From<DecoratedComment<L>> for SourceComment<L> {
+    fn from(decorated: DecoratedComment<L>) -> Self {
+        Self {
+            lines_before: decorated.lines_before,
+            lines_after: decorated.lines_after,
+            piece: decorated.comment,
+            kind: decorated.kind,
+        }
+    }
+}
+
+/// The position of a comment in the source document.
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
+pub enum CommentTextPosition {
+    /// A comment that is separated by at least one line break from the following token
+    ///
+    /// ```javascript
+    /// a; /* this */ // or this
+    /// b;
+    EndOfLine,
+
+    /// A Comment that is separated by at least one line break from the preceding token
+    ///
+    /// ```javascript
+    /// a;
+    /// /* comment */ /* or this */
+    /// b;
+    /// ```
+    OwnLine,
+
+    /// A comment that is placed on the same line as the preceding and following token.
+    ///
+    /// ```javascript
+    /// a /* comment */ + b
+    /// ```
+    SameLine,
+}
+
+impl CommentTextPosition {
+    pub const fn is_same_line(&self) -> bool {
+        matches!(self, CommentTextPosition::SameLine)
+    }
+
+    pub const fn is_own_line(&self) -> bool {
+        matches!(self, CommentTextPosition::OwnLine)
+    }
+
+    pub const fn is_end_of_line(&self) -> bool {
+        matches!(self, CommentTextPosition::EndOfLine)
+    }
+}
+
+#[derive(Debug)]
+pub enum CommentPlacement<L: Language> {
+    /// Overrides the positioning of the comment to be a leading node comment.
+    Leading {
+        node: SyntaxNode<L>,
+        comment: SourceComment<L>,
+    },
+    /// Overrides the positioning of the comment to be a trailing node comment.
+    Trailing {
+        node: SyntaxNode<L>,
+        comment: SourceComment<L>,
+    },
+
+    /// Makes this comment a dangling comment of `node`
+    Dangling {
+        node: SyntaxNode<L>,
+        comment: SourceComment<L>,
+    },
+
+    /// Uses the default positioning rules for the comment.
+    /// TODO document rules
+    Default(DecoratedComment<L>),
+}
+
+impl<L: Language> CommentPlacement<L> {
+    #[inline]
+    pub fn leading(node: SyntaxNode<L>, comment: impl Into<SourceComment<L>>) -> Self {
+        Self::Leading {
+            node,
+            comment: comment.into(),
+        }
+    }
+
+    pub fn dangling(node: SyntaxNode<L>, comment: impl Into<SourceComment<L>>) -> Self {
+        Self::Dangling {
+            node,
+            comment: comment.into(),
+        }
+    }
+
+    #[inline]
+    pub fn trailing(node: SyntaxNode<L>, comment: impl Into<SourceComment<L>>) -> Self {
+        Self::Trailing {
+            node,
+            comment: comment.into(),
+        }
+    }
+
+    #[inline]
+    pub fn or_else<F>(self, or_else: F) -> Self
+    where
+        F: FnOnce(DecoratedComment<L>) -> CommentPlacement<L>,
+    {
+        match self {
+            CommentPlacement::Default(comment) => or_else(comment),
+            placement => placement,
+        }
+    }
+}
+
 /// Defines how to format comments for a specific [Language].
-pub trait CommentStyle<L: Language> {
+pub trait CommentStyle: Default {
+    type Language: Language;
+
     /// Returns `true` if a comment with the given `text` is a `rome-ignore format:` suppression comment.
-    fn is_suppression(&self, text: &str) -> bool;
+    fn is_suppression(_text: &str) -> bool {
+        false
+    }
 
     /// Returns the (kind)[CommentKind] of the comment
-    fn get_comment_kind(&self, comment: &SyntaxTriviaPieceComments<L>) -> CommentKind;
-
-    /// Returns `true` if a token with the passed `kind` marks the start of a group. Common group tokens are:
-    /// * left parentheses: `(`, `[`, `{`
-    fn is_group_start_token(&self, kind: L::Kind) -> bool;
-
-    /// Returns `true` if a token with the passed `kind` marks the end of a group. Common group end tokens are:
-    /// * right parentheses: `)`, `]`, `}`
-    /// * end of statement token: `;`
-    /// * element separator: `,` or `.`.
-    /// * end of file token: `EOF`
-    fn is_group_end_token(&self, kind: L::Kind) -> bool;
+    fn get_comment_kind(comment: &SyntaxTriviaPieceComments<Self::Language>) -> CommentKind;
+
+    fn place_comment(
+        &self,
+        comment: DecoratedComment<Self::Language>,
+    ) -> CommentPlacement<Self::Language> {
+        CommentPlacement::Default(comment)
+    }
 }
 
 /// Type that stores the comments of a tree and gives access to:
@@ -143,7 +319,7 @@ pub trait CommentStyle<L: Language> {
 /// * the dangling comments of a token
 ///
 /// Cloning `comments` is cheap as it only involves bumping a reference counter.
-#[derive(Debug, Default, Clone)]
+#[derive(Debug, Clone, Default)]
 pub struct Comments<L: Language> {
     /// The use of a [Rc] is necessary to achieve that [Comments] has a lifetime that is independent from the [crate::Formatter].
     /// Having independent lifetimes is necessary to support the use case where a (formattable object)[crate::Format]
@@ -165,65 +341,117 @@ pub struct Comments<L: Language> {
 
 impl<L: Language> Comments<L> {
     /// Extracts all the suppressions from `root` and its child nodes.
-    pub fn from_node<FormatLanguage>(root: &SyntaxNode<L>, language: &FormatLanguage) -> Self
+    pub fn from_node<Style>(
+        root: &SyntaxNode<L>,
+        style: &Style,
+        source_map: Option<&TransformSourceMap>,
+    ) -> Self
     where
-        FormatLanguage: crate::FormatLanguage<SyntaxLanguage = L>,
+        Style: CommentStyle<Language = L>,
     {
-        let mut suppressed_nodes = HashSet::new();
-        let mut current_node = None;
-
-        for event in root.preorder_with_tokens(Direction::Next) {
-            match event {
-                WalkEvent::Enter(SyntaxElement::Node(node)) => {
-                    // Lists cannot have a suppression comment attached, it must
-                    // belong to either the entire parent node or one of the children
-                    if node.kind().is_root() || node.kind().is_list() {
-                        continue;
-                    }
-
-                    if current_node.is_none() {
-                        current_node = Some(node);
-                    }
-                }
-                WalkEvent::Leave(SyntaxElement::Node(node)) => {
-                    if current_node == Some(node) {
-                        current_node = None;
-                    }
-                }
-                WalkEvent::Enter(SyntaxElement::Token(token)) => {
-                    if let Some(current_node) = current_node.take() {
-                        for comment in token
-                            .leading_trivia()
-                            .pieces()
-                            .filter_map(|piece| piece.as_comments())
-                        {
-                            if language.comment_style().is_suppression(comment.text()) {
-                                suppressed_nodes.insert(current_node);
-                                break;
-                            }
-                        }
-                    }
-                }
-                WalkEvent::Leave(SyntaxElement::Token(_)) => {
-                    // Token already handled as part of the enter event.
-                }
-            }
-        }
+        let builder = CommentsBuilderVisitor::new(style, source_map);
 
-        let data = CommentsData {
-            suppressed_nodes,
-            #[cfg(debug_assertions)]
-            checked_suppressions: RefCell::default(),
-        };
+        let (comments, skipped) = builder.visit(root);
 
         Self {
-            data: Rc::new(data),
+            data: Rc::new(CommentsData {
+                root: Some(root.clone()),
+                is_suppression: Style::is_suppression,
+
+                comments,
+                with_skipped: skipped,
+                #[cfg(debug_assertions)]
+                checked_suppressions: RefCell::new(Default::default()),
+            }),
         }
     }
 
+    /// Returns `true` if the given `node` has any leading or trailing comments.
+    #[inline]
+    pub fn has_comments(&self, node: &SyntaxNode<L>) -> bool {
+        self.data.comments.has(&node.key())
+    }
+
+    /// Returns `true` if the given [node] has any leading comments.
+    /// By default, a comment is a node's leading comment if:
+    /// * the previous sibling is a token
+    /// * there's a line break before the commend ending before this comment and the comment.
+    #[inline]
+    pub fn has_leading_comments(&self, node: &SyntaxNode<L>) -> bool {
+        !self.leading_comments(node).is_empty()
+    }
+
+    /// Tests if the node has any leading comment that will be placed on its own line.
+    pub fn has_leading_own_line_comment(&self, node: &SyntaxNode<L>) -> bool {
+        self.leading_comments(node)
+            .iter()
+            .any(|comment| comment.lines_after() > 0)
+    }
+
+    /// Returns the [node]'s leading comments.
+    #[inline]
+    pub fn leading_comments(&self, node: &SyntaxNode<L>) -> &[SourceComment<L>] {
+        self.data.comments.leading(&node.key())
+    }
+
+    /// Returns `true` if node has any dangling comments.
+    pub fn has_dangling_comments(&self, node: &SyntaxNode<L>) -> bool {
+        !self.dangling_comments(node).is_empty()
+    }
+
+    /// Returns the dangling comments of `node`
+    pub fn dangling_comments(&self, node: &SyntaxNode<L>) -> &[SourceComment<L>] {
+        self.data.comments.dangling(&node.key())
+    }
+
+    /// Returns the [node]'s trailing comments.
+    #[inline]
+    pub fn trailing_comments(&self, node: &SyntaxNode<L>) -> &[SourceComment<L>] {
+        self.data.comments.trailing(&node.key())
+    }
+
+    pub fn has_trailing_line_comment(&self, node: &SyntaxNode<L>) -> bool {
+        self.trailing_comments(node)
+            .iter()
+            .any(|comment| comment.kind().is_line())
+    }
+
+    /// Returns `true` if the given [node] has any trailing comments.
+    /// By default, a comment is a node's trailing comment if:
+    /// * the next sibling is a token
+    /// * there's **no** line break between the node and this comment.
+    #[inline]
+    pub fn has_trailing_comments(&self, node: &SyntaxNode<L>) -> bool {
+        !self.trailing_comments(node).is_empty()
+    }
+
+    /// Returns an iterator over the leading and trailing comments of `node`.
+    pub fn leading_trailing_comments(
+        &self,
+        node: &SyntaxNode<L>,
+    ) -> impl Iterator<Item = &SourceComment<L>> {
+        self.leading_comments(node)
+            .iter()
+            .chain(self.trailing_comments(node).iter())
+    }
+
+    /// Returns an iterator over the leading, dangling, and trailing comments of `node`.
+    pub fn leading_dangling_trailing_comments<'a>(
+        &'a self,
+        node: &'a SyntaxNode<L>,
+    ) -> impl Iterator<Item = &SourceComment<L>> + 'a {
+        self.data.comments.parts(&node.key())
+    }
+
+    /// Returns `true` if that node has skipped token trivia attached.
+    #[inline]
+    pub fn has_skipped(&self, token: &SyntaxToken<L>) -> bool {
+        self.data.with_skipped.contains(&token.key())
+    }
+
     /// Returns `true` if the passed `node` has a leading suppression comment.
     ///
-    /// Suppression comments only apply if they are at the start of a node and they suppress the most
+    /// Suppression comments only apply if they at the start of a node and they suppress the most
     /// outer node.
     ///
     /// # Examples
@@ -237,7 +465,10 @@ impl<L: Language> Comments<L> {
     /// call expression is nested inside of the expression statement.
     pub fn is_suppressed(&self, node: &SyntaxNode<L>) -> bool {
         self.mark_suppression_checked(node);
-        self.data.suppressed_nodes.contains(node)
+        let is_suppression = self.data.is_suppression;
+
+        self.leading_dangling_trailing_comments(node)
+            .any(|comment| is_suppression(comment.piece().text()))
     }
 
     /// Marks that it isn't necessary for the given node to check if it has been suppressed or not.
@@ -262,6 +493,8 @@ impl<L: Language> Comments<L> {
     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() {
@@ -287,10 +520,14 @@ Node:
     }
 }
 
-#[derive(Debug, Default)]
 struct CommentsData<L: Language> {
-    /// Stores the nodes that have at least one leading suppression comment.
-    suppressed_nodes: HashSet<SyntaxNode<L>>,
+    root: Option<SyntaxNode<L>>,
+
+    is_suppression: fn(&str) -> bool,
+
+    /// Stores all leading node comments by node
+    comments: CommentsMap<SyntaxElementKey, SourceComment<L>>,
+    with_skipped: FxHashSet<SyntaxElementKey>,
 
     /// Stores all nodes for which [Comments::is_suppressed] has been called.
     /// This index of nodes that have been checked if they have a suppression comments is used to
@@ -301,5 +538,101 @@ struct CommentsData<L: Language> {
     /// as verbatim if its formatting fails which has the same result as formatting it as suppressed node
     /// (thus, guarantees that the formatting isn't changed).
     #[cfg(debug_assertions)]
-    checked_suppressions: RefCell<HashSet<SyntaxNode<L>>>,
+    checked_suppressions: RefCell<FxHashSet<SyntaxNode<L>>>,
+}
+
+impl<L: Language> Default for CommentsData<L> {
+    fn default() -> Self {
+        Self {
+            root: None,
+            is_suppression: |_| false,
+            comments: Default::default(),
+            with_skipped: Default::default(),
+            #[cfg(debug_assertions)]
+            checked_suppressions: Default::default(),
+        }
+    }
+}
+
+impl<L: Language> std::fmt::Debug for CommentsData<L> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let mut comments = Vec::new();
+
+        if let Some(root) = &self.root {
+            for node in root.descendants() {
+                for leading in self.comments.leading(&node.key()) {
+                    comments.push(DebugComment::Leading {
+                        node: node.clone(),
+                        comment: leading,
+                    });
+                }
+
+                for dangling in self.comments.dangling(&node.key()) {
+                    comments.push(DebugComment::Dangling {
+                        node: node.clone(),
+                        comment: dangling,
+                    });
+                }
+
+                for trailing in self.comments.trailing(&node.key()) {
+                    comments.push(DebugComment::Trailing {
+                        node: node.clone(),
+                        comment: trailing,
+                    });
+                }
+            }
+        }
+
+        comments.sort_by_key(|comment| comment.start());
+
+        f.debug_list().entries(comments).finish()
+    }
+}
+
+/// Helper for printing a comment of [Comments]
+enum DebugComment<'a, L: Language> {
+    Leading {
+        comment: &'a SourceComment<L>,
+        node: SyntaxNode<L>,
+    },
+    Trailing {
+        comment: &'a SourceComment<L>,
+        node: SyntaxNode<L>,
+    },
+    Dangling {
+        comment: &'a SourceComment<L>,
+        node: SyntaxNode<L>,
+    },
+}
+
+impl<L: Language> DebugComment<'_, L> {
+    fn start(&self) -> TextSize {
+        match self {
+            DebugComment::Leading { comment, .. }
+            | DebugComment::Trailing { comment, .. }
+            | DebugComment::Dangling { comment, .. } => comment.piece.text_range().start(),
+        }
+    }
+}
+
+impl<L: Language> std::fmt::Debug for DebugComment<'_, L> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            DebugComment::Leading { node, comment } => f
+                .debug_struct("Leading")
+                .field("node", node)
+                .field("comment", comment)
+                .finish(),
+            DebugComment::Dangling { node, comment } => f
+                .debug_struct("Dangling")
+                .field("node", node)
+                .field("comment", comment)
+                .finish(),
+            DebugComment::Trailing { node, comment } => f
+                .debug_struct("Trailing")
+                .field("node", node)
+                .field("comment", comment)
+                .finish(),
+        }
+    }
 }
diff --git a/crates/rome_formatter/src/comments/builder.rs b/crates/rome_formatter/src/comments/builder.rs
new file mode 100644
index 00000000000..c7484535e07
--- /dev/null
+++ b/crates/rome_formatter/src/comments/builder.rs
@@ -0,0 +1,999 @@
+use crate::comments::map::CommentsMap;
+use crate::comments::CommentTextPosition;
+use crate::source_map::{DeletedRangeEntry, DeletedRanges};
+use crate::{
+    CommentPlacement, CommentStyle, DecoratedComment, SourceComment, TextRange, TextSize,
+    TransformSourceMap,
+};
+use rome_rowan::syntax::SyntaxElementKey;
+use rome_rowan::{
+    Direction, Language, SyntaxElement, SyntaxKind, SyntaxNode, SyntaxToken, WalkEvent,
+};
+use rustc_hash::FxHashSet;
+
+pub(super) struct CommentsBuilderVisitor<'a, Style: CommentStyle> {
+    builder: CommentsBuilder<Style::Language>,
+    style: &'a Style,
+    parentheses: SourceParentheses<'a>,
+
+    // State
+    pending_comments: Vec<DecoratedComment<Style::Language>>,
+    preceding_node: Option<SyntaxNode<Style::Language>>,
+    following_node_index: Option<usize>,
+    parents: Vec<SyntaxNode<Style::Language>>,
+    last_token: Option<SyntaxToken<Style::Language>>,
+}
+
+impl<'a, Style> CommentsBuilderVisitor<'a, Style>
+where
+    Style: CommentStyle,
+{
+    pub(super) fn new(style: &'a Style, source_map: Option<&'a TransformSourceMap>) -> Self {
+        Self {
+            style,
+            builder: Default::default(),
+            parentheses: SourceParentheses::from_source_map(source_map),
+
+            pending_comments: Default::default(),
+            preceding_node: Default::default(),
+            following_node_index: Default::default(),
+            parents: Default::default(),
+            last_token: Default::default(),
+        }
+    }
+
+    pub(super) fn visit(
+        mut self,
+        root: &SyntaxNode<Style::Language>,
+    ) -> (
+        CommentsMap<SyntaxElementKey, SourceComment<Style::Language>>,
+        FxHashSet<SyntaxElementKey>,
+    ) {
+        for event in root.preorder_with_tokens(Direction::Next) {
+            match event {
+                WalkEvent::Enter(SyntaxElement::Node(node)) => {
+                    self.visit_node(WalkEvent::Enter(node))
+                }
+
+                WalkEvent::Leave(SyntaxElement::Node(node)) => {
+                    self.visit_node(WalkEvent::Leave(node))
+                }
+
+                WalkEvent::Enter(SyntaxElement::Token(token)) => self.visit_token(token),
+                WalkEvent::Leave(SyntaxElement::Token(_)) => {
+                    // Handled as part of enter
+                }
+            }
+        }
+
+        assert!(
+            self.parents.is_empty(),
+            "Expected all enclosing nodes to have been processed but contains {:#?}",
+            self.parents
+        );
+        self.flush_comments(None);
+
+        self.builder.finish()
+    }
+
+    fn visit_node(&mut self, event: WalkEvent<SyntaxNode<Style::Language>>) {
+        match event {
+            WalkEvent::Enter(node) => {
+                // Lists cannot have comments attached. They either belong to the entire parent or to
+                // the first child. So we ignore lists all together
+                if node.kind().is_list() {
+                    return;
+                }
+
+                let is_root = matches!(self.following_node_index, Some(0));
+
+                // Associate comments with the most outer node
+                // Set following here because it is the "following node" of the next token's leading trivia.
+                if self.following_node_index.is_none() || is_root {
+                    // Flush in case the node doesn't have any tokens.
+                    self.flush_comments(Some(&node));
+                    self.following_node_index = Some(self.parents.len());
+                }
+
+                self.parents.push(node);
+            }
+
+            WalkEvent::Leave(node) => {
+                if node.kind().is_list() {
+                    return;
+                }
+
+                self.parents.pop().unwrap();
+
+                // We're passed this node, flush any pending comments for its children
+                self.following_node_index = None;
+                self.flush_comments(None);
+
+                // We're passed this node, so it must precede the sibling that comes next.
+                self.preceding_node = Some(node);
+            }
+        }
+    }
+
+    fn visit_token(&mut self, token: SyntaxToken<Style::Language>) {
+        let mut comments_start = self.pending_comments.len();
+
+        // The index of the last trailing comment in `pending_comments`.
+        let mut trailing_end: Option<usize> = None;
+
+        // Number of lines before the next comment, token, or skipped token trivia
+        let mut lines_before = 0;
+
+        // Trailing comments are all `SameLine` comments EXCEPT if any is followed by a line break,
+        // a leading comment (that always have line breaks), or there's a line break before the token.
+        let mut position = CommentTextPosition::SameLine;
+
+        // Process the trailing trivia of the last token
+        if let Some(last_token) = self.last_token.take() {
+            for piece in last_token.trailing_trivia().pieces() {
+                if piece.is_newline() {
+                    lines_before += 1;
+                    // All comments following from here are own line comments
+                    position = CommentTextPosition::OwnLine;
+
+                    if trailing_end.is_none() {
+                        trailing_end = Some(self.pending_comments.len());
+                    }
+                } else if let Some(comment) = piece.as_comments() {
+                    self.queue_comment(DecoratedComment {
+                        enclosing: self.enclosing_node().clone(),
+                        preceding: self.preceding_node.clone(),
+                        following: None,
+                        following_token: token.clone(),
+                        lines_before,
+                        lines_after: 0, // Will be initialized after
+                        position,
+                        kind: Style::get_comment_kind(&comment),
+                        comment,
+                    });
+
+                    lines_before = 0;
+                }
+
+                if let Some(parens_source_range) = self
+                    .parentheses
+                    .r_paren_source_range(piece.text_range().end())
+                {
+                    self.flush_before_r_paren_comments(
+                        parens_source_range,
+                        &last_token,
+                        position,
+                        lines_before,
+                        comments_start,
+                        trailing_end,
+                    );
+
+                    lines_before = 0;
+                    position = CommentTextPosition::SameLine;
+                    comments_start = 0;
+                    trailing_end = None;
+                }
+            }
+        }
+
+        // Process the leading trivia of the current token. the trailing trivia is handled as part of the next token
+        for leading in token.leading_trivia().pieces() {
+            if leading.is_newline() {
+                lines_before += 1;
+                // All comments following from here are own line comments
+                position = CommentTextPosition::OwnLine;
+
+                if trailing_end.is_none() {
+                    trailing_end = Some(self.pending_comments.len());
+                }
+            } else if leading.is_skipped() {
+                self.builder.mark_has_skipped(&token);
+
+                lines_before = 0;
+                break;
+            } else if let Some(comment) = leading.as_comments() {
+                let kind = Style::get_comment_kind(&comment);
+
+                self.queue_comment(DecoratedComment {
+                    enclosing: self.enclosing_node().clone(),
+                    preceding: self.preceding_node.clone(),
+                    following: None,
+                    following_token: token.clone(),
+                    lines_before,
+                    lines_after: 0,
+                    position,
+                    kind,
+                    comment,
+                });
+
+                lines_before = 0;
+            }
+        }
+
+        let trailing_end = trailing_end.unwrap_or(self.pending_comments.len());
+
+        self.last_token = Some(token);
+
+        let mut comments = self.pending_comments[comments_start..]
+            .iter_mut()
+            .enumerate()
+            .peekable();
+
+        // Update the lines after of all comments as well as the positioning of end of line comments.
+        while let Some((index, comment)) = comments.next() {
+            // Update the position of all trailing comments to be end of line as we've seen a line break since.
+            if index < trailing_end && position.is_own_line() {
+                comment.position = CommentTextPosition::EndOfLine;
+            }
+
+            comment.lines_after = comments
+                .peek()
+                .map_or(lines_before, |(_, next)| next.lines_before);
+        }
+
+        // Set following node to `None` because it now becomes the enclosing node.
+        if let Some(following_node) = self.following_node() {
+            self.flush_comments(Some(&following_node.clone()));
+            self.following_node_index = None;
+
+            // The following node is only set after entering a node
+            // That means, following node is only set for the first token of a node.
+            // Unset preceding node if this is the first token because the preceding node belongs to the parent.
+            self.preceding_node = None;
+        }
+    }
+
+    fn enclosing_node(&self) -> &SyntaxNode<Style::Language> {
+        let element = match self.following_node_index {
+            None => self.parents.last(),
+            Some(index) if index == 0 => Some(&self.parents[0]),
+            Some(index) => Some(&self.parents[index - 1]),
+        };
+
+        element.expect("Expected enclosing nodes to at least contain the root node.")
+    }
+
+    fn following_node(&self) -> Option<&SyntaxNode<Style::Language>> {
+        self.following_node_index.map(|index| {
+            self.parents
+                .get(index)
+                .expect("Expected following node index to point to a valid parent node")
+        })
+    }
+
+    fn queue_comment(&mut self, comment: DecoratedComment<Style::Language>) {
+        self.pending_comments.push(comment);
+    }
+
+    fn flush_comments(&mut self, following: Option<&SyntaxNode<Style::Language>>) {
+        for mut comment in self.pending_comments.drain(..) {
+            comment.following = following.cloned();
+
+            let placement = self.style.place_comment(comment);
+            self.builder.add_comment(placement);
+        }
+    }
+
+    /// Processes comments appearing right before a `)` of a parenthesized expressions.
+    #[cold]
+    fn flush_before_r_paren_comments(
+        &mut self,
+        parens_source_range: TextRange,
+        last_token: &SyntaxToken<Style::Language>,
+        position: CommentTextPosition,
+        lines_before: u32,
+        start: usize,
+        trailing_end: Option<usize>,
+    ) {
+        let enclosing = self.enclosing_node().clone();
+
+        let trailing_end = trailing_end.unwrap_or(self.pending_comments.len());
+        let mut comments = self.pending_comments[start..]
+            .iter_mut()
+            .enumerate()
+            .peekable();
+
+        let parenthesized_node = self
+            .parentheses
+            .outer_most_parenthesized_node(last_token, parens_source_range);
+
+        // SAFETY: Safe, because the above loop at least returns the parent of the token. If this isn't the case,
+        // then it's likely that the source map is corrupted.
+        let preceding = parenthesized_node.expect("Last token to have a parent node.");
+
+        // Using the `enclosing` as default but it's mainly to satisfy Rust. The only case where it is used
+        // is if someone formats a Parenthesized expression as the root. Something we explicitly disallow
+        // in rome_js_formatter
+        let enclosing = preceding.parent().unwrap_or(enclosing);
+
+        // Update the lines after of all comments as well as the positioning of end of line comments.
+        while let Some((index, comment)) = comments.next() {
+            // Update the position of all trailing comments to be end of line as we've seen a line break since.
+            if index < trailing_end && position.is_own_line() {
+                comment.position = CommentTextPosition::EndOfLine;
+            }
+
+            comment.preceding = Some(preceding.clone());
+            comment.enclosing = enclosing.clone();
+            comment.lines_after = comments
+                .peek()
+                .map_or(lines_before, |(_, next)| next.lines_before);
+        }
+
+        self.flush_comments(None);
+    }
+}
+
+struct CommentsBuilder<L: Language> {
+    comments: CommentsMap<SyntaxElementKey, SourceComment<L>>,
+    skipped: FxHashSet<SyntaxElementKey>,
+}
+
+impl<L: Language> CommentsBuilder<L> {
+    fn add_comment(&mut self, placement: CommentPlacement<L>) {
+        match placement {
+            CommentPlacement::Leading { node, comment } => {
+                self.push_leading_comment(&node, comment);
+            }
+            CommentPlacement::Trailing { node, comment } => {
+                self.push_trailing_comment(&node, comment);
+            }
+            CommentPlacement::Dangling { node, comment } => {
+                self.push_dangling_comment(&node, comment)
+            }
+            CommentPlacement::Default(mut comment) => {
+                match comment.position {
+                    CommentTextPosition::EndOfLine => {
+                        match (comment.take_preceding_node(), comment.take_following_node()) {
+                            (Some(preceding), Some(_)) => {
+                                // Attach comments with both preceding and following node to the preceding
+                                // because there's a line break separating it from the following node.
+                                // ```javascript
+                                // a; // comment
+                                // b
+                                // ```
+                                self.push_trailing_comment(&preceding, comment);
+                            }
+                            (Some(preceding), None) => {
+                                self.push_trailing_comment(&preceding, comment);
+                            }
+                            (None, Some(following)) => {
+                                self.push_leading_comment(&following, comment);
+                            }
+                            (None, None) => {
+                                self.push_dangling_comment(
+                                    &comment.enclosing_node().clone(),
+                                    comment,
+                                );
+                            }
+                        }
+                    }
+                    CommentTextPosition::OwnLine => {
+                        match (comment.take_preceding_node(), comment.take_following_node()) {
+                            // Following always wins for a leading comment
+                            // ```javascript
+                            // a;
+                            // // comment
+                            // b
+                            // ```
+                            // attach the comment to the `b` expression statement
+                            (_, Some(following)) => {
+                                self.push_leading_comment(&following, comment);
+                            }
+                            (Some(preceding), None) => {
+                                self.push_trailing_comment(&preceding, comment);
+                            }
+                            (None, None) => {
+                                self.push_dangling_comment(
+                                    &comment.enclosing_node().clone(),
+                                    comment,
+                                );
+                            }
+                        }
+                    }
+                    CommentTextPosition::SameLine => {
+                        match (comment.take_preceding_node(), comment.take_following_node()) {
+                            (Some(preceding), Some(following)) => {
+                                // Only make it a trailing comment if it directly follows the preceding node but not if it is separated
+                                // by one or more tokens
+                                // ```javascript
+                                // a /* comment */ b;   //  Comment is a trailing comment
+                                // a, /* comment */ b;  // Comment should be a leading comment
+                                // ```
+                                if preceding.text_range().end()
+                                    == comment.piece().as_piece().token().text_range().end()
+                                {
+                                    self.push_trailing_comment(&preceding, comment);
+                                } else {
+                                    self.push_leading_comment(&following, comment);
+                                }
+                            }
+                            (Some(preceding), None) => {
+                                self.push_trailing_comment(&preceding, comment);
+                            }
+                            (None, Some(following)) => {
+                                self.push_leading_comment(&following, comment);
+                            }
+                            (None, None) => {
+                                self.push_dangling_comment(
+                                    &comment.enclosing_node().clone(),
+                                    comment,
+                                );
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    fn mark_has_skipped(&mut self, token: &SyntaxToken<L>) {
+        self.skipped.insert(token.key());
+    }
+
+    fn push_leading_comment(&mut self, node: &SyntaxNode<L>, comment: impl Into<SourceComment<L>>) {
+        self.comments.push_leading(node.key(), comment.into());
+    }
+
+    fn push_dangling_comment(
+        &mut self,
+        node: &SyntaxNode<L>,
+        comment: impl Into<SourceComment<L>>,
+    ) {
+        self.comments.push_dangling(node.key(), comment.into());
+    }
+
+    fn push_trailing_comment(
+        &mut self,
+        node: &SyntaxNode<L>,
+        comment: impl Into<SourceComment<L>>,
+    ) {
+        self.comments.push_trailing(node.key(), comment.into());
+    }
+
+    fn finish(
+        self,
+    ) -> (
+        CommentsMap<SyntaxElementKey, SourceComment<L>>,
+        FxHashSet<SyntaxElementKey>,
+    ) {
+        (self.comments, self.skipped)
+    }
+}
+
+impl<L: Language> Default for CommentsBuilder<L> {
+    fn default() -> Self {
+        Self {
+            comments: CommentsMap::new(),
+            skipped: FxHashSet::default(),
+        }
+    }
+}
+
+enum SourceParentheses<'a> {
+    Empty,
+    SourceMap {
+        map: &'a TransformSourceMap,
+        next: Option<DeletedRangeEntry<'a>>,
+        tail: DeletedRanges<'a>,
+    },
+}
+
+impl<'a> SourceParentheses<'a> {
+    fn from_source_map(source_map: Option<&'a TransformSourceMap>) -> Self {
+        match source_map {
+            None => Self::Empty,
+            Some(source_map) => {
+                let mut deleted = source_map.deleted_ranges();
+                SourceParentheses::SourceMap {
+                    map: source_map,
+                    next: deleted.next(),
+                    tail: deleted,
+                }
+            }
+        }
+    }
+
+    /// Returns the range of `node` including its parentheses if any. Otherwise returns the range as is
+    fn parenthesized_range<L: Language>(&self, node: &SyntaxNode<L>) -> TextRange {
+        match self {
+            SourceParentheses::Empty => node.text_trimmed_range(),
+            SourceParentheses::SourceMap { map, .. } => map.trimmed_source_range(node),
+        }
+    }
+
+    /// Tests if the next offset is at a position where the original source document used to have an `)`.
+    ///
+    /// Must be called with offsets in increasing order.
+    ///
+    /// Returns the source range of the `)` if there's any `)` in the deleted range at this offset. Returns `None` otherwise
+
+    fn r_paren_source_range(&mut self, offset: TextSize) -> Option<TextRange> {
+        match self {
+            SourceParentheses::Empty => None,
+            SourceParentheses::SourceMap { next, tail, .. } => {
+                while let Some(range) = next {
+                    #[allow(clippy::comparison_chain)]
+                    if range.transformed == offset {
+                        // A deleted range can contain multiple tokens. See if there's any `)` in the deleted
+                        // range and compute its source range.
+                        return range.text.find(')').map(|r_paren_position| {
+                            let start = range.source + TextSize::from(r_paren_position as u32);
+                            TextRange::at(start, TextSize::from(1))
+                        });
+                    } else if range.transformed > offset {
+                        return None;
+                    } else {
+                        *next = tail.next();
+                    }
+                }
+
+                None
+            }
+        }
+    }
+
+    /// Searches the outer most node that still is inside of the parentheses specified by the `parentheses_source_range`.
+    fn outer_most_parenthesized_node<L: Language>(
+        &self,
+        token: &SyntaxToken<L>,
+        parentheses_source_range: TextRange,
+    ) -> Option<SyntaxNode<L>> {
+        match self {
+            SourceParentheses::Empty => token.parent(),
+            SourceParentheses::SourceMap { map, .. } => {
+                debug_assert_eq!(&map.text()[parentheses_source_range], ")");
+
+                // How this works: We search the outer most node that, in the source document ends right after the `)`.
+                // The issue is, it is possible that multiple nodes end right after the `)`
+                //
+                // ```javascript
+                // !(
+                //     a
+                //     /* comment */
+                //  )
+                // ```
+                // The issue is, that in the transformed document, the `ReferenceIdentifier`, `IdentifierExpression`, `UnaryExpression`, and `ExpressionStatement`
+                // all end at the end position of `)`.
+                // However, not all the nodes start at the same position. That's why this code also tracks the start.
+                // We first find the closest node that directly ends at the position of the right paren. We then continue
+                // upwards to find the most outer node that starts at the same position as that node. (In this case,
+                // `ReferenceIdentifier` -> `IdentifierExpression`.
+                let mut start_offset = None;
+                let r_paren_source_end = parentheses_source_range.end();
+
+                let ancestors = token.ancestors().take_while(|node| {
+                    let source_range = self.parenthesized_range(node);
+
+                    if let Some(start) = start_offset {
+                        TextRange::new(start, r_paren_source_end).contains_range(source_range)
+                    } else if source_range.end() == r_paren_source_end {
+                        start_offset = Some(source_range.start());
+                        true
+                    } else {
+                        source_range.end() < r_paren_source_end
+                    }
+                });
+
+                ancestors.last()
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::comments::builder::CommentsBuilderVisitor;
+    use crate::comments::map::CommentsMap;
+    use crate::comments::CommentTextPosition;
+    use crate::{
+        CommentKind, CommentPlacement, CommentStyle, DecoratedComment, SourceComment, TextSize,
+        TransformSourceMap, TransformSourceMapBuilder,
+    };
+    use rome_js_parser::parse_module;
+    use rome_js_syntax::{
+        JsIdentifierExpression, JsLanguage, JsParameters, JsParenthesizedExpression,
+        JsPropertyObjectMember, JsReferenceIdentifier, JsShorthandPropertyObjectMember,
+        JsSyntaxKind, JsSyntaxNode, JsUnaryExpression,
+    };
+    use rome_rowan::syntax::SyntaxElementKey;
+    use rome_rowan::{AstNode, BatchMutation, SyntaxNode, SyntaxTriviaPieceComments, TextRange};
+    use std::cell::RefCell;
+
+    #[test]
+    fn leading_comment() {
+        let (root, decorated, comments) = extract_comments(
+            r#"const foo = {
+  a: 'a',
+  /* comment for this line */
+  b
+};"#,
+        );
+
+        assert_eq!(decorated.len(), 1);
+
+        let comment = decorated.last().unwrap();
+
+        assert_eq!(comment.position(), CommentTextPosition::OwnLine);
+        assert_eq!(comment.lines_before(), 1);
+        assert_eq!(comment.lines_after(), 1);
+        assert_eq!(
+            comment
+                .preceding_node()
+                .map(|node| node.text_trimmed().to_string())
+                .as_deref(),
+            Some("a: 'a'")
+        );
+        assert_eq!(
+            comment
+                .following_node()
+                .map(|node| node.text_trimmed().to_string())
+                .as_deref(),
+            Some("b")
+        );
+        assert_eq!(
+            comment.enclosing_node().kind(),
+            JsSyntaxKind::JS_OBJECT_EXPRESSION
+        );
+
+        let b = root
+            .descendants()
+            .find_map(JsShorthandPropertyObjectMember::cast)
+            .unwrap();
+
+        assert!(!comments.leading(&b.syntax().key()).is_empty());
+    }
+
+    #[test]
+    fn trailing_comment() {
+        let (root, decorated, comments) = extract_comments(
+            r#"const foo = {
+  a: 'a' /* comment for this line */,
+  b
+};"#,
+        );
+
+        assert_eq!(decorated.len(), 1);
+
+        let comment = decorated.last().unwrap();
+
+        assert_eq!(comment.position(), CommentTextPosition::SameLine);
+        assert_eq!(comment.lines_after(), 0);
+        assert_eq!(comment.lines_before(), 0);
+        assert_eq!(
+            comment
+                .preceding_node()
+                .map(|node| node.text_trimmed().to_string())
+                .as_deref(),
+            Some("a: 'a'")
+        );
+        assert_eq!(
+            comment
+                .following_node()
+                .map(|node| node.text_trimmed().to_string())
+                .as_deref(),
+            Some("b")
+        );
+        assert_eq!(
+            comment.enclosing_node().kind(),
+            JsSyntaxKind::JS_OBJECT_EXPRESSION
+        );
+
+        let a = root
+            .descendants()
+            .find_map(JsPropertyObjectMember::cast)
+            .unwrap();
+
+        assert!(!comments.trailing(&a.syntax().key()).is_empty());
+    }
+
+    #[test]
+    fn end_of_line_comment() {
+        let (root, decorated, comments) = extract_comments(
+            r#"const foo = {
+  a: 'a', /* comment for this line */
+  b
+};"#,
+        );
+
+        assert_eq!(decorated.len(), 1);
+
+        let comment = decorated.last().unwrap();
+
+        assert_eq!(comment.position(), CommentTextPosition::EndOfLine);
+        assert_eq!(comment.lines_before(), 0);
+        assert_eq!(comment.lines_after(), 1);
+        assert_eq!(
+            comment
+                .preceding_node()
+                .map(|node| node.text_trimmed().to_string())
+                .as_deref(),
+            Some("a: 'a'")
+        );
+        assert_eq!(
+            comment
+                .following_node()
+                .map(|node| node.text_trimmed().to_string())
+                .as_deref(),
+            Some("b")
+        );
+        assert_eq!(
+            comment.enclosing_node().kind(),
+            JsSyntaxKind::JS_OBJECT_EXPRESSION
+        );
+
+        let a = root
+            .descendants()
+            .find_map(JsPropertyObjectMember::cast)
+            .unwrap();
+
+        assert!(!comments.trailing(&a.syntax().key()).is_empty());
+    }
+
+    #[test]
+    fn dangling_arrow() {
+        let (root, decorated_comments, comments) = extract_comments("(/* comment */)  => true");
+
+        assert_eq!(decorated_comments.len(), 1);
+
+        let decorated = &decorated_comments[0];
+        assert_eq!(decorated.position(), CommentTextPosition::SameLine);
+        assert_eq!(decorated.lines_before(), 0);
+        assert_eq!(decorated.lines_after(), 0);
+        assert_eq!(decorated.preceding_node(), None);
+        assert_eq!(decorated.following_node(), None);
+        assert_eq!(
+            decorated.enclosing_node().kind(),
+            JsSyntaxKind::JS_PARAMETERS
+        );
+
+        let parameters = root.descendants().find_map(JsParameters::cast).unwrap();
+        assert!(!comments.dangling(&parameters.syntax().key()).is_empty());
+    }
+
+    #[test]
+    fn dangling_comments() {
+        let (root, decorated_comments, comments) = extract_comments(
+            r#"
+            function (/* test */) {}
+            "#,
+        );
+
+        assert_eq!(decorated_comments.len(), 1);
+
+        let decorated = &decorated_comments[0];
+        assert_eq!(decorated.position(), CommentTextPosition::SameLine);
+        assert_eq!(decorated.lines_before(), 0);
+        assert_eq!(decorated.lines_after(), 0);
+        assert_eq!(decorated.preceding_node(), None);
+        assert_eq!(decorated.following_node(), None);
+        assert_eq!(
+            decorated.enclosing_node().kind(),
+            JsSyntaxKind::JS_PARAMETERS
+        );
+
+        let parameters = root.descendants().find_map(JsParameters::cast).unwrap();
+        assert!(!comments.dangling(&parameters.syntax().key()).is_empty());
+    }
+
+    #[test]
+    fn r_paren() {
+        let source = r#"!(
+    a
+    /* comment */
+)
+/* comment */
+b;"#;
+
+        let mut source_map_builder = TransformSourceMapBuilder::with_source(source.to_string());
+        let l_paren_range = TextRange::new(TextSize::from(1), TextSize::from(2));
+        let r_paren_range = TextRange::new(TextSize::from(27), TextSize::from(28));
+
+        assert_eq!(&source[l_paren_range], "(");
+        assert_eq!(&source[r_paren_range], ")");
+
+        source_map_builder.add_deleted_range(l_paren_range);
+        source_map_builder.add_deleted_range(r_paren_range);
+        source_map_builder.extend_trimmed_node_range(
+            TextRange::new(TextSize::from(7), TextSize::from(8)),
+            TextRange::new(l_paren_range.start(), r_paren_range.end()),
+        );
+
+        let source_map = source_map_builder.finish();
+
+        let root = parse_module(source, 0).syntax();
+
+        // A lot of code that simply removes the parenthesized expression and moves the parens
+        // trivia to the identifiers leading / trailing trivia.
+        let parenthesized = root
+            .descendants()
+            .find_map(JsParenthesizedExpression::cast)
+            .unwrap();
+
+        let reference_identifier = root
+            .descendants()
+            .find_map(JsReferenceIdentifier::cast)
+            .unwrap();
+
+        let mut mutation = BatchMutation::new(root);
+
+        let identifier_expression =
+            JsIdentifierExpression::cast(parenthesized.expression().unwrap().into_syntax())
+                .unwrap();
+        let l_paren = parenthesized.l_paren_token().unwrap();
+        let r_paren = parenthesized.r_paren_token().unwrap();
+
+        let identifier_token = reference_identifier.value_token().unwrap();
+        let new_identifier_token = identifier_token
+            .with_leading_trivia_pieces(
+                l_paren
+                    .leading_trivia()
+                    .pieces()
+                    .chain(l_paren.trailing_trivia().pieces())
+                    .chain(identifier_token.leading_trivia().pieces())
+                    .collect::<Vec<_>>(),
+            )
+            .with_trailing_trivia_pieces(
+                identifier_token
+                    .trailing_trivia()
+                    .pieces()
+                    .chain(r_paren.leading_trivia().pieces())
+                    .chain(r_paren.trailing_trivia().pieces())
+                    .collect::<Vec<_>>(),
+            );
+
+        let new_reference_identifier = reference_identifier.with_value_token(new_identifier_token);
+
+        let new_identifier_expression = identifier_expression.with_name(new_reference_identifier);
+
+        mutation.replace_element_discard_trivia(
+            parenthesized.into_syntax().into(),
+            new_identifier_expression.into_syntax().into(),
+        );
+
+        let transformed = mutation.commit();
+
+        let style = TestCommentStyle::default();
+        let comments_builder = CommentsBuilderVisitor::new(&style, Some(&source_map));
+        let (comments, _) = comments_builder.visit(&transformed);
+
+        let decorated_comments = style.finish();
+
+        assert_eq!(decorated_comments.len(), 2);
+
+        let argument_trailing = &decorated_comments[0];
+        assert_eq!(argument_trailing.position(), CommentTextPosition::OwnLine);
+        assert_eq!(argument_trailing.lines_before(), 1);
+        assert_eq!(argument_trailing.lines_after(), 1);
+        assert_eq!(
+            argument_trailing
+                .preceding_node()
+                .map(|preceding| preceding.kind()),
+            Some(JsSyntaxKind::JS_IDENTIFIER_EXPRESSION)
+        );
+        assert_eq!(argument_trailing.following_node(), None);
+        assert_eq!(
+            argument_trailing.enclosing_node().kind(),
+            JsSyntaxKind::JS_UNARY_EXPRESSION
+        );
+
+        let identifier_leading = &decorated_comments[1];
+        assert_eq!(identifier_leading.position(), CommentTextPosition::OwnLine);
+        assert_eq!(identifier_leading.lines_before(), 1);
+        assert_eq!(identifier_leading.lines_after(), 1);
+        assert_eq!(
+            identifier_leading.preceding_node().map(SyntaxNode::kind),
+            Some(JsSyntaxKind::JS_EXPRESSION_STATEMENT)
+        );
+        assert_eq!(
+            identifier_leading.following_node().map(SyntaxNode::kind),
+            Some(JsSyntaxKind::JS_EXPRESSION_STATEMENT)
+        );
+        assert_eq!(
+            identifier_leading.enclosing_node().kind(),
+            JsSyntaxKind::JS_MODULE
+        );
+
+        let unary = transformed
+            .descendants()
+            .find_map(JsUnaryExpression::cast)
+            .unwrap();
+        assert!(!comments
+            .trailing(&unary.argument().unwrap().syntax().key())
+            .is_empty());
+    }
+
+    #[test]
+    fn comment_only_program() {
+        let (root, decorated, comments) = extract_comments(
+            r#"/* test */
+
+/* test */"#,
+        );
+
+        assert_eq!(decorated.len(), 2);
+
+        let first = &decorated[0];
+        assert_eq!(first.position(), CommentTextPosition::EndOfLine);
+        assert_eq!(first.lines_before(), 0);
+        assert_eq!(first.lines_after(), 2);
+        assert_eq!(first.preceding_node(), None);
+        assert_eq!(
+            first.following_node().map(SyntaxNode::kind),
+            Some(JsSyntaxKind::JS_MODULE)
+        );
+        assert_eq!(first.enclosing_node().kind(), JsSyntaxKind::JS_MODULE);
+
+        let second = &decorated[1];
+        assert_eq!(second.position(), CommentTextPosition::OwnLine);
+        assert_eq!(second.lines_before(), 2);
+        assert_eq!(second.lines_after(), 0);
+        assert_eq!(second.preceding_node(), None);
+        assert_eq!(
+            first.following_node().map(SyntaxNode::kind),
+            Some(JsSyntaxKind::JS_MODULE)
+        );
+        assert_eq!(second.enclosing_node().kind(), JsSyntaxKind::JS_MODULE);
+
+        assert!(!comments.leading(&root.key()).is_empty());
+    }
+
+    fn extract_comments(
+        source: &str,
+    ) -> (
+        JsSyntaxNode,
+        Vec<DecoratedComment<JsLanguage>>,
+        CommentsMap<SyntaxElementKey, SourceComment<JsLanguage>>,
+    ) {
+        extract_with_source_map(source, None)
+    }
+
+    fn extract_with_source_map(
+        source: &str,
+        source_map: Option<&TransformSourceMap>,
+    ) -> (
+        JsSyntaxNode,
+        Vec<DecoratedComment<JsLanguage>>,
+        CommentsMap<SyntaxElementKey, SourceComment<JsLanguage>>,
+    ) {
+        let tree = parse_module(source, 0);
+
+        let style = TestCommentStyle::default();
+        let builder = CommentsBuilderVisitor::new(&style, source_map);
+        let (comments, _) = builder.visit(&tree.syntax());
+
+        (tree.syntax(), style.finish(), comments)
+    }
+
+    #[derive(Default)]
+    struct TestCommentStyle {
+        recorded_comments: RefCell<Vec<DecoratedComment<JsLanguage>>>,
+    }
+
+    impl CommentStyle for TestCommentStyle {
+        type Language = JsLanguage;
+
+        fn is_suppression(_: &str) -> bool {
+            false
+        }
+
+        fn get_comment_kind(_: &SyntaxTriviaPieceComments<Self::Language>) -> CommentKind {
+            CommentKind::Block
+        }
+
+        fn place_comment(
+            &self,
+            comment: DecoratedComment<Self::Language>,
+        ) -> CommentPlacement<Self::Language> {
+            self.recorded_comments.borrow_mut().push(comment.clone());
+
+            CommentPlacement::Default(comment)
+        }
+    }
+
+    impl TestCommentStyle {
+        fn finish(self) -> Vec<DecoratedComment<JsLanguage>> {
+            self.recorded_comments.take()
+        }
+    }
+}
diff --git a/crates/rome_formatter/src/lib.rs b/crates/rome_formatter/src/lib.rs
index 3e3cc9a8510..d5e12a4353b 100644
--- a/crates/rome_formatter/src/lib.rs
+++ b/crates/rome_formatter/src/lib.rs
@@ -53,7 +53,10 @@ pub use builders::{
     if_group_fits_on_line, indent, labelled, line_suffix, soft_block_indent, soft_line_break,
     soft_line_break_or_space, soft_line_indent_or_space, space, text, BestFitting,
 };
-pub use comments::{CommentKind, CommentStyle, Comments, SourceComment};
+pub use comments::{
+    CommentKind, CommentPlacement, CommentStyle, CommentTextPosition, Comments, DecoratedComment,
+    SourceComment,
+};
 pub use format_element::{normalize_newlines, FormatElement, Text, Verbatim, LINE_TERMINATORS};
 pub use group_id::GroupId;
 use indexmap::IndexSet;
@@ -232,14 +235,10 @@ pub trait FormatOptions {
 /// The context customizes the comments formatting and stores the comments of the CST.
 pub trait CstFormatContext: FormatContext {
     type Language: Language;
-    type Style: CommentStyle<Self::Language>;
-
-    /// Rule for formatting leading comments.
-    type LeadingCommentRule: FormatRule<SourceComment<Self::Language>, Context = Self> + Default;
+    type Style: CommentStyle<Language = Self::Language>;
 
-    /// Customizes how comments are formatted
-    #[deprecated(note = "Prefer FormatLanguage::comment_style")]
-    fn comment_style(&self) -> Self::Style;
+    /// Rule for formatting comments.
+    type CommentRule: FormatRule<SourceComment<Self::Language>, Context = Self> + Default;
 
     /// Returns a reference to the program's comments.
     fn comments(&self) -> &Comments<Self::Language>;
@@ -825,14 +824,11 @@ pub trait FormatLanguage {
     type Context: CstFormatContext<Language = Self::SyntaxLanguage>;
 
     /// The type specifying how to format comments.
-    type CommentStyle: CommentStyle<Self::SyntaxLanguage>;
+    type CommentStyle: CommentStyle<Language = Self::SyntaxLanguage>;
 
     /// The rule type that can format a [SyntaxNode] of this language
     type FormatRule: FormatRule<SyntaxNode<Self::SyntaxLanguage>, Context = Self::Context> + Default;
 
-    /// Customizes how comments are formatted
-    fn comment_style(&self) -> Self::CommentStyle;
-
     /// Performs an optional pre-processing of the tree. This can be useful to remove nodes
     /// that otherwise complicate formatting.
     ///
@@ -877,7 +873,7 @@ pub fn format_node<L: FormatLanguage>(
             None => (root.clone(), None),
         };
 
-        let comments = Comments::from_node(&root, &language);
+        let comments = Comments::from_node(&root, &L::CommentStyle::default(), source_map.as_ref());
         let format_node = FormatRefWithRule::new(&root, L::FormatRule::default());
 
         let context = language.create_context(comments, source_map);