From 3b3654df56a5d151be7c7545e3d51c2afde19785 Mon Sep 17 00:00:00 2001
From: nym21 <expat-epochs0i@icloud.com>
Date: Sun, 13 Jul 2025 19:30:50 +0200
Subject: [PATCH] vec: add local and shared stored_len to raw variant

---
 crates/brk_vec/src/variants/compressed.rs | 32 +++++++++++------------
 crates/brk_vec/src/variants/raw.rs        | 27 ++++++++++++++-----
 2 files changed, 36 insertions(+), 23 deletions(-)

diff --git a/crates/brk_vec/src/variants/compressed.rs b/crates/brk_vec/src/variants/compressed.rs
index 5ced0c22c..d81485a0b 100644
--- a/crates/brk_vec/src/variants/compressed.rs
+++ b/crates/brk_vec/src/variants/compressed.rs
@@ -57,7 +57,7 @@ where
     }
 
     pub fn import(parent: &Path, name: &str, version: Version) -> Result<Self> {
-        let inner = RawVec::import(parent, name, version)?;
+        let mut inner = RawVec::import(parent, name, version)?;
 
         let pages_meta = {
             let path = inner
@@ -66,12 +66,19 @@ where
             if inner.is_empty() {
                 let _ = fs::remove_file(&path);
             }
-            Arc::new(ArcSwap::new(Arc::new(CompressedPagesMetadata::read(
-                &path,
-            )?)))
+            CompressedPagesMetadata::read(&path)?
         };
 
-        Ok(Self { inner, pages_meta })
+        inner.set_stored_len(if let Some(last) = pages_meta.last() {
+            (pages_meta.len() - 1) * Self::PER_PAGE + last.values_len as usize
+        } else {
+            0
+        });
+
+        Ok(Self {
+            inner,
+            pages_meta: Arc::new(ArcSwap::new(Arc::new(pages_meta))),
+        })
     }
 
     fn decode_page(&self, page_index: usize, mmap: &Mmap) -> Result<Vec<T>> {
@@ -130,14 +137,6 @@ where
         page_index * Self::PER_PAGE
     }
 
-    fn stored_len__(pages_meta: &Guard<Arc<CompressedPagesMetadata>>) -> usize {
-        if let Some(last) = pages_meta.last() {
-            (pages_meta.len() - 1) * Self::PER_PAGE + last.values_len as usize
-        } else {
-            0
-        }
-    }
-
     #[inline]
     pub fn iter(&self) -> CompressedVecIterator<'_, I, T> {
         self.into_iter()
@@ -186,7 +185,7 @@ where
 
     #[inline]
     fn stored_len(&self) -> usize {
-        Self::stored_len__(&self.pages_meta.load())
+        self.inner.stored_len()
     }
 
     #[inline]
@@ -486,14 +485,15 @@ where
 
     fn into_iter(self) -> Self::IntoIter {
         let pages_meta = self.pages_meta.load();
-        let stored_len = CompressedVec::<I, T>::stored_len__(&pages_meta);
+        let stored_len = self.stored_len();
+
         CompressedVecIterator {
             vec: self,
             mmap: self.create_mmap().unwrap(),
             decoded_page: None,
             pages_meta,
-            stored_len,
             index: 0,
+            stored_len,
         }
     }
 }
diff --git a/crates/brk_vec/src/variants/raw.rs b/crates/brk_vec/src/variants/raw.rs
index f978e840f..5847979a8 100644
--- a/crates/brk_vec/src/variants/raw.rs
+++ b/crates/brk_vec/src/variants/raw.rs
@@ -28,7 +28,8 @@ pub struct RawVec<I, T> {
     parent: PathBuf,
     name: &'static str,
     pushed: Vec<T>,
-    stored_len: Arc<AtomicUsize>,
+    local_stored_len: Option<usize>,
+    shared_stored_len: Arc<AtomicUsize>,
     phantom: PhantomData<I>,
 }
 
@@ -94,11 +95,17 @@ where
             name: Box::leak(Box::new(name.to_string())),
             parent: parent.to_owned(),
             pushed: vec![],
-            stored_len: Arc::new(AtomicUsize::new(stored_len)),
+            local_stored_len: Some(stored_len),
+            shared_stored_len: Arc::new(AtomicUsize::new(stored_len)),
             phantom: PhantomData,
         })
     }
 
+    pub fn set_stored_len(&mut self, len: usize) {
+        self.local_stored_len.replace(len);
+        self.shared_stored_len.store(len, Ordering::Relaxed);
+    }
+
     #[inline]
     pub fn iter(&self) -> RawVecIterator<'_, I, T> {
         self.into_iter()
@@ -151,7 +158,8 @@ where
 
     #[inline]
     fn stored_len(&self) -> usize {
-        self.stored_len.load(Ordering::SeqCst)
+        self.local_stored_len
+            .unwrap_or_else(|| self.shared_stored_len.load(Ordering::Relaxed))
     }
 
     #[inline]
@@ -195,7 +203,11 @@ where
         let mut file = file_opt.unwrap_or(self.open_file()?);
         self.file_write_all(&mut file, &bytes)?;
 
-        self.stored_len.fetch_add(pushed_len, Ordering::SeqCst);
+        if let Some(local_stored_len) = self.local_stored_len.as_mut() {
+            *local_stored_len += pushed_len;
+        }
+        self.shared_stored_len
+            .fetch_add(pushed_len, Ordering::Relaxed);
 
         Ok(())
     }
@@ -212,7 +224,7 @@ where
             return Ok(());
         }
 
-        self.stored_len.store(index, Ordering::SeqCst);
+        self.set_stored_len(index);
 
         let len = index * Self::SIZE_OF_T + HEADER_OFFSET;
 
@@ -223,7 +235,7 @@ where
     }
 
     fn reset(&mut self) -> Result<()> {
-        self.stored_len.store(0, Ordering::SeqCst);
+        self.set_stored_len(0);
         self.reset_()
     }
 }
@@ -267,7 +279,8 @@ impl<I, T> Clone for RawVec<I, T> {
             name: self.name,
             pushed: vec![],
             phantom: PhantomData,
-            stored_len: self.stored_len.clone(),
+            local_stored_len: None,
+            shared_stored_len: self.shared_stored_len.clone(),
         }
     }
 }