[codex] Optimize source map streaming hot paths

src/cached_source.rs

@@ -1,5 +1,6 @@
 use std::{
   borrow::Cow,
+  cell::OnceCell,
   hash::{Hash, Hasher},
   sync::{Arc, OnceLock},
 };
@@ -166,7 +167,8 @@ impl Source for CachedSource {
 }
 
 struct CachedSourceChunks<'source> {
-  chunks: Box<dyn Chunks + 'source>,
+  inner: &'source BoxSource,
+  chunks: OnceCell<Box<dyn Chunks + 'source>>,
   cache: Arc<CachedData>,
   source: Cow<'source, str>,
 }
@@ -176,7 +178,8 @@ impl<'a> CachedSourceChunks<'a> {
     let source = cache_source.source().into_string_lossy();
 
     Self {
-      chunks: cache_source.inner.stream_chunks(),
+      inner: &cache_source.inner,
+      chunks: OnceCell::new(),
       cache: cache_source.cache.clone(),
       source,
     }
@@ -220,10 +223,11 @@ impl Chunks for CachedSourceChunks<'_> {
         }
       }
       None => {
+        let chunks = self.chunks.get_or_init(|| self.inner.stream_chunks());
         let (generated_info, map) = stream_and_get_source_and_map(
           options,
           object_pool,
-          self.chunks.as_ref(),
+          chunks.as_ref(),
           on_chunk,
           on_source,
           on_name,

src/concat_source.rs

@@ -154,6 +154,9 @@ impl ConcatSource {
         children.extend(other_children.iter().cloned());
         return;
       }
+
+      children.push(box_source.clone());
+      return;
     }
 
     // Check if the source itself is a ConcatSource

src/encoder.rs

@@ -5,12 +5,22 @@ const B64_CHARS: &[u8] =
 
 #[inline(always)]
 pub fn encode_vlq(out: &mut Vec<u8>, a: u32, b: u32) {
+  if a == b {
+    out.push(b'A');
+    return;
+  }
+
   let mut num = if a >= b {
     (a - b) << 1
   } else {
     ((b - a) << 1) + 1
   };
 
+  if num < 0b100000 {
+    out.push(B64_CHARS[num as usize]);
+    return;
+  }
+
   loop {
     let mut digit = num & 0b11111;
     num >>= 5;
@@ -68,6 +78,14 @@ pub(crate) struct FullMappingsEncoder {
   mappings: Vec<u8>,
 }
 
+#[derive(Clone, Copy)]
+pub(crate) struct OriginalLocationInfo {
+  pub source_index: u32,
+  pub original_line: u32,
+  pub original_column: u32,
+  pub name_index: Option<u32>,
+}
+
 impl FullMappingsEncoder {
   pub fn new() -> Self {
     Self {
@@ -86,10 +104,73 @@ impl FullMappingsEncoder {
 }
 
 impl FullMappingsEncoder {
-  fn encode(&mut self, mapping: &Mapping) {
-    if self.active_mapping && self.current_line == mapping.generated_line {
+  #[inline(always)]
+  pub(crate) fn encode_original_no_name(
+    &mut self,
+    generated_line: u32,
+    generated_column: u32,
+    original_line: u32,
+    original_column: u32,
+  ) {
+    if self.active_mapping
+      && self.current_line == generated_line
+      && self.current_source_index == 0
+      && self.current_original_line == original_line
+      && self.current_original_column == original_column
+      && !self.active_name
+    {
+      return;
+    }
+
+    if self.current_line < generated_line {
+      let count = (generated_line - self.current_line) as usize;
+      self.mappings.extend(std::iter::repeat_n(b';', count));
+      self.current_line = generated_line;
+      self.current_column = 0;
+      self.initial = false;
+    } else if self.initial {
+      self.initial = false;
+    } else {
+      self.mappings.push(b',');
+    }
+
+    encode_vlq(&mut self.mappings, generated_column, self.current_column);
+    self.current_column = generated_column;
+    self.active_mapping = true;
+    self.active_name = false;
+    if self.current_source_index == 0 {
+      self.mappings.push(b'A');
+    } else {
+      encode_vlq(&mut self.mappings, 0, self.current_source_index);
+      self.current_source_index = 0;
+    }
+    encode_vlq(
+      &mut self.mappings,
+      original_line,
+      self.current_original_line,
+    );
+    self.current_original_line = original_line;
+    if original_column == self.current_original_column {
+      self.mappings.push(b'A');
+    } else {
+      encode_vlq(
+        &mut self.mappings,
+        original_column,
+        self.current_original_column,
+      );
+      self.current_original_column = original_column;
+    }
+  }
+
+  pub(crate) fn encode_raw(
+    &mut self,
+    generated_line: u32,
+    generated_column: u32,
+    original: Option<OriginalLocationInfo>,
+  ) {
+    if self.active_mapping && self.current_line == generated_line {
       // A mapping is still active
-      if mapping.original.as_ref().is_some_and(|original| {
+      if original.is_some_and(|original| {
         original.source_index == self.current_source_index
           && original.original_line == self.current_original_line
           && original.original_column == self.current_original_column
@@ -101,16 +182,16 @@ impl FullMappingsEncoder {
       }
     } else {
       // No mapping is active
-      if mapping.original.is_none() {
+      if original.is_none() {
         // avoid writing unnecessary generated mappings
         return;
       }
     }
 
-    if self.current_line < mapping.generated_line {
-      let count = (mapping.generated_line - self.current_line) as usize;
+    if self.current_line < generated_line {
+      let count = (generated_line - self.current_line) as usize;
       self.mappings.extend(std::iter::repeat_n(b';', count));
-      self.current_line = mapping.generated_line;
+      self.current_line = generated_line;
       self.current_column = 0;
       self.initial = false;
     } else if self.initial {
@@ -119,13 +200,9 @@ impl FullMappingsEncoder {
       self.mappings.push(b',');
     }
 
-    encode_vlq(
-      &mut self.mappings,
-      mapping.generated_column,
-      self.current_column,
-    );
-    self.current_column = mapping.generated_column;
-    if let Some(original) = &mapping.original {
+    encode_vlq(&mut self.mappings, generated_column, self.current_column);
+    self.current_column = generated_column;
+    if let Some(original) = original {
       self.active_mapping = true;
       if original.source_index == self.current_source_index {
         self.mappings.push(b'A');
@@ -165,8 +242,24 @@ impl FullMappingsEncoder {
     }
   }
 
+  fn encode(&mut self, mapping: &Mapping) {
+    self.encode_raw(
+      mapping.generated_line,
+      mapping.generated_column,
+      mapping
+        .original
+        .as_ref()
+        .map(|original| OriginalLocationInfo {
+          source_index: original.source_index,
+          original_line: original.original_line,
+          original_column: original.original_column,
+          name_index: original.name_index,
+        }),
+    );
+  }
+
   #[allow(unsafe_code)]
-  fn drain(&mut self) -> String {
+  pub(crate) fn drain(&mut self) -> String {
     unsafe {
       // SAFETY: The `mappings` field in the source map consists solely of ASCII characters.
       String::from_utf8_unchecked(std::mem::take(&mut self.mappings))

src/helpers.rs

@@ -161,6 +161,15 @@ pub fn utf16_len(s: &str) -> usize {
   simd_utf16_len::utf16_len(s)
 }
 
+#[inline]
+pub(crate) fn utf16_len_or_len(s: &str, is_ascii: bool) -> usize {
+  if is_ascii {
+    s.len()
+  } else {
+    utf16_len(s)
+  }
+}
+
 pub struct PotentialTokens<'a> {
   text: &'a str,
 }

src/original_source.rs

@@ -7,8 +7,8 @@ use std::{
 use crate::{
   helpers::{
     get_generated_source_info, get_map, split_into_lines,
-    split_into_potential_tokens, utf16_len, Chunks, GeneratedInfo,
-    StreamChunks,
+    split_into_potential_tokens, utf16_len, utf16_len_or_len, Chunks,
+    GeneratedInfo, StreamChunks,
   },
   object_pool::ObjectPool,
   source::{Mapping, OriginalLocation},
@@ -142,6 +142,7 @@ impl Chunks for OriginalSourceChunks<'_> {
     on_source(0, Cow::Borrowed(&self.0.name), Some(&self.0.value));
     if options.columns {
       // With column info we need to read all lines and split them
+      let is_ascii = self.0.value.is_ascii();
       let mut line = 1;
       let mut column = 0;
       for token in split_into_potential_tokens(self.0.value.as_ref()) {
@@ -176,7 +177,7 @@ impl Chunks for OriginalSourceChunks<'_> {
           line += 1;
           column = 0;
         } else {
-          column += utf16_len(token) as u32;
+          column += utf16_len_or_len(token, is_ascii) as u32;
         }
       }
       GeneratedInfo {