swc/
dropped_comments_preserver.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
use 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};

/// Preserves comments that would otherwise be dropped.
///
/// If during compilation an ast node associated with
/// a comment is dropped, the comment will not appear in the final emitted
/// output. This can create problems in the JavaScript ecosystem, particularly
/// around istanbul coverage and other tooling that relies on comment
/// directives.
///
/// This transformer shifts orphaned comments to the next closest known span
/// while making a best-effort to preserve the "general orientation" of
/// comments.

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);
        }
    }

    /// We'll be shifting all comments to known span positions, so drain the
    /// current comments first to limit the amount of look ups needed into
    /// the hashmaps.
    ///
    /// This way, we only need to take the comments once, and then add them back
    /// once.
    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
    }

    /// Shift all comments to known leading positions.
    /// This prevents trailing comments from ending up associated with
    /// nodes that will not emit trailing comments, while
    /// preserving any comments that might show up after all code positions.
    ///
    /// This maintains the highest fidelity between existing comment positions
    /// of pre and post compiled code.
    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
    }

    /// These comments trail all known span lo byte positions.
    /// Therefore, by shifting them to trail the highest known hi position, we
    /// ensure that any remaining trailing comments are emitted in a
    /// similar location
    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(),
        );
    }
}