perf: refactor rope to iter

Author: SyMind

src/cached_source.rs

@@ -90,13 +90,15 @@ impl Source for CachedSource {
       // 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())
+        std::mem::transmute::<Vec<&str>, Vec<&'static str>>(
+          self.rope().collect(),
+        )
       }
     });
     SourceValue::String(Cow::Owned(rope.join("")))
   }
 
-  fn rope(&self) -> Vec<&str> {
+  fn rope(&self) -> Box<dyn Iterator<Item = &str> + '_> {
     self.inner.rope()
   }
 

src/concat_source.rs

@@ -175,15 +175,12 @@ impl Source for ConcatSource {
     }
   }
 
-  fn rope(&self) -> Vec<&str> {
+  fn rope(&self) -> Box<dyn Iterator<Item = &str> + '_> {
     let children = self.optimized_children();
     if children.len() == 1 {
       children[0].rope()
     } else {
-      children
-        .iter()
-        .flat_map(|child| child.rope())
-        .collect::<Vec<_>>()
+      Box::new(children.iter().flat_map(|child| child.rope()))
     }
   }
 

src/original_source.rs

@@ -56,8 +56,8 @@ impl Source for OriginalSource {
     SourceValue::String(Cow::Borrowed(&self.value))
   }
 
-  fn rope(&self) -> Vec<&str> {
-    vec![self.value.as_ref()]
+  fn rope(&self) -> Box<dyn Iterator<Item = &str> + '_> {
+    Box::new(std::iter::once(self.value.as_ref()))
   }
 
   fn buffer(&self) -> Cow<[u8]> {

src/raw_source.rs

@@ -64,8 +64,8 @@ impl Source for RawStringSource {
     SourceValue::String(Cow::Borrowed(&self.0))
   }
 
-  fn rope(&self) -> Vec<&str> {
-    vec![self.0.as_ref()]
+  fn rope(&self) -> Box<dyn Iterator<Item = &str> + '_> {
+    Box::new(std::iter::once(self.0.as_ref()))
   }
 
   fn buffer(&self) -> Cow<[u8]> {
@@ -214,8 +214,8 @@ impl Source for RawBufferSource {
     SourceValue::Buffer(Cow::Borrowed(&self.value))
   }
 
-  fn rope(&self) -> Vec<&str> {
-    vec![self.get_or_init_value_as_string()]
+  fn rope(&self) -> Box<dyn Iterator<Item = &str> + '_> {
+    Box::new(std::iter::once(self.get_or_init_value_as_string()))
   }
 
   fn buffer(&self) -> Cow<[u8]> {

src/replace_source.rs

@@ -161,109 +161,24 @@ impl ReplaceSource {
 
 impl Source for ReplaceSource {
   fn source(&self) -> SourceValue {
-    let rope = self.rope();
-    if rope.len() == 1 {
-      SourceValue::String(Cow::Borrowed(rope[0]))
-    } else {
-      SourceValue::String(Cow::Owned(rope.join("")))
+    if self.replacements.is_empty() {
+      return self.inner.source();
     }
-  }
-
-  fn rope(&self) -> Vec<&str> {
-    let inner_rope = self.inner.rope();
-    let mut rope =
-      Vec::with_capacity(inner_rope.len() + self.replacements.len() * 2);
-
-    let mut pos: usize = 0;
-    let mut replacement_idx: usize = 0;
-    let mut replacement_end: Option<usize> = None;
-    let mut next_replacement: Option<usize> = (replacement_idx
-      < self.replacements.len())
-    .then(|| self.replacements[replacement_idx].start as usize);
-
-    'chunk_loop: for chunk in inner_rope {
-      let mut chunk_pos = 0;
-      let end_pos = pos + chunk.len();
-
-      // Skip over when it has been replaced
-      if let Some(replacement_end) =
-        replacement_end.filter(|replacement_end| *replacement_end > pos)
-      {
-        // Skip over the whole chunk
-        if replacement_end >= end_pos {
-          pos = end_pos;
-          continue;
-        }
-        // Partially skip over chunk
-        chunk_pos = replacement_end - pos;
-        pos += chunk_pos;
-      }
-
-      // Is a replacement in the chunk?
-      while let Some(next_replacement_pos) = next_replacement
-        .filter(|next_replacement_pos| *next_replacement_pos < end_pos)
-      {
-        if next_replacement_pos > pos {
-          // Emit chunk until replacement
-          let offset = next_replacement_pos - pos;
-          let chunk_slice = &chunk[chunk_pos..(chunk_pos + offset)];
-          rope.push(chunk_slice);
-          chunk_pos += offset;
-          pos = next_replacement_pos;
-        }
-        // Insert replacement content split into chunks by lines
-        let replacement = &self.replacements[replacement_idx];
-        rope.push(&replacement.content);
-
-        // Remove replaced content by settings this variable
-        replacement_end = if let Some(replacement_end) = replacement_end {
-          Some(replacement_end.max(replacement.end as usize))
-        } else {
-          Some(replacement.end as usize)
-        };
-
-        // Move to next replacement
-        replacement_idx += 1;
-        next_replacement = if replacement_idx < self.replacements.len() {
-          Some(self.replacements[replacement_idx].start as usize)
-        } else {
-          None
-        };
-
-        // Skip over when it has been replaced
-        let offset = chunk.len() as i64 - end_pos as i64
-          + replacement_end.unwrap() as i64
-          - chunk_pos as i64;
-        if offset > 0 {
-          // Skip over whole chunk
-          if replacement_end
-            .is_some_and(|replacement_end| replacement_end >= end_pos)
-          {
-            pos = end_pos;
-            continue 'chunk_loop;
-          }
-
-          // Partially skip over chunk
-          chunk_pos += offset as usize;
-          pos += offset as usize;
-        }
-      }
-
-      // Emit remaining chunk
-      if chunk_pos < chunk.len() {
-        rope.push(&chunk[chunk_pos..]);
-      }
-      pos = end_pos;
+    let mut string = String::with_capacity(self.size());
+    for chunk in self.rope() {
+      string.push_str(chunk);
     }
+    SourceValue::String(Cow::Owned(string))
+  }
 
-    // Handle remaining replacements one by one
-    while replacement_idx < self.replacements.len() {
-      let content = &self.replacements[replacement_idx].content;
-      rope.push(content);
-      replacement_idx += 1;
+  fn rope(&self) -> Box<dyn Iterator<Item = &str> + '_> {
+    if self.replacements.is_empty() {
+      return self.inner.rope();
     }
-
-    rope
+    Box::new(ReplaceSourceRopeIterator::new(
+      self.inner.rope(),
+      &self.replacements,
+    ))
   }
 
   fn buffer(&self) -> Cow<[u8]> {
@@ -897,6 +812,191 @@ impl PartialEq for ReplaceSource {
 
 impl Eq for ReplaceSource {}
 
+/// Iterator for ReplaceSource rope that applies replacements on the fly
+pub struct ReplaceSourceRopeIterator<'a> {
+  inner_rope: Box<dyn Iterator<Item = &'a str> + 'a>,
+  replacements: &'a [Replacement],
+
+  // Current state
+  pos: usize,
+  replacement_idx: usize,
+  replacement_end: Option<usize>,
+
+  // Current chunk state
+  current_chunk: Option<&'a str>,
+  chunk_pos: usize,
+  chunk_start_pos: usize,
+
+  // Buffer for remaining replacements
+  remaining_replacements: std::collections::VecDeque<&'a str>,
+
+  // Flag to indicate if we're done with inner rope
+  inner_exhausted: bool,
+}
+
+impl<'a> ReplaceSourceRopeIterator<'a> {
+  fn new(
+    inner_rope: Box<dyn Iterator<Item = &'a str> + 'a>,
+    replacements: &'a [Replacement],
+  ) -> Self {
+    Self {
+      inner_rope,
+      replacements,
+      pos: 0,
+      replacement_idx: 0,
+      replacement_end: None,
+      current_chunk: None,
+      chunk_pos: 0,
+      chunk_start_pos: 0,
+      remaining_replacements: std::collections::VecDeque::new(),
+      inner_exhausted: false,
+    }
+  }
+
+  fn get_next_replacement_pos(&self) -> Option<usize> {
+    if self.replacement_idx < self.replacements.len() {
+      Some(self.replacements[self.replacement_idx].start as usize)
+    } else {
+      None
+    }
+  }
+
+  fn advance_to_next_chunk(&mut self) -> bool {
+    if let Some(chunk) = self.inner_rope.next() {
+      self.current_chunk = Some(chunk);
+      self.chunk_pos = 0;
+      self.chunk_start_pos = self.pos;
+      true
+    } else {
+      self.inner_exhausted = true;
+      false
+    }
+  }
+
+  fn process_current_chunk(&mut self) -> Option<&'a str> {
+    let chunk = self.current_chunk?;
+    let chunk_end_pos = self.chunk_start_pos + chunk.len();
+
+    // Skip over when it has been replaced
+    if let Some(replacement_end) =
+      self.replacement_end.filter(|&end| end > self.pos)
+    {
+      if replacement_end >= chunk_end_pos {
+        // Skip entire chunk
+        self.pos = chunk_end_pos;
+        self.current_chunk = None;
+        return None;
+      }
+
+      // Partially skip chunk
+      let skip_amount = replacement_end - self.pos;
+      self.chunk_pos += skip_amount;
+      self.pos += skip_amount;
+    }
+
+    // Check for replacements in current chunk
+    if let Some(next_replacement_pos) = self
+      .get_next_replacement_pos()
+      .filter(|&pos| pos < chunk_end_pos)
+    {
+      if next_replacement_pos > self.pos {
+        // Emit chunk until replacement
+        let offset = next_replacement_pos - self.pos;
+        let chunk_slice = &chunk[self.chunk_pos..(self.chunk_pos + offset)];
+        self.chunk_pos += offset;
+        self.pos = next_replacement_pos;
+        return Some(chunk_slice);
+      }
+
+      // Add replacement content
+      let replacement = &self.replacements[self.replacement_idx];
+      self.remaining_replacements.push_back(&replacement.content);
+
+      // Update replacement_end
+      self.replacement_end = Some(
+        self
+          .replacement_end
+          .map(|end| end.max(replacement.end as usize))
+          .unwrap_or(replacement.end as usize),
+      );
+
+      // Move to next replacement
+      self.replacement_idx += 1;
+
+      // Skip replaced content
+      let offset = chunk.len() as i64 - chunk_end_pos as i64
+        + self.replacement_end.unwrap() as i64
+        - self.chunk_pos as i64;
+
+      if offset > 0 {
+        if self.replacement_end.is_some_and(|end| end >= chunk_end_pos) {
+          // Skip entire remaining chunk
+          self.pos = chunk_end_pos;
+          self.current_chunk = None;
+          return self.remaining_replacements.pop_front();
+        }
+
+        // Partially skip chunk
+        self.chunk_pos += offset as usize;
+        self.pos += offset as usize;
+      }
+
+      // Return the replacement content
+      return self.remaining_replacements.pop_front();
+    }
+
+    // Emit remaining chunk
+    if self.chunk_pos < chunk.len() {
+      let remaining = &chunk[self.chunk_pos..];
+      self.pos = chunk_end_pos;
+      self.current_chunk = None;
+      Some(remaining)
+    } else {
+      self.pos = chunk_end_pos;
+      self.current_chunk = None;
+      None
+    }
+  }
+}
+
+impl<'a> Iterator for ReplaceSourceRopeIterator<'a> {
+  type Item = &'a str;
+
+  fn next(&mut self) -> Option<Self::Item> {
+    // First, check if we have buffered replacement content
+    if let Some(replacement) = self.remaining_replacements.pop_front() {
+      return Some(replacement);
+    }
+
+    // Process current chunk if we have one
+    if self.current_chunk.is_some() {
+      if let Some(result) = self.process_current_chunk() {
+        return Some(result);
+      }
+    }
+
+    // Advance to next chunk if inner rope is not exhausted
+    while !self.inner_exhausted {
+      if !self.advance_to_next_chunk() {
+        break;
+      }
+
+      if let Some(result) = self.process_current_chunk() {
+        return Some(result);
+      }
+    }
+
+    // Handle remaining replacements when inner rope is exhausted
+    if self.replacement_idx < self.replacements.len() {
+      let replacement = &self.replacements[self.replacement_idx];
+      self.replacement_idx += 1;
+      Some(&replacement.content)
+    } else {
+      None
+    }
+  }
+}
+
 #[cfg(test)]
 mod tests {
   use rustc_hash::FxHasher;