perf(module-graph): compact connection overlay

Author: hardfist

crates/rspack_core/src/module_graph/rollback/dense_dependency_id_overlay_map.rs

@@ -1,10 +1,65 @@
 use super::OverlayValue;
 use crate::DependencyId;
 
+const UNCHANGED: u32 = 0;
+
+#[derive(Debug, Clone)]
+struct DenseOverlay<V> {
+  /// Maps a dependency id to an index in `values`. Zero means the base value is unchanged.
+  value_indices: Vec<u32>,
+  /// Stores only values that differ from the base map.
+  values: Vec<OverlayValue<V>>,
+}
+
+impl<V> Default for DenseOverlay<V> {
+  fn default() -> Self {
+    Self {
+      value_indices: Vec::new(),
+      values: Vec::new(),
+    }
+  }
+}
+
+impl<V> DenseOverlay<V> {
+  #[inline]
+  fn get(&self, index: usize) -> Option<&OverlayValue<V>> {
+    let value_index = *self.value_indices.get(index)?;
+    if value_index == UNCHANGED {
+      return None;
+    }
+    self.values.get(value_index as usize - 1)
+  }
+
+  #[inline]
+  fn get_mut(&mut self, index: usize) -> Option<&mut OverlayValue<V>> {
+    let value_index = *self.value_indices.get(index)?;
+    if value_index == UNCHANGED {
+      return None;
+    }
+    self.values.get_mut(value_index as usize - 1)
+  }
+
+  #[inline]
+  fn insert(&mut self, index: usize, value: OverlayValue<V>) {
+    if self.value_indices.len() <= index {
+      self.value_indices.resize(index + 1, UNCHANGED);
+    }
+
+    let value_index = &mut self.value_indices[index];
+    if *value_index == UNCHANGED {
+      self.values.push(value);
+      *value_index = u32::try_from(self.values.len())
+        .expect("dense dependency overlay cannot contain more than u32::MAX values");
+    } else {
+      self.values[*value_index as usize - 1] = value;
+    }
+  }
+}
+
 #[derive(Debug, Clone)]
 pub struct DenseDependencyIdOverlayMap<V> {
   base: Vec<Option<V>>,
-  overlay: Option<Vec<Option<OverlayValue<V>>>>,
+  overlay: Option<DenseOverlay<V>>,
 }
 
 impl<V> Default for DenseDependencyIdOverlayMap<V> {
@@ -19,7 +74,7 @@ impl<V> Default for DenseDependencyIdOverlayMap<V> {
 impl<V> DenseDependencyIdOverlayMap<V> {
   #[inline]
   pub fn checkpoint(&mut self) {
-    self.overlay.get_or_insert_with(Vec::new);
+    self.overlay.get_or_insert_default();
   }
 
   #[inline]
@@ -30,10 +85,8 @@ impl<V> DenseDependencyIdOverlayMap<V> {
   #[inline]
   pub fn insert(&mut self, key: DependencyId, value: V) {
     let index = key.as_u32() as usize;
-    if self.overlay.is_some() {
-      Self::ensure_len(self.overlay(), index);
-      self.overlay.as_mut().expect("overlay checked above")[index] =
-        Some(OverlayValue::Value(value));
+    if let Some(overlay) = &mut self.overlay {
+      overlay.insert(index, OverlayValue::Value(value));
     } else {
       Self::ensure_len(&mut self.base, index);
       self.base[index] = Some(value);
@@ -43,9 +96,8 @@ impl<V> DenseDependencyIdOverlayMap<V> {
   #[inline]
   pub fn remove(&mut self, key: &DependencyId) {
     let index = key.as_u32() as usize;
-    if self.overlay.is_some() {
-      Self::ensure_len(self.overlay(), index);
-      self.overlay.as_mut().expect("overlay checked above")[index] = Some(OverlayValue::Tombstone);
+    if let Some(overlay) = &mut self.overlay {
+      overlay.insert(index, OverlayValue::Tombstone);
     } else if let Some(value) = self.base.get_mut(index) {
       *value = None;
     }
@@ -55,7 +107,7 @@ impl<V> DenseDependencyIdOverlayMap<V> {
   pub fn get(&self, key: &DependencyId) -> Option<&V> {
     let index = key.as_u32() as usize;
     if let Some(overlay) = &self.overlay
-      && let Some(Some(value)) = overlay.get(index)
+      && let Some(value) = overlay.get(index)
     {
       return match value {
         OverlayValue::Value(value) => Some(value),
@@ -74,7 +126,7 @@ impl<V> DenseDependencyIdOverlayMap<V> {
     if self.overlay.is_some() {
       self.materialize_overlay_value(index);
       let overlay = self.overlay.as_mut().expect("overlay checked above");
-      match overlay.get_mut(index).and_then(Option::as_mut) {
+      match overlay.get_mut(index) {
         Some(OverlayValue::Value(value)) => Some(value),
         _ => None,
       }
@@ -89,20 +141,18 @@ impl<V> DenseDependencyIdOverlayMap<V> {
     V: Clone,
   {
     let overlay = self.overlay.as_ref().expect("overlay checked above");
-    if matches!(overlay.get(index), Some(Some(_))) {
+    if overlay.get(index).is_some() {
       return;
     }
 
     if let Some(value) = self.base.get(index).and_then(Option::as_ref).cloned() {
-      Self::ensure_len(self.overlay(), index);
-      self.overlay.as_mut().expect("overlay checked above")[index] =
-        Some(OverlayValue::Value(value));
+      self.overlay().insert(index, OverlayValue::Value(value));
     }
   }
 
   #[inline]
-  fn overlay(&mut self) -> &mut Vec<Option<OverlayValue<V>>> {
-    self.overlay.get_or_insert_with(Vec::new)
+  fn overlay(&mut self) -> &mut DenseOverlay<V> {
+    self.overlay.get_or_insert_default()
   }
 
   #[inline]
@@ -172,4 +222,34 @@ mod tests {
 
     assert_eq!(map.get(&a), Some(&1));
   }
+
+  #[test]
+  fn repeated_overlay_writes_reuse_the_value_slot() {
+    let mut map = DenseDependencyIdOverlayMap::default();
+    let id = DependencyId::from(7);
+
+    map.checkpoint();
+    map.insert(id, 1);
+    map.insert(id, 2);
+    map.remove(&id);
+
+    let overlay = map.overlay.as_ref().expect("should have overlay");
+    assert_eq!(overlay.value_indices.len(), 8);
+    assert_eq!(overlay.values.len(), 1);
+    assert_eq!(map.get(&id), None);
+  }
+
+  #[test]
+  fn sparse_overlay_stores_only_changed_values() {
+    let mut map = DenseDependencyIdOverlayMap::default();
+    let id = DependencyId::from(1024);
+
+    map.checkpoint();
+    map.insert(id, 1);
+
+    let overlay = map.overlay.as_ref().expect("should have overlay");
+    assert_eq!(overlay.value_indices.len(), 1025);
+    assert_eq!(overlay.values.len(), 1);
+    assert_eq!(map.get(&id), Some(&1));
+  }
 }