fix: replace size compute (#204)

Author: SyMind

.github/workflows/Bench.yaml

@@ -29,7 +29,7 @@ jobs:
       - name: Install codspeed
         uses: taiki-e/install-action@v2
         with:
-          tool: cargo-codspeed
+          tool: cargo-codspeed@4.0.5
 
       - name: Build Benchmark
         run: cargo codspeed build --features codspeed

benches/benchmark_repetitive_react_components.rs

@@ -9,9 +9,9 @@ pub use criterion::*;
 pub use codspeed_criterion_compat::*;
 
 use rspack_sources::{
-  BoxSource, CachedSource, ConcatSource, MapOptions, ObjectPool,
-  OriginalSource, RawStringSource, ReplaceSource, ReplacementEnforce, Source,
-  SourceExt, SourceMap, SourceMapSource, SourceMapSourceOptions,
+  BoxSource, ConcatSource, MapOptions, ObjectPool, OriginalSource,
+  RawStringSource, ReplaceSource, ReplacementEnforce, Source, SourceExt,
+  SourceMap, SourceMapSource, SourceMapSourceOptions,
 };
 
 static REPETITIVE_1K_REACT_COMPONENTS_SOURCE: LazyLock<BoxSource> =

src/cached_source.rs

@@ -1,6 +1,5 @@
 use std::{
   borrow::Cow,
-  cell::OnceCell,
   hash::{Hash, Hasher},
   sync::{Arc, OnceLock},
 };
@@ -21,6 +20,7 @@ use crate::{
 struct CachedData {
   hash: OnceLock<u64>,
   size: OnceLock<usize>,
+  chunks: OnceLock<Vec<&'static str>>,
   columns_map: OnceLock<Option<SourceMap>>,
   line_only_map: OnceLock<Option<SourceMap>>,
 }
@@ -79,19 +79,50 @@ impl CachedSource {
       cache: Arc::new(CachedData::default()),
     }
   }
+
+  fn get_or_init_chunks(&self) -> &[&str] {
+    self.cache.chunks.get_or_init(|| {
+      let mut chunks = Vec::new();
+      self.inner.rope(&mut |chunk| {
+        chunks.push(chunk);
+      });
+      #[allow(unsafe_code)]
+      // SAFETY: CachedSource guarantees that the underlying source outlives the cache,
+      // so transmuting Vec<&str> to Vec<&'static str> is safe in this context.
+      // This allows us to store string slices in the cache without additional allocations.
+      unsafe {
+        std::mem::transmute::<Vec<&str>, Vec<&'static str>>(chunks)
+      }
+    })
+  }
 }
 
 impl Source for CachedSource {
   fn source(&self) -> SourceValue {
-    self.inner.source()
+    let chunks = self.get_or_init_chunks();
+    let mut string = String::with_capacity(self.size());
+    for chunk in chunks {
+      string.push_str(chunk);
+    }
+    SourceValue::String(Cow::Owned(string))
+  }
+
+  fn rope<'a>(&'a self, on_chunk: &mut dyn FnMut(&'a str)) {
+    let chunks = self.get_or_init_chunks();
+    chunks.iter().for_each(|chunk| on_chunk(chunk));
   }
 
   fn buffer(&self) -> Cow<[u8]> {
     self.inner.buffer()
   }
 
   fn size(&self) -> usize {
-    *self.cache.size.get_or_init(|| self.inner.size())
+    *self.cache.size.get_or_init(|| {
+      if let Some(chunks) = self.cache.chunks.get() {
+        return chunks.iter().fold(0, |acc, chunk| acc + chunk.len());
+      }
+      self.inner.size()
+    })
   }
 
   fn map(
@@ -114,10 +145,6 @@ impl Source for CachedSource {
     }
   }
 
-  fn write_to_string(&self, string: &mut String) {
-    self.inner.write_to_string(string);
-  }
-
   fn to_writer(&self, writer: &mut dyn std::io::Write) -> std::io::Result<()> {
     self.inner.to_writer(writer)
   }
@@ -126,17 +153,17 @@ impl Source for CachedSource {
 struct CachedSourceChunks<'source> {
   chunks: Box<dyn Chunks + 'source>,
   cache: Arc<CachedData>,
-  inner: &'source dyn Source,
-  source: OnceCell<Cow<'source, str>>,
+  source: Cow<'source, str>,
 }
 
 impl<'a> CachedSourceChunks<'a> {
   fn new(cache_source: &'a CachedSource) -> Self {
+    let source = cache_source.source().into_string_lossy();
+
     Self {
       chunks: cache_source.inner.stream_chunks(),
       cache: cache_source.cache.clone(),
-      inner: &cache_source.inner,
-      source: OnceCell::new(),
+      source,
     }
   }
 }
@@ -157,22 +184,19 @@ impl Chunks for CachedSourceChunks<'_> {
     };
     match cell.get() {
       Some(map) => {
-        let source = self
-          .source
-          .get_or_init(|| self.inner.source().into_string_lossy());
         if let Some(map) = map {
           stream_chunks_of_source_map(
             options,
             object_pool,
-            source.as_ref(),
+            self.source.as_ref(),
             map,
             on_chunk,
             on_source,
             on_name,
           )
         } else {
           stream_chunks_of_raw_source(
-            source.as_ref(),
+            self.source.as_ref(),
             options,
             on_chunk,
             on_source,

src/concat_source.rs

@@ -165,14 +165,24 @@ impl Source for ConcatSource {
   fn source(&self) -> SourceValue {
     let children = self.optimized_children();
     if children.len() == 1 {
-      children[0].source()
-    } else {
-      // Use to_writer to avoid multiple heap allocations that would occur
-      // when concatenating nested ConcatSource instances directly
-      let mut string = String::with_capacity(self.size());
-      self.write_to_string(&mut string);
-      SourceValue::String(Cow::Owned(string))
+      return children[0].source();
     }
+
+    let mut string = String::with_capacity(self.size());
+    let mut on_chunk = |chunk| {
+      string.push_str(chunk);
+    };
+    children.iter().for_each(|child| {
+      child.rope(&mut on_chunk);
+    });
+    SourceValue::String(Cow::Owned(string))
+  }
+
+  fn rope<'a>(&'a self, on_chunk: &mut dyn FnMut(&'a str)) {
+    let children = self.optimized_children();
+    children.iter().for_each(|child| {
+      child.rope(on_chunk);
+    });
   }
 
   fn buffer(&self) -> Cow<[u8]> {
@@ -206,12 +216,6 @@ impl Source for ConcatSource {
     result
   }
 
-  fn write_to_string(&self, string: &mut String) {
-    for child in self.optimized_children() {
-      child.write_to_string(string);
-    }
-  }
-
   fn to_writer(&self, writer: &mut dyn std::io::Write) -> std::io::Result<()> {
     for child in self.optimized_children() {
       child.to_writer(writer)?;

src/helpers.rs

@@ -460,10 +460,9 @@ fn stream_chunks_of_source_map_full<'a>(
   on_source: OnSource<'_, 'a>,
   on_name: OnName<'_, 'a>,
 ) -> GeneratedInfo {
-  let a = split_into_lines(source);
-  let lines: Vec<WithUtf16<'a, 'a>> = a
+  let lines = split_into_lines(source)
     .map(|line| WithUtf16::new(object_pool, line))
-    .collect::<Vec<_>>();
+    .collect::<Vec<WithUtf16<'a, 'a>>>();
 
   if lines.is_empty() {
     return GeneratedInfo {

src/original_source.rs

@@ -56,6 +56,10 @@ impl Source for OriginalSource {
     SourceValue::String(Cow::Borrowed(&self.value))
   }
 
+  fn rope<'a>(&'a self, on_chunk: &mut dyn FnMut(&'a str)) {
+    on_chunk(self.value.as_ref())
+  }
+
   fn buffer(&self) -> Cow<[u8]> {
     Cow::Borrowed(self.value.as_bytes())
   }
@@ -73,10 +77,6 @@ impl Source for OriginalSource {
     get_map(object_pool, chunks.as_ref(), options)
   }
 
-  fn write_to_string(&self, string: &mut String) {
-    string.push_str(self.value.as_ref());
-  }
-
   fn to_writer(&self, writer: &mut dyn std::io::Write) -> std::io::Result<()> {
     writer.write_all(self.value.as_bytes())
   }

src/raw_source.rs

@@ -64,6 +64,10 @@ impl Source for RawStringSource {
     SourceValue::String(Cow::Borrowed(&self.0))
   }
 
+  fn rope<'a>(&'a self, on_chunk: &mut dyn FnMut(&'a str)) {
+    on_chunk(self.0.as_ref())
+  }
+
   fn buffer(&self) -> Cow<[u8]> {
     Cow::Borrowed(self.0.as_bytes())
   }
@@ -76,10 +80,6 @@ impl Source for RawStringSource {
     None
   }
 
-  fn write_to_string(&self, string: &mut String) {
-    string.push_str(self.0.as_ref());
-  }
-
   fn to_writer(&self, writer: &mut dyn std::io::Write) -> std::io::Result<()> {
     writer.write_all(self.0.as_bytes())
   }
@@ -210,6 +210,10 @@ impl Source for RawBufferSource {
     SourceValue::Buffer(Cow::Borrowed(&self.value))
   }
 
+  fn rope<'a>(&'a self, on_chunk: &mut dyn FnMut(&'a str)) {
+    on_chunk(self.get_or_init_value_as_string())
+  }
+
   fn buffer(&self) -> Cow<[u8]> {
     Cow::Borrowed(&self.value)
   }
@@ -222,10 +226,6 @@ impl Source for RawBufferSource {
     None
   }
 
-  fn write_to_string(&self, string: &mut String) {
-    string.push_str(self.get_or_init_value_as_string());
-  }
-
   fn to_writer(&self, writer: &mut dyn std::io::Write) -> std::io::Result<()> {
     writer.write_all(&self.value)
   }