perf: Rope enum

Author: SyMind

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

@@ -13,7 +13,7 @@ use crate::{
   },
   object_pool::ObjectPool,
   source::SourceValue,
-  BoxSource, MapOptions, Source, SourceExt, SourceMap,
+  BoxSource, MapOptions, Rope, Source, SourceExt, SourceMap,
 };
 
 #[derive(Default)]
@@ -82,14 +82,17 @@ impl CachedSource {
 
   fn get_or_init_chunks(&self) -> &[&str] {
     self.cache.chunks.get_or_init(|| {
+      let rope = self.inner.rope();
+      let chunks = match rope {
+        Rope::Light(s) => vec![s],
+        Rope::Full(iter) => iter.collect::<Vec<_>>(),
+      };
       #[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>>(
-          self.rope().collect(),
-        )
+        std::mem::transmute::<Vec<&str>, Vec<&'static str>>(chunks)
       }
     })
   }
@@ -98,19 +101,16 @@ impl CachedSource {
 impl Source for CachedSource {
   fn source(&self) -> SourceValue {
     let chunks = self.get_or_init_chunks();
-    if chunks.len() == 1 {
-      return SourceValue::String(Cow::Borrowed(chunks[0]));
-    }
     let mut string = String::with_capacity(self.size());
     for chunk in chunks {
       string.push_str(chunk);
     }
     SourceValue::String(Cow::Owned(string))
   }
 
-  fn rope(&self) -> Box<dyn Iterator<Item = &str> + '_> {
+  fn rope(&self) -> Rope {
     let chunks = self.get_or_init_chunks();
-    Box::new(chunks.iter().cloned())
+    Rope::Full(Box::new(chunks.iter().cloned()))
   }
 
   fn buffer(&self) -> Cow<[u8]> {

src/concat_source.rs

@@ -12,7 +12,7 @@ use crate::{
   linear_map::LinearMap,
   object_pool::ObjectPool,
   source::{Mapping, OriginalLocation},
-  BoxSource, MapOptions, RawStringSource, Source, SourceExt, SourceMap,
+  BoxSource, MapOptions, RawStringSource, Rope, Source, SourceExt, SourceMap,
   SourceValue,
 };
 
@@ -163,24 +163,29 @@ impl ConcatSource {
 
 impl Source for ConcatSource {
   fn source(&self) -> SourceValue {
-    let children = self.optimized_children();
-    if children.len() == 1 {
-      children[0].source()
-    } else {
-      let mut string = String::with_capacity(self.size());
-      for chunk in self.rope() {
-        string.push_str(chunk);
+    match self.rope() {
+      Rope::Light(s) => SourceValue::String(Cow::Borrowed(s)),
+      Rope::Full(iter) => {
+        let mut string = String::with_capacity(self.size());
+        for chunk in iter {
+          string.push_str(chunk);
+        }
+        SourceValue::String(Cow::Owned(string))
       }
-      SourceValue::String(Cow::Owned(string))
     }
   }
 
-  fn rope(&self) -> Box<dyn Iterator<Item = &str> + '_> {
+  fn rope(&self) -> Rope {
     let children = self.optimized_children();
     if children.len() == 1 {
       children[0].rope()
     } else {
-      Box::new(children.iter().flat_map(|child| child.rope()))
+      Rope::Full(Box::new(children.iter().flat_map(
+        |child| match child.rope() {
+          Rope::Light(s) => Box::new(std::iter::once(s)),
+          Rope::Full(iter) => iter,
+        },
+      )))
     }
   }
 

src/lib.rs

@@ -23,7 +23,7 @@ pub use original_source::OriginalSource;
 pub use raw_source::{RawBufferSource, RawStringSource};
 pub use replace_source::{ReplaceSource, ReplacementEnforce};
 pub use source::{
-  BoxSource, MapOptions, Mapping, OriginalLocation, Source, SourceExt,
+  BoxSource, MapOptions, Mapping, OriginalLocation, Rope, Source, SourceExt,
   SourceMap, SourceValue,
 };
 pub use source_map_source::{

src/original_source.rs

@@ -11,7 +11,7 @@ use crate::{
   },
   object_pool::ObjectPool,
   source::{Mapping, OriginalLocation},
-  MapOptions, Source, SourceMap, SourceValue,
+  MapOptions, Rope, Source, SourceMap, SourceValue,
 };
 
 /// Represents source code, it will create source map for the source code,
@@ -56,8 +56,8 @@ impl Source for OriginalSource {
     SourceValue::String(Cow::Borrowed(&self.value))
   }
 
-  fn rope(&self) -> Box<dyn Iterator<Item = &str> + '_> {
-    Box::new(std::iter::once(self.value.as_ref()))
+  fn rope(&self) -> Rope {
+    Rope::Light(self.value.as_ref())
   }
 
   fn buffer(&self) -> Cow<[u8]> {

src/raw_source.rs

@@ -10,7 +10,7 @@ use crate::{
     GeneratedInfo, StreamChunks,
   },
   object_pool::ObjectPool,
-  MapOptions, Source, SourceMap, SourceValue,
+  MapOptions, Rope, Source, SourceMap, SourceValue,
 };
 
 /// A string variant of [RawStringSource].
@@ -64,8 +64,8 @@ impl Source for RawStringSource {
     SourceValue::String(Cow::Borrowed(&self.0))
   }
 
-  fn rope(&self) -> Box<dyn Iterator<Item = &str> + '_> {
-    Box::new(std::iter::once(self.0.as_ref()))
+  fn rope(&self) -> Rope {
+    Rope::Light(self.0.as_ref())
   }
 
   fn buffer(&self) -> Cow<[u8]> {
@@ -210,8 +210,8 @@ impl Source for RawBufferSource {
     SourceValue::Buffer(Cow::Borrowed(&self.value))
   }
 
-  fn rope(&self) -> Box<dyn Iterator<Item = &str> + '_> {
-    Box::new(std::iter::once(self.get_or_init_value_as_string()))
+  fn rope(&self) -> Rope {
+    Rope::Light(self.get_or_init_value_as_string())
   }
 
   fn buffer(&self) -> Cow<[u8]> {

src/replace_source.rs

@@ -12,8 +12,8 @@ use crate::{
   linear_map::LinearMap,
   object_pool::ObjectPool,
   source_content_lines::SourceContentLines,
-  BoxSource, MapOptions, Mapping, OriginalLocation, OriginalSource, Source,
-  SourceExt, SourceMap, SourceValue,
+  BoxSource, MapOptions, Mapping, OriginalLocation, OriginalSource, Rope,
+  Source, SourceExt, SourceMap, SourceValue,
 };
 
 /// Decorates a Source with replacements and insertions of source code,
@@ -161,24 +161,30 @@ impl ReplaceSource {
 
 impl Source for ReplaceSource {
   fn source(&self) -> SourceValue {
-    if self.replacements.is_empty() {
-      return self.inner.source();
-    }
-    let mut string = String::with_capacity(self.size());
-    for chunk in self.rope() {
-      string.push_str(chunk);
+    match self.rope() {
+      Rope::Light(s) => SourceValue::String(Cow::Borrowed(s)),
+      Rope::Full(iter) => {
+        let mut string = String::with_capacity(self.size());
+        for chunk in iter {
+          string.push_str(chunk);
+        }
+        SourceValue::String(Cow::Owned(string))
+      }
     }
-    SourceValue::String(Cow::Owned(string))
   }
 
-  fn rope(&self) -> Box<dyn Iterator<Item = &str> + '_> {
+  fn rope(&self) -> Rope {
     if self.replacements.is_empty() {
       return self.inner.rope();
     }
-    Box::new(ReplaceSourceRopeIterator::new(
-      self.inner.rope(),
+    let inner_chunks = match self.inner.rope() {
+      Rope::Light(s) => Box::new(std::iter::once(s)),
+      Rope::Full(iter) => iter,
+    };
+    Rope::Full(Box::new(ReplaceSourceRopeIterator::new(
+      inner_chunks,
       &self.replacements,
-    ))
+    )))
   }
 
   fn buffer(&self) -> Cow<[u8]> {
@@ -239,8 +245,13 @@ impl Source for ReplaceSource {
   }
 
   fn to_writer(&self, writer: &mut dyn std::io::Write) -> std::io::Result<()> {
-    for text in self.rope() {
-      writer.write_all(text.as_bytes())?;
+    match self.rope() {
+      Rope::Light(s) => writer.write_all(s.as_bytes())?,
+      Rope::Full(iter) => {
+        for chunk in iter {
+          writer.write_all(chunk.as_bytes())?
+        }
+      }
     }
     Ok(())
   }

src/source.rs

@@ -107,6 +107,11 @@ impl<'a> SourceValue<'a> {
   }
 }
 
+pub enum Rope<'a> {
+  Light(&'a str),
+  Full(Box<dyn Iterator<Item = &'a str> + 'a>),
+}
+
 /// [Source] abstraction, [webpack-sources docs](https://github.com/webpack/webpack-sources/#source).
 pub trait Source:
   StreamChunks + DynHash + AsAny + DynEq + DynClone + fmt::Debug + Sync + Send
@@ -115,7 +120,7 @@ pub trait Source:
   fn source(&self) -> SourceValue;
 
   /// Return a lightweight "rope" view of the source as borrowed string slices.
-  fn rope(&self) -> Box<dyn Iterator<Item = &str> + '_>;
+  fn rope(&self) -> Rope;
 
   /// Get the source buffer.
   fn buffer(&self) -> Cow<[u8]>;
@@ -144,7 +149,7 @@ impl Source for BoxSource {
     self.as_ref().source()
   }
 
-  fn rope(&self) -> Box<dyn Iterator<Item = &str> + '_> {
+  fn rope(&self) -> Rope {
     self.as_ref().rope()
   }
 

src/source_map_source.rs

@@ -5,12 +5,9 @@ use std::{
 };
 
 use crate::{
-  helpers::{
-    get_map, stream_chunks_of_combined_source_map, stream_chunks_of_source_map,
-    Chunks, StreamChunks,
-  },
-  object_pool::ObjectPool,
-  MapOptions, Source, SourceMap, SourceValue,
+  MapOptions, Rope, Source, SourceMap, SourceValue, helpers::{
+    Chunks, StreamChunks, get_map, stream_chunks_of_combined_source_map, stream_chunks_of_source_map
+  }, object_pool::ObjectPool
 };
 
 /// Options for [SourceMapSource::new].
@@ -94,8 +91,8 @@ impl Source for SourceMapSource {
     SourceValue::String(Cow::Borrowed(&self.value))
   }
 
-  fn rope(&self) -> Box<dyn Iterator<Item = &str> + '_> {
-    Box::new(std::iter::once(self.value.as_ref()))
+  fn rope(&self) -> Rope {
+    Rope::Light(&self.value)
   }
 
   fn buffer(&self) -> Cow<[u8]> {