Skip to main content

swc_ecma_codegen/text_writer/
basic_impl.rs

1use std::io::Write;
2
3use memchr::memchr2;
4use rustc_hash::FxBuildHasher;
5use swc_allocator::api::global::HashSet;
6use swc_common::{sync::Lrc, BytePos, LineCol, SourceMap, Span};
7
8use super::{BindingStorage, Result, ScopeBindingRecord, ScopeRecord, WriteJs};
9
10///
11/// -----
12///
13/// Ported from `createTextWriter` of the typescript compiler.
14///
15/// https://github.com/Microsoft/TypeScript/blob/45eaf42006/src/compiler/utilities.ts#L2548
16pub struct JsWriter<'a, W: Write> {
17    indent: usize,
18    indent_str: &'static str,
19    line_start: bool,
20    line_count: usize,
21    line_pos: usize,
22    new_line: &'a str,
23    srcmap: Option<&'a mut Vec<(BytePos, LineCol)>>,
24    srcmap_done: HashSet<(BytePos, u32, u32), FxBuildHasher>,
25    /// Used to avoid including whitespaces created by indention.
26    pending_srcmap: Option<BytePos>,
27    wr: W,
28    scopes: Option<&'a mut Vec<ScopeRecord>>,
29    scope_stack: Vec<u32>,
30}
31
32impl<'a, W: Write> JsWriter<'a, W> {
33    pub fn new(
34        cm: Lrc<SourceMap>,
35        new_line: &'a str,
36        wr: W,
37        srcmap: Option<&'a mut Vec<(BytePos, LineCol)>>,
38    ) -> Self {
39        Self::new_with_scopes(cm, new_line, wr, srcmap, None)
40    }
41
42    pub fn new_with_scopes(
43        _: Lrc<SourceMap>,
44        new_line: &'a str,
45        wr: W,
46        srcmap: Option<&'a mut Vec<(BytePos, LineCol)>>,
47        scopes: Option<&'a mut Vec<ScopeRecord>>,
48    ) -> Self {
49        JsWriter {
50            indent: Default::default(),
51            indent_str: "    ",
52            line_start: true,
53            line_count: 0,
54            line_pos: Default::default(),
55            new_line,
56            srcmap,
57            wr,
58            pending_srcmap: Default::default(),
59            srcmap_done: Default::default(),
60            scopes,
61            scope_stack: Default::default(),
62        }
63    }
64
65    pub fn preamble(&mut self, s: &str) -> Result {
66        self.raw_write(s)?;
67        self.update_pos(s);
68
69        Ok(())
70    }
71
72    /// Sets the indentation string. Defaults to four spaces.
73    pub fn set_indent_str(&mut self, indent_str: &'static str) {
74        self.indent_str = indent_str;
75    }
76
77    #[inline]
78    fn write_indent_string(&mut self) -> Result {
79        for _ in 0..self.indent {
80            self.raw_write(self.indent_str)?;
81        }
82        if self.srcmap.is_some() {
83            self.line_pos += self.indent_str.len() * self.indent;
84        }
85
86        Ok(())
87    }
88
89    #[inline]
90    fn raw_write(&mut self, data: &str) -> Result {
91        self.wr.write_all(data.as_bytes())?;
92
93        Ok(())
94    }
95
96    #[inline(always)]
97    fn write(&mut self, span: Option<Span>, data: &str) -> Result {
98        if !data.is_empty() {
99            if self.line_start {
100                self.write_indent_string()?;
101                self.line_start = false;
102
103                if let Some(pending) = self.pending_srcmap.take() {
104                    self.srcmap(pending);
105                }
106            }
107
108            if let Some(span) = span {
109                self.srcmap(span.lo());
110            }
111
112            self.raw_write(data)?;
113            self.update_pos(data);
114
115            if let Some(span) = span {
116                self.srcmap(span.hi());
117            }
118        }
119
120        Ok(())
121    }
122
123    #[inline]
124    fn update_pos(&mut self, s: &str) {
125        if self.srcmap.is_none() || s.is_empty() {
126            return;
127        }
128
129        let bytes = s.as_bytes();
130
131        // Most tokens are single-byte ASCII characters with no line terminator.
132        if bytes.len() == 1 {
133            match bytes[0] {
134                b'\n' | b'\r' => {
135                    self.line_count += 1;
136                    self.line_pos = 0;
137                }
138                _ => {
139                    self.line_pos += 1;
140                }
141            }
142
143            return;
144        }
145
146        // Fast path for common ASCII tokens without line breaks.
147        if memchr2(b'\n', b'\r', bytes).is_none() {
148            if s.is_ascii() {
149                self.line_pos += bytes.len();
150            } else {
151                self.line_pos += s.encode_utf16().count();
152            }
153            return;
154        }
155
156        let line_start_of_s = compute_line_starts(s);
157        self.line_count += line_start_of_s.line_count;
158
159        if s.is_ascii() {
160            self.line_pos = bytes.len() - line_start_of_s.byte_pos;
161        } else {
162            self.line_pos = s[line_start_of_s.byte_pos..].encode_utf16().count();
163        }
164    }
165
166    #[inline]
167    fn srcmap(&mut self, byte_pos: BytePos) {
168        if byte_pos.is_dummy() && byte_pos != BytePos(u32::MAX) {
169            return;
170        }
171
172        if let Some(ref mut srcmap) = self.srcmap {
173            let key = (byte_pos, self.line_count as _, self.line_pos as _);
174            if self.srcmap_done.insert(key) {
175                let loc = LineCol {
176                    line: key.1,
177                    col: key.2,
178                };
179
180                srcmap.push((byte_pos, loc));
181            }
182        }
183    }
184
185    #[inline]
186    fn generated_pos(&self) -> LineCol {
187        LineCol {
188            line: self.line_count as u32,
189            col: self.line_pos as u32,
190        }
191    }
192}
193
194impl<W: Write> WriteJs for JsWriter<'_, W> {
195    #[inline]
196    fn increase_indent(&mut self) -> Result {
197        self.indent += 1;
198        Ok(())
199    }
200
201    #[inline]
202    fn decrease_indent(&mut self) -> Result {
203        self.indent -= 1;
204        Ok(())
205    }
206
207    #[inline]
208    fn write_semi(&mut self, span: Option<Span>) -> Result {
209        self.write(span, ";")?;
210        Ok(())
211    }
212
213    #[inline]
214    fn write_space(&mut self) -> Result {
215        self.write(None, " ")?;
216        Ok(())
217    }
218
219    #[inline]
220    fn write_keyword(&mut self, span: Option<Span>, s: &'static str) -> Result {
221        self.write(span, s)?;
222        Ok(())
223    }
224
225    #[inline]
226    fn write_operator(&mut self, span: Option<Span>, s: &str) -> Result {
227        self.write(span, s)?;
228        Ok(())
229    }
230
231    #[inline]
232    fn write_param(&mut self, s: &str) -> Result {
233        self.write(None, s)?;
234        Ok(())
235    }
236
237    #[inline]
238    fn write_property(&mut self, s: &str) -> Result {
239        self.write(None, s)?;
240        Ok(())
241    }
242
243    #[inline]
244    fn write_line(&mut self) -> Result {
245        let pending = self.pending_srcmap.take();
246        if !self.line_start {
247            self.raw_write(self.new_line)?;
248            if self.srcmap.is_some() {
249                self.line_count += 1;
250                self.line_pos = 0;
251            }
252            self.line_start = true;
253
254            if let Some(pending) = pending {
255                self.srcmap(pending)
256            }
257        }
258
259        Ok(())
260    }
261
262    #[inline]
263    fn write_lit(&mut self, span: Span, s: &str) -> Result {
264        if !s.is_empty() {
265            self.srcmap(span.lo());
266            self.write(None, s)?;
267            self.srcmap(span.hi());
268        }
269
270        Ok(())
271    }
272
273    #[inline]
274    fn write_comment(&mut self, s: &str) -> Result {
275        self.write(None, s)?;
276        Ok(())
277    }
278
279    #[inline]
280    fn write_str_lit(&mut self, span: Span, s: &str) -> Result {
281        if !s.is_empty() {
282            self.srcmap(span.lo());
283            self.write(None, s)?;
284            self.srcmap(span.hi());
285        }
286
287        Ok(())
288    }
289
290    #[inline]
291    fn write_str(&mut self, s: &str) -> Result {
292        self.write(None, s)?;
293        Ok(())
294    }
295
296    #[inline]
297    fn write_symbol(&mut self, span: Span, s: &str) -> Result {
298        self.write(Some(span), s)?;
299        Ok(())
300    }
301
302    #[inline]
303    fn write_punct(
304        &mut self,
305        span: Option<Span>,
306        s: &'static str,
307        _commit_pending_semi: bool,
308    ) -> Result {
309        self.write(span, s)?;
310        Ok(())
311    }
312
313    #[inline]
314    fn care_about_srcmap(&self) -> bool {
315        self.srcmap.is_some()
316    }
317
318    #[inline]
319    fn add_srcmap(&mut self, pos: BytePos) -> Result {
320        if self.srcmap.is_some() {
321            if self.line_start {
322                self.pending_srcmap = Some(pos);
323            } else {
324                self.srcmap(pos);
325            }
326        }
327        Ok(())
328    }
329
330    #[inline]
331    fn commit_pending_semi(&mut self) -> Result {
332        Ok(())
333    }
334
335    #[inline(always)]
336    fn can_ignore_invalid_unicodes(&mut self) -> bool {
337        false
338    }
339
340    #[inline(always)]
341    fn has_scope_tracking(&self) -> bool {
342        self.scopes.is_some()
343    }
344
345    #[inline]
346    fn start_scope(
347        &mut self,
348        name: Option<&str>,
349        kind: super::ScopeKind,
350        is_stack_frame: bool,
351        is_hidden: bool,
352        original_span: Option<Span>,
353    ) -> Result {
354        let start = self.generated_pos();
355        let Some(scopes) = self.scopes.as_deref_mut() else {
356            return Ok(());
357        };
358
359        let idx = scopes.len() as u32;
360        let parent = self.scope_stack.last().copied();
361        scopes.push(ScopeRecord {
362            parent,
363            generated_start: start,
364            generated_end: None,
365            original_span,
366            kind,
367            is_stack_frame,
368            is_hidden,
369            name: name.map(str::to_owned),
370            bindings: vec![],
371        });
372        self.scope_stack.push(idx);
373
374        Ok(())
375    }
376
377    #[inline]
378    fn end_scope(&mut self) -> Result {
379        let Some(idx) = self.scope_stack.pop() else {
380            return Ok(());
381        };
382
383        let end = self.generated_pos();
384        if let Some(scopes) = self.scopes.as_deref_mut() {
385            if let Some(scope) = scopes.get_mut(idx as usize) {
386                scope.generated_end = Some(end);
387            }
388        }
389
390        Ok(())
391    }
392
393    #[inline]
394    fn add_scope_variable(
395        &mut self,
396        name: &str,
397        expression: Option<&str>,
398        storage: BindingStorage,
399    ) -> Result {
400        let Some(scopes) = self.scopes.as_deref() else {
401            return Ok(());
402        };
403
404        let target_scope = match storage {
405            BindingStorage::Hoisted => self
406                .scope_stack
407                .iter()
408                .rev()
409                .copied()
410                .find(|idx| scopes[*idx as usize].kind.is_hoist_target())
411                .or_else(|| self.scope_stack.last().copied()),
412            BindingStorage::Lexical => self.scope_stack.last().copied(),
413        };
414
415        let Some(target_scope) = target_scope else {
416            return Ok(());
417        };
418
419        let Some(scopes) = self.scopes.as_deref_mut() else {
420            return Ok(());
421        };
422        let scope = &mut scopes[target_scope as usize];
423        if scope.bindings.iter().any(|binding| binding.name == name) {
424            return Ok(());
425        }
426
427        scope.bindings.push(ScopeBindingRecord {
428            name: name.to_string(),
429            expression: expression.map(str::to_owned),
430            storage,
431        });
432
433        Ok(())
434    }
435}
436
437#[derive(Debug)]
438struct LineStart {
439    line_count: usize,
440    byte_pos: usize,
441}
442fn compute_line_starts(s: &str) -> LineStart {
443    compute_line_starts_from_bytes(s.as_bytes())
444}
445
446#[inline]
447fn compute_line_starts_from_bytes(bytes: &[u8]) -> LineStart {
448    let mut count = 0;
449    let mut line_start = 0;
450    let mut pos = 0;
451
452    while pos < bytes.len() {
453        match bytes[pos] {
454            b'\r' => {
455                count += 1;
456                pos += 1;
457                if pos < bytes.len() && bytes[pos] == b'\n' {
458                    pos += 1;
459                } else {
460                    // already moved over '\r'
461                }
462                line_start = pos;
463            }
464            b'\n' => {
465                count += 1;
466                pos += 1;
467                line_start = pos;
468            }
469            _ => {
470                pos += 1;
471            }
472        }
473    }
474
475    LineStart {
476        line_count: count,
477        byte_pos: line_start,
478    }
479}
480
481#[cfg(test)]
482mod test {
483    use std::sync::Arc;
484
485    use swc_common::SourceMap;
486
487    use super::{compute_line_starts_from_bytes, JsWriter};
488    use crate::text_writer::{BindingStorage, ScopeKind, WriteJs};
489
490    #[test]
491    fn changes_indent_str() {
492        let source_map = Arc::new(SourceMap::default());
493        let mut output = Vec::new();
494        let mut writer = JsWriter::new(source_map, "\n", &mut output, None);
495        writer.set_indent_str("\t");
496        writer.increase_indent().unwrap();
497        writer.write_indent_string().unwrap();
498        writer.increase_indent().unwrap();
499        writer.write_indent_string().unwrap();
500        assert_eq!(output, "\t\t\t".as_bytes());
501    }
502
503    #[test]
504    fn collects_scopes_and_bindings() {
505        let source_map = Arc::new(SourceMap::default());
506        let mut output = Vec::new();
507        let mut scopes = vec![];
508        let mut writer =
509            JsWriter::new_with_scopes(source_map, "\n", &mut output, None, Some(&mut scopes));
510
511        writer
512            .start_scope(Some("global"), ScopeKind::Global, false, false, None)
513            .unwrap();
514        writer
515            .add_scope_variable("g", Some("g"), BindingStorage::Hoisted)
516            .unwrap();
517
518        writer
519            .start_scope(Some("f"), ScopeKind::Function, true, false, None)
520            .unwrap();
521        writer
522            .add_scope_variable("p", Some("p"), BindingStorage::Lexical)
523            .unwrap();
524        writer
525            .start_scope(None, ScopeKind::Block, false, false, None)
526            .unwrap();
527        writer
528            .add_scope_variable("v", Some("v"), BindingStorage::Hoisted)
529            .unwrap();
530
531        writer.end_scope().unwrap();
532        writer.end_scope().unwrap();
533        writer.end_scope().unwrap();
534
535        assert_eq!(scopes.len(), 3);
536        assert_eq!(scopes[0].name.as_deref(), Some("global"));
537        assert_eq!(scopes[1].name.as_deref(), Some("f"));
538        assert_eq!(scopes[1].bindings.len(), 2);
539        assert_eq!(scopes[1].bindings[0].name, "p");
540        assert_eq!(scopes[1].bindings[1].name, "v");
541        assert_eq!(scopes[2].bindings.len(), 0);
542    }
543
544    #[test]
545    fn hoisted_binding_targets_nearest_hoist_scope() {
546        let source_map = Arc::new(SourceMap::default());
547        let mut output = Vec::new();
548        let mut scopes = vec![];
549        let mut writer =
550            JsWriter::new_with_scopes(source_map, "\n", &mut output, None, Some(&mut scopes));
551
552        writer
553            .start_scope(Some("global"), ScopeKind::Global, false, false, None)
554            .unwrap();
555        writer
556            .start_scope(Some("fn"), ScopeKind::Function, true, false, None)
557            .unwrap();
558        writer
559            .start_scope(None, ScopeKind::Block, false, false, None)
560            .unwrap();
561
562        writer
563            .add_scope_variable("h", Some("h"), BindingStorage::Hoisted)
564            .unwrap();
565        writer
566            .add_scope_variable("l", Some("l"), BindingStorage::Lexical)
567            .unwrap();
568
569        writer.end_scope().unwrap();
570        writer.end_scope().unwrap();
571        writer.end_scope().unwrap();
572
573        assert_eq!(scopes.len(), 3);
574        assert_eq!(scopes[1].bindings.len(), 1);
575        assert_eq!(scopes[1].bindings[0].name, "h");
576        assert_eq!(scopes[2].bindings.len(), 1);
577        assert_eq!(scopes[2].bindings[0].name, "l");
578    }
579
580    #[test]
581    fn deduplicates_scope_variables_by_name() {
582        let source_map = Arc::new(SourceMap::default());
583        let mut output = Vec::new();
584        let mut scopes = vec![];
585        let mut writer =
586            JsWriter::new_with_scopes(source_map, "\n", &mut output, None, Some(&mut scopes));
587
588        writer
589            .start_scope(Some("global"), ScopeKind::Global, false, false, None)
590            .unwrap();
591        writer
592            .add_scope_variable("dup", Some("first"), BindingStorage::Lexical)
593            .unwrap();
594        writer
595            .add_scope_variable("dup", Some("second"), BindingStorage::Lexical)
596            .unwrap();
597        writer.end_scope().unwrap();
598
599        assert_eq!(scopes.len(), 1);
600        assert_eq!(scopes[0].bindings.len(), 1);
601        assert_eq!(scopes[0].bindings[0].name, "dup");
602        assert_eq!(scopes[0].bindings[0].expression.as_deref(), Some("first"));
603    }
604
605    #[test]
606    fn tracks_position_for_ascii_without_newline() {
607        let source_map = Arc::new(SourceMap::default());
608        let mut output = Vec::new();
609        let mut srcmap = vec![];
610        let mut writer = JsWriter::new(source_map, "\n", &mut output, Some(&mut srcmap));
611
612        writer.write_str("abc").unwrap();
613
614        assert_eq!(writer.line_count, 0);
615        assert_eq!(writer.line_pos, 3);
616    }
617
618    #[test]
619    fn tracks_position_for_crlf_and_unicode_tail() {
620        let source_map = Arc::new(SourceMap::default());
621        let mut output = Vec::new();
622        let mut srcmap = vec![];
623        let mut writer = JsWriter::new(source_map, "\n", &mut output, Some(&mut srcmap));
624
625        writer.write_str("a\r\nβ").unwrap();
626
627        assert_eq!(writer.line_count, 1);
628        assert_eq!(writer.line_pos, 1);
629    }
630
631    #[test]
632    fn compute_line_starts_handles_crlf_pairs() {
633        let line_start = compute_line_starts_from_bytes(b"a\r\nb\r\n");
634
635        assert_eq!(line_start.line_count, 2);
636        assert_eq!(line_start.byte_pos, 6);
637    }
638}