swc/
dropped_comments_preserver.rsuse swc_common::{
comments::{Comment, Comments, SingleThreadedComments},
BytePos, Span, DUMMY_SP,
};
use swc_ecma_ast::{Module, Script};
use swc_ecma_visit::{as_folder, noop_visit_mut_type, Fold, VisitMut, VisitMutWith};
pub fn dropped_comments_preserver(
comments: Option<SingleThreadedComments>,
) -> impl Fold + VisitMut {
as_folder(DroppedCommentsPreserver {
comments,
is_first_span: true,
known_spans: Vec::new(),
})
}
struct DroppedCommentsPreserver {
comments: Option<SingleThreadedComments>,
is_first_span: bool,
known_spans: Vec<Span>,
}
type CommentEntries = Vec<(BytePos, Vec<Comment>)>;
impl VisitMut for DroppedCommentsPreserver {
noop_visit_mut_type!();
fn visit_mut_module(&mut self, module: &mut Module) {
module.visit_mut_children_with(self);
self.known_spans
.sort_by(|span_a, span_b| span_a.lo.cmp(&span_b.lo));
self.shift_comments_to_known_spans();
}
fn visit_mut_script(&mut self, script: &mut Script) {
script.visit_mut_children_with(self);
self.known_spans
.sort_by(|span_a, span_b| span_a.lo.cmp(&span_b.lo));
self.shift_comments_to_known_spans();
}
fn visit_mut_span(&mut self, span: &mut Span) {
if span.is_dummy() || self.is_first_span {
self.is_first_span = false;
return;
}
self.known_spans.push(*span);
span.visit_mut_children_with(self)
}
}
impl DroppedCommentsPreserver {
fn shift_comments_to_known_spans(&self) {
if let Some(comments) = &self.comments {
let trailing_comments = self.shift_leading_comments(comments);
self.shift_trailing_comments(trailing_comments);
}
}
fn collect_existing_comments(&self, comments: &SingleThreadedComments) -> CommentEntries {
let (mut leading_comments, mut trailing_comments) = comments.borrow_all_mut();
let mut existing_comments: CommentEntries = leading_comments
.drain()
.chain(trailing_comments.drain())
.collect();
existing_comments.sort_by(|(bp_a, _), (bp_b, _)| bp_a.cmp(bp_b));
existing_comments
}
fn shift_leading_comments(&self, comments: &SingleThreadedComments) -> CommentEntries {
let mut existing_comments = self.collect_existing_comments(comments);
for span in self.known_spans.iter() {
let (comments_to_move, next_byte_positions): (CommentEntries, CommentEntries) =
existing_comments
.drain(..)
.partition(|(bp, _)| *bp <= span.lo);
existing_comments.extend(next_byte_positions);
let collected_comments = comments_to_move.into_iter().flat_map(|(_, c)| c).collect();
self.comments
.add_leading_comments(span.lo, collected_comments)
}
existing_comments
}
fn shift_trailing_comments(&self, remaining_comment_entries: CommentEntries) {
let last_trailing = self
.known_spans
.iter()
.copied()
.fold(
DUMMY_SP,
|acc, span| if span.hi > acc.hi { span } else { acc },
);
self.comments.add_trailing_comments(
last_trailing.hi,
remaining_comment_entries
.into_iter()
.flat_map(|(_, c)| c)
.collect(),
);
}
}