From ff1a008e27b42cb2d14109c37504d797f8576635 Mon Sep 17 00:00:00 2001 From: rymnc <43716372+rymnc@users.noreply.github.com> Date: Fri, 14 Aug 2026 03:13:39 +0530 Subject: [PATCH 1/2] fix(rotortree): cut some cruft, improve some perf --- crates/rotortree/src/chunked_level.rs | 400 ++++++++---------- crates/rotortree/src/storage/checkpoint.rs | 45 -- crates/rotortree/src/storage/mod.rs | 14 +- crates/rotortree/src/storage/recovery.rs | 15 +- crates/rotortree/src/tree.rs | 213 +++++----- crates/rotortree/tests/storage_integration.rs | 1 - 6 files changed, 283 insertions(+), 405 deletions(-) diff --git a/crates/rotortree/src/chunked_level.rs b/crates/rotortree/src/chunked_level.rs index b7e131e..8a8f010 100644 --- a/crates/rotortree/src/chunked_level.rs +++ b/crates/rotortree/src/chunked_level.rs @@ -10,6 +10,8 @@ use std::{ vec::Vec, }; +use core::mem::MaybeUninit; + use crate::{ Hash, TreeError, @@ -37,9 +39,10 @@ impl Chunk { Arc::make_mut(&mut self.0) } + /// Allocate a chunk and let `f` write its `CHUNK_SIZE` hashes in place. #[inline] - pub(crate) fn new_memory(data: [Hash; CHUNK_SIZE]) -> Self { - Self(Arc::new(data)) + pub(crate) fn from_fn(f: impl FnOnce(&mut [MaybeUninit])) -> Self { + Self(new_arc_from_fn(f)) } #[cfg(test)] @@ -48,6 +51,41 @@ impl Chunk { } } +/// Allocate a chunk and let `f` write its contents in place. +#[inline] +fn new_arc_from_fn(f: impl FnOnce(&mut [MaybeUninit])) -> Arc<[Hash; CHUNK_SIZE]> { + let mut arc = Arc::<[Hash; CHUNK_SIZE]>::new_uninit(); + // SAFETY: freshly allocated so uniquely owned, and `MaybeUninit` has + // the same layout as `Hash`. + let out = unsafe { + &mut *(Arc::get_mut(&mut arc).unwrap_unchecked().as_mut_ptr() + as *mut [MaybeUninit; CHUNK_SIZE]) + }; + f(out); + // SAFETY: `f` is contracted to initialise all CHUNK_SIZE elements. + unsafe { arc.assume_init() } +} + +/// Minimum chunks to parallelize. +#[cfg(feature = "parallel")] +const PAR_MIN_CHUNKS: usize = 16; + +#[inline] +fn build_chunks(n: usize, build: impl Fn(usize) -> Chunk + Send + Sync) -> Vec { + #[cfg(feature = "parallel")] + if n >= PAR_MIN_CHUNKS { + use rayon::prelude::*; + return (0..n).into_par_iter().map(build).collect(); + } + (0..n).map(build).collect() +} + +#[inline] +pub(crate) fn as_uninit_mut(s: &mut [Hash]) -> &mut [MaybeUninit] { + // SAFETY: `MaybeUninit` has the same layout as `Hash`. + unsafe { &mut *(s as *mut [Hash] as *mut [MaybeUninit]) } +} + #[cfg(feature = "storage")] #[derive(Clone)] pub(crate) struct Chunk(ChunkInner); @@ -87,9 +125,10 @@ impl Chunk { } } + /// Allocate a chunk and let `f` write its `CHUNK_SIZE` hashes in place. #[inline] - pub(crate) fn new_memory(data: [Hash; CHUNK_SIZE]) -> Self { - Self(ChunkInner::Memory(Arc::new(data))) + pub(crate) fn from_fn(f: impl FnOnce(&mut [MaybeUninit])) -> Self { + Self(ChunkInner::Memory(new_arc_from_fn(f))) } pub(crate) fn new_mapped( @@ -114,8 +153,7 @@ impl Chunk { } } -/// A single level of the tree stored as segmented chunks plus a -/// fixed-size tail buffer. +/// A single level of the tree, stored as segmented chunks. #[derive(Clone)] pub(crate) struct ChunkedLevel { /// Immutable segments of committed chunks, shared with snapshots. @@ -123,10 +161,6 @@ pub(crate) struct ChunkedLevel { /// Mutable buffer of committed chunks not yet frozen into a segment. /// At most `CHUNKS_PER_SEGMENT - 1` items. pending: Vec, - /// Fixed-size tail buffer (partially filled). - tail: [Hash; CHUNK_SIZE], - /// Number of valid entries in `tail`. - tail_len: usize, /// Total number of hashes in this level. len: usize, } @@ -136,8 +170,6 @@ impl ChunkedLevel { Self { segments: Vec::new(), pending: Vec::new(), - tail: [[0u8; 32]; CHUNK_SIZE], - tail_len: 0, len: 0, } } @@ -145,12 +177,7 @@ impl ChunkedLevel { /// Construct a level from checkpoint data, partitioning chunks into /// segments and pending. #[cfg(feature = "storage")] - pub(crate) fn from_parts( - chunks: Vec, - tail: [Hash; CHUNK_SIZE], - tail_len: usize, - len: usize, - ) -> Self { + pub(crate) fn from_parts(chunks: Vec, len: usize) -> Self { let full_segments = chunks.len() / CHUNKS_PER_SEGMENT; let mut segments = Vec::with_capacity(full_segments); let mut drain = chunks.into_iter(); @@ -167,8 +194,6 @@ impl ChunkedLevel { Self { segments, pending, - tail, - tail_len, len, } } @@ -179,25 +204,31 @@ impl ChunkedLevel { self.len } - /// Total number of committed chunks (segments + pending). + /// Total number of chunks, i.e. `len` rounded up to a whole chunk. + #[cfg(any(feature = "storage", test))] #[inline] pub(crate) fn chunk_count(&self) -> usize { self.segments.len() * CHUNKS_PER_SEGMENT + self.pending.len() } - /// Resolve a chunk index to a slice reference. + /// The only place that knows chunks live in frozen segments below + /// `committed` and in `pending` above it. #[inline(always)] - fn chunk_slice(&self, chunk_idx: usize) -> &[Hash; CHUNK_SIZE] { + pub(crate) fn chunk(&self, chunk_idx: usize) -> &Chunk { let committed = self.segments.len() * CHUNKS_PER_SEGMENT; if chunk_idx < committed { - let seg_idx = chunk_idx / CHUNKS_PER_SEGMENT; - let seg_off = chunk_idx % CHUNKS_PER_SEGMENT; - self.segments[seg_idx][seg_off].as_slice() + &self.segments[chunk_idx / CHUNKS_PER_SEGMENT] + [chunk_idx % CHUNKS_PER_SEGMENT] } else { - self.pending[chunk_idx - committed].as_slice() + &self.pending[chunk_idx - committed] } } + #[inline(always)] + fn chunk_slice(&self, chunk_idx: usize) -> &[Hash; CHUNK_SIZE] { + self.chunk(chunk_idx).as_slice() + } + /// Read a hash at the given index. #[inline] pub(crate) fn get(&self, index: usize) -> Result { @@ -207,57 +238,38 @@ impl ChunkedLevel { size: self.len as u64, }); } - let chunk_idx = index / CHUNK_SIZE; - let offset = index % CHUNK_SIZE; - if chunk_idx < self.chunk_count() { - Ok(self.chunk_slice(chunk_idx)[offset]) - } else { - Ok(self.tail[offset]) - } + Ok(self.chunk_slice(index / CHUNK_SIZE)[index % CHUNK_SIZE]) } - /// Copy a contiguous group of hashes into `out`. - /// Fast path when the group falls within a single chunk or tail. + /// Copy `[start, start + count)` into `out`. #[inline(always)] pub(crate) fn get_group(&self, start: usize, count: usize, out: &mut [Hash]) { - let chunk_idx = start / CHUNK_SIZE; - let offset = start % CHUNK_SIZE; - if offset + count <= CHUNK_SIZE { - let src = if chunk_idx < self.chunk_count() { - &self.chunk_slice(chunk_idx)[offset..offset + count] - } else { - &self.tail[offset..offset + count] - }; - out[..count].copy_from_slice(src); - } else { - for (i, item) in out.iter_mut().enumerate().take(count) { - *item = self.get(start + i).expect("checked prev; qed"); - } + let mut at = 0; + for run in self.runs(start, count) { + out[at..at + run.len()].copy_from_slice(run); + at += run.len(); } } - /// Borrow a contiguous group of `count` hashes starting at `start`, - /// when it lies wholly within a single chunk or the tail. - /// - /// Returns `None` if the group straddles a chunk/tail boundary, in which - /// case the caller must fall back to a copying read. For the batched - /// parent path, full groups of arity N in {2,4,8,16} are chunk-aligned - /// (`CHUNK_SIZE % N == 0`), so this always returns `Some`. + /// Borrow `[start, start + count)` as contiguous runs, one per chunk. #[inline] - pub(crate) fn group_slice(&self, start: usize, count: usize) -> Option<&[Hash]> { - if start + count > self.len { - return None; - } - let chunk_idx = start / CHUNK_SIZE; - let offset = start % CHUNK_SIZE; - if offset + count > CHUNK_SIZE { - return None; - } - if chunk_idx < self.chunk_count() { - Some(&self.chunk_slice(chunk_idx)[offset..offset + count]) - } else { - Some(&self.tail[offset..offset + count]) - } + pub(crate) fn runs(&self, start: usize, count: usize) -> impl Iterator { + debug_assert!( + start + count <= self.len, + "runs: {start}+{count} > len {}", + self.len + ); + let mut done = 0; + core::iter::from_fn(move || { + if done == count { + return None; + } + let idx = start + done; + let offset = idx % CHUNK_SIZE; + let take = (CHUNK_SIZE - offset).min(count - done); + done += take; + Some(&self.chunk_slice(idx / CHUNK_SIZE)[offset..offset + take]) + }) } /// Write a hash at the given index @@ -278,125 +290,119 @@ impl ChunkedLevel { "set_preallocated: index {index} >= len {}", self.len ); - let chunk_idx = index / CHUNK_SIZE; - let offset = index % CHUNK_SIZE; + self.chunk_slice_mut(index / CHUNK_SIZE)[index % CHUNK_SIZE] = value; + } + + /// Resolve a chunk index to a mutable slice, copy-on-writing it. + #[inline(always)] + fn chunk_slice_mut(&mut self, chunk_idx: usize) -> &mut [Hash; CHUNK_SIZE] { let committed = self.segments.len() * CHUNKS_PER_SEGMENT; if chunk_idx < committed { let seg_idx = chunk_idx / CHUNKS_PER_SEGMENT; let seg_off = chunk_idx % CHUNKS_PER_SEGMENT; - Arc::make_mut(&mut self.segments[seg_idx])[seg_off].make_mut()[offset] = - value; - } else if chunk_idx - committed < self.pending.len() { - self.pending[chunk_idx - committed].make_mut()[offset] = value; + Arc::make_mut(&mut self.segments[seg_idx])[seg_off].make_mut() } else { - self.tail[offset] = value; + self.pending[chunk_idx - committed].make_mut() } } - /// Append a hash. Promotes the tail when it reaches - /// `CHUNK_SIZE`. - #[cfg(test)] - #[inline] - pub(crate) fn push(&mut self, value: Hash) -> Result<(), TreeError> { - self.tail[self.tail_len] = value; - self.tail_len = self.tail_len.checked_add(1).ok_or(TreeError::MathError)?; - self.len = self.len.checked_add(1).ok_or(TreeError::MathError)?; - if self.tail_len == CHUNK_SIZE { - self.promote_tail(); + /// Overwrite `[start, start + count)` with hashes produced by `fill`. + pub(crate) fn write_with( + &mut self, + start: usize, + count: usize, + fill: impl Fn(usize, &mut [MaybeUninit]), + ) { + debug_assert!( + start + count <= self.len, + "write_with: {start}+{count} > len {}", + self.len + ); + let mut done = 0; + while done < count { + let idx = start + done; + let offset = idx % CHUNK_SIZE; + let take = (CHUNK_SIZE - offset).min(count - done); + let chunk = self.chunk_slice_mut(idx / CHUNK_SIZE); + fill(done, as_uninit_mut(&mut chunk[offset..offset + take])); + done += take; } - Ok(()) } - pub(crate) fn extend(&mut self, values: &[Hash]) -> Result<(), TreeError> { - if values.is_empty() { + /// Append `count` hashes produced by `fill`, writing each chunk straight + /// into its final allocation. + pub(crate) fn extend_with(&mut self, count: usize, fill: F) -> Result<(), TreeError> + where + F: Fn(usize, &mut [MaybeUninit]) + Sync, + { + if count == 0 { return Ok(()); } - let new_len = self - .len - .checked_add(values.len()) - .ok_or(TreeError::MathError)?; - - let mut remaining = values; - - // fill current tail - if self.tail_len > 0 { - let space = CHUNK_SIZE - self.tail_len; - let to_copy = space.min(remaining.len()); - self.tail[self.tail_len..self.tail_len + to_copy] - .copy_from_slice(&remaining[..to_copy]); - self.tail_len += to_copy; - remaining = &remaining[to_copy..]; - if self.tail_len == CHUNK_SIZE { - self.promote_tail(); - } - } - - // full chunks — bypass tail - let full_chunks = remaining.len() / CHUNK_SIZE; - if full_chunks > 0 { - self.pending.reserve(full_chunks.min(CHUNKS_PER_SEGMENT)); - for i in 0..full_chunks { - let start = i * CHUNK_SIZE; - let chunk: [Hash; CHUNK_SIZE] = remaining[start..start + CHUNK_SIZE] - .try_into() - .expect("slice len == CHUNK_SIZE; qed"); - self.push_chunk(Chunk::new_memory(chunk)); + let new_len = self.len.checked_add(count).ok_or(TreeError::MathError)?; + + let offset = self.len % CHUNK_SIZE; + let done = if offset > 0 { + let take = (CHUNK_SIZE - offset).min(count); + let chunk = self.chunk_slice_mut(self.len / CHUNK_SIZE); + fill(0, as_uninit_mut(&mut chunk[offset..offset + take])); + take + } else { + 0 + }; + + // fill into allocations + if done < count { + let n = (count - done).div_ceil(CHUNK_SIZE); + self.pending.reserve(n.min(CHUNKS_PER_SEGMENT)); + let built = build_chunks(n, |ci| { + let base = done + ci * CHUNK_SIZE; + let take = (count - base).min(CHUNK_SIZE); + Chunk::from_fn(|out| { + fill(base, &mut out[..take]); + out[take..].fill(MaybeUninit::new([0u8; 32])); + }) + }); + for chunk in built { + self.push_chunk(chunk); } - remaining = &remaining[full_chunks * CHUNK_SIZE..]; - } - - // tail remainder - if !remaining.is_empty() { - self.tail[..remaining.len()].copy_from_slice(remaining); - self.tail_len = remaining.len(); } self.len = new_len; Ok(()) } - pub(crate) fn ensure_len(&mut self, target: usize) -> Result<(), TreeError> { - if self.len >= target { - return Ok(()); - } - let needed = target - self.len; - - let tail_space = CHUNK_SIZE - self.tail_len; - let fill_tail = tail_space.min(needed); - debug_assert!( - self.tail[self.tail_len..self.tail_len + fill_tail] - .iter() - .all(|h| *h == [0u8; 32]), - "ensure_len: tail slots must be zeroed" - ); - self.tail_len += fill_tail; - let mut filled = fill_tail; - if self.tail_len == CHUNK_SIZE { - self.promote_tail(); - } + #[cfg(test)] + #[inline] + pub(crate) fn push(&mut self, value: Hash) -> Result<(), TreeError> { + self.extend(&[value]) + } - let remaining = needed - filled; - let full_chunks = remaining / CHUNK_SIZE; - if full_chunks > 0 { - for _ in 0..full_chunks { - self.push_chunk(Chunk::new_memory([[0u8; 32]; CHUNK_SIZE])); + #[inline] + pub(crate) fn extend(&mut self, values: &[Hash]) -> Result<(), TreeError> { + self.extend_with(values.len(), |off, out| { + // SAFETY: `extend_with` only ever asks for disjoint sub-ranges of + // the `count` it was given, so `off + out.len() <= values.len()`; + // `Hash` and `MaybeUninit` share a layout; and `out` is either + // a fresh allocation or `self.tail`, neither of which can overlap + // the caller's `values`. + unsafe { + core::ptr::copy_nonoverlapping( + values.as_ptr().add(off), + out.as_mut_ptr().cast::(), + out.len(), + ); } - filled += full_chunks * CHUNK_SIZE; - } - - let leftover = needed - filled; - self.tail_len += leftover; - - self.len = target; - Ok(()) + }) } - /// Promote the full tail into a chunk, freezing pending if full - fn promote_tail(&mut self) { - debug_assert_eq!(self.tail_len, CHUNK_SIZE); - self.push_chunk(Chunk::new_memory(self.tail)); - self.tail = [[0u8; 32]; CHUNK_SIZE]; - self.tail_len = 0; + /// Grow to `target`, zero-filling the new slots. + pub(crate) fn ensure_len(&mut self, target: usize) -> Result<(), TreeError> { + if self.len >= target { + return Ok(()); + } + self.extend_with(target - self.len, |_, out| { + out.fill(MaybeUninit::new([0u8; 32])); + }) } /// Push a chunk to pending, freezing into a segment when full @@ -421,34 +427,9 @@ impl ChunkedLevel { /// Collect chunks from index `already` onward #[cfg(feature = "storage")] pub(crate) fn chunks_since(&self, already: usize) -> Vec { - let total = self.chunk_count(); - if already >= total { - return Vec::new(); - } - let committed = self.segments.len() * CHUNKS_PER_SEGMENT; - let mut result = Vec::with_capacity(total - already); - - // Collect from segments - if already < committed { - let start_seg = already / CHUNKS_PER_SEGMENT; - let start_off = already % CHUNKS_PER_SEGMENT; - for (seg_i, segment) in self.segments.iter().enumerate().skip(start_seg) { - let from = if seg_i == start_seg { start_off } else { 0 }; - for chunk in &segment[from..] { - result.push(chunk.clone()); - } - } - } - - // Collect from pending - let pending_start = already.saturating_sub(committed); - if pending_start < self.pending.len() { - for chunk in &self.pending[pending_start..] { - result.push(chunk.clone()); - } - } - - result + (already..self.chunk_count()) + .map(|i| self.chunk(i).clone()) + .collect() } /// Remap the first `count` chunks to mmap-backed chunks (one region per shard) @@ -466,18 +447,7 @@ impl ChunkedLevel { return; } - let committed = self.segments.len() * CHUNKS_PER_SEGMENT; - let mut unmapped: Vec = - Vec::with_capacity(total.saturating_sub(remap_count)); - for chunk_idx in remap_count..total { - if chunk_idx < committed { - let seg_idx = chunk_idx / CHUNKS_PER_SEGMENT; - let seg_off = chunk_idx % CHUNKS_PER_SEGMENT; - unmapped.push(self.segments[seg_idx][seg_off].clone()); - } else { - unmapped.push(self.pending[chunk_idx - committed].clone()); - } - } + let unmapped: Vec = self.chunks_since(remap_count); self.segments.clear(); self.pending.clear(); @@ -491,24 +461,4 @@ impl ChunkedLevel { .for_each(|chunk| self.push_chunk(chunk)); } - /// Access the tail buffer - #[cfg(feature = "storage")] - pub(crate) fn tail_data(&self) -> &[Hash; CHUNK_SIZE] { - &self.tail - } - - #[cfg(test)] - pub(crate) fn tail_len(&self) -> usize { - self.tail_len - } - - #[cfg(test)] - pub(crate) fn get_chunk(&self, idx: usize) -> &Chunk { - let committed = self.segments.len() * CHUNKS_PER_SEGMENT; - if idx < committed { - &self.segments[idx / CHUNKS_PER_SEGMENT][idx % CHUNKS_PER_SEGMENT] - } else { - &self.pending[idx - committed] - } - } } diff --git a/crates/rotortree/src/storage/checkpoint.rs b/crates/rotortree/src/storage/checkpoint.rs index 67521d4..52a3800 100644 --- a/crates/rotortree/src/storage/checkpoint.rs +++ b/crates/rotortree/src/storage/checkpoint.rs @@ -197,51 +197,6 @@ pub(crate) fn read_meta(data_dir: &Path) -> Result, Stora })) } -/// Write all level tails atomically (tmp -> fsync -> rename). -pub(crate) fn write_tails( - data_dir: &Path, - tails: &[[Hash; CHUNK_SIZE]], - max_depth: usize, -) -> io::Result<()> { - let total_size = max_depth * CHUNK_BYTE_SIZE; - let mut buf = vec![0u8; total_size]; - - for (i, tail) in tails.iter().enumerate() { - let base = i * CHUNK_BYTE_SIZE; - buf[base..base + CHUNK_BYTE_SIZE].copy_from_slice(tail.as_flattened()); - } - - atomic_write(&data_dir.join("tails.bin"), &buf) -} - -/// Read all tails from disk. Returns `None` if the file is missing or wrong size -pub(crate) fn read_tails( - data_dir: &Path, - max_depth: usize, -) -> io::Result>> { - let path = data_dir.join("tails.bin"); - let data = match fs::read(&path) { - Ok(d) => d, - Err(e) if e.kind() == io::ErrorKind::NotFound => return Ok(None), - Err(e) => return Err(e), - }; - let expected = max_depth * CHUNK_BYTE_SIZE; - if data.len() != expected { - return Ok(None); - } - - let mut tails = Vec::with_capacity(max_depth); - for i in 0..max_depth { - let base = i * CHUNK_BYTE_SIZE; - let chunk = &data[base..base + CHUNK_BYTE_SIZE]; - let mut tail = [[0u8; 32]; CHUNK_SIZE]; - tail.as_flattened_mut().copy_from_slice(chunk); - tails.push(tail); - } - - Ok(Some(tails)) -} - pub(crate) fn level_dir_path(data_dir: &Path, level_idx: usize) -> PathBuf { data_dir.join(format!("level_{level_idx}")) } diff --git a/crates/rotortree/src/storage/mod.rs b/crates/rotortree/src/storage/mod.rs index 1481b9a..3589d92 100644 --- a/crates/rotortree/src/storage/mod.rs +++ b/crates/rotortree/src/storage/mod.rs @@ -65,8 +65,8 @@ struct LevelCheckpointData { new_chunks: Vec, /// start from from_chunk: usize, + /// Complete chunks only: a partial last chunk is rewritten each checkpoint. total_chunks: usize, - tail: [Hash; CHUNK_SIZE], } /// checkpoint snapshot @@ -222,7 +222,10 @@ impl Shared let mut level_data = Vec::with_capacity(active_levels); for level_idx in 0..active_levels { - let total_chunks = state.inner.levels[level_idx].chunk_count(); + // Count only whole chunks as done. The last chunk may be + // partially filled and still growing + let total_chunks = + state.inner.levels[level_idx].len() / CHUNK_SIZE; let already = if level_idx < state.checkpointed_chunks.len() { state.checkpointed_chunks[level_idx] } else { @@ -232,13 +235,10 @@ impl Shared let new_chunks: Vec = state.inner.levels[level_idx].chunks_since(already); - let tail = *state.inner.levels[level_idx].tail_data(); - level_data.push(LevelCheckpointData { new_chunks, from_chunk: already, total_chunks, - tail, }); } @@ -280,10 +280,6 @@ impl Shared file.sync_data()?; } - let tails: Vec<[Hash; CHUNK_SIZE]> = - snap.level_data.iter().map(|ld| ld.tail).collect(); - checkpoint::write_tails(&self.data_dir, &tails, MAX_DEPTH)?; - #[allow(clippy::cast_possible_truncation)] checkpoint::write_meta( &self.data_dir, diff --git a/crates/rotortree/src/storage/recovery.rs b/crates/rotortree/src/storage/recovery.rs index 7110517..9907f75 100644 --- a/crates/rotortree/src/storage/recovery.rs +++ b/crates/rotortree/src/storage/recovery.rs @@ -168,21 +168,16 @@ where } } - let tails = match checkpoint::read_tails(data_dir, MAX_DEPTH)? { - Some(t) => t, - None => return recover(wal_file, hasher), - }; - let mut inner = TreeInner::::new(); - for level_idx in 0..=depth.min(MAX_DEPTH - 1) { - let len = level_lens[level_idx]; + let active = depth.min(MAX_DEPTH - 1) + 1; + for (level_idx, &len) in level_lens.iter().enumerate().take(active) { if len == 0 { continue; } - let num_chunks = len / CHUNK_SIZE; - let tail_len = len % CHUNK_SIZE; + // Every chunk on disk, including a zero-padded partial last one. + let num_chunks = len.div_ceil(CHUNK_SIZE); let regions = if num_chunks > 0 { checkpoint::mmap_level_shards(data_dir, level_idx, num_chunks)? @@ -202,7 +197,7 @@ where .collect() }; - inner.set_level_from_parts(level_idx, chunks, tails[level_idx], tail_len, len); + inner.set_level_from_parts(level_idx, chunks, len); } inner.root = if leaf_count > 0 { diff --git a/crates/rotortree/src/tree.rs b/crates/rotortree/src/tree.rs index 507c543..48940ca 100644 --- a/crates/rotortree/src/tree.rs +++ b/crates/rotortree/src/tree.rs @@ -2,11 +2,17 @@ use crate::{ Hash, Hasher, TreeError, - chunked_level::ChunkedLevel, + chunked_level::{ + CHUNK_SIZE, + ChunkedLevel, + }, }; -/// Number of parents per rayon task -#[cfg(feature = "parallel")] +/// Parents gathered per `hash_many_into` call. +const PARENT_WINDOW: usize = 16; + +/// Number of parents per rayon task in [`TreeInner::recompute_root`]. +#[cfg(all(feature = "parallel", feature = "storage"))] const PAR_CHUNK_SIZE: usize = 64; #[cfg(feature = "parallel")] @@ -117,11 +123,9 @@ impl TreeInner { &mut self, level_idx: usize, chunks: Vec, - tail: [Hash; crate::chunked_level::CHUNK_SIZE], - tail_len: usize, len: usize, ) { - self.levels[level_idx] = ChunkedLevel::from_parts(chunks, tail, tail_len, len); + self.levels[level_idx] = ChunkedLevel::from_parts(chunks, len); } /// Recompute the root hash from level 0 data bottom-up @@ -431,22 +435,22 @@ impl LeanIMT let next_level = level + 1; if next_level < levels.len() { + levels[next_level].ensure_len(num_parents)?; // Split so the child level is borrowed immutably while the parent // level is written mutably. let (head, tail) = levels.split_at_mut(next_level); let child = &head[level]; let parent = &mut tail[0]; - // Leading full groups (count == N) batch through hash_many_into. - // The final group is full iff `level_len` is a multiple of N. - let full_parents = (level_len / N).max(start_parent); - Self::_batch_full_groups(child, parent, start_parent, full_parents, hasher); - - // Trailing partial group / lift (empty range when level_len % N == 0). - for parent_idx in full_parents..num_parents { - let p = Self::_compute_parent(child, parent_idx, level_len, hasher)?; - parent.set_preallocated(parent_idx, p); - } + parent.write_with(start_parent, num_parents - start_parent, |off, out| { + Self::_compute_parents_batched( + child, + start_parent + off, + level_len, + hasher, + out, + ); + }); // At the root level there is exactly one parent; it is the root. if is_root_level { @@ -465,59 +469,57 @@ impl LeanIMT Ok(()) } - /// Hash the full groups `start_parent..full_parents` of `child` into - /// `parent`, batching eligible runs through [`Hasher::hash_many_into`]. + /// Fill `out` with parents `start_parent .. start_parent + out.len()`, + /// batching full groups through [`Hasher::hash_many_into`] and falling back + /// to the scalar [`Self::_compute_parent`] for the trailing partial group + /// and the lift. #[inline] - fn _batch_full_groups( + fn _compute_parents_batched( child: &ChunkedLevel, - parent: &mut ChunkedLevel, start_parent: usize, - full_parents: usize, + level_len: usize, hasher: &H, + out: &mut [core::mem::MaybeUninit], ) { - /// Parents gathered per `hash_many_into` call. The Blake3 override - /// re-splits this into `simd_degree`-sized SIMD calls internally. - const WINDOW: usize = 16; - - let mut parent_idx = start_parent; - let mut refs: [&[Hash]; WINDOW] = [&[]; WINDOW]; - let mut out: [Hash; WINDOW] = [[0u8; 32]; WINDOW]; - while parent_idx < full_parents { - let take = (full_parents - parent_idx).min(WINDOW); - let mut filled = 0; - for slot in refs.iter_mut().take(take) { - let start = (parent_idx + filled) * N; - match child.group_slice(start, N) { - Some(g) => *slot = g, - // Group straddles a boundary (cannot happen for eligible N - // since CHUNK_SIZE % N == 0, but stay correct anyway): - // stop the window here and let the scalar tail finish it. - None => break, + let full = if CHUNK_SIZE.is_multiple_of(N) { + (level_len / N).saturating_sub(start_parent).min(out.len()) + } else { + 0 + }; + + let mut refs: [&[Hash]; PARENT_WINDOW] = [&[]; PARENT_WINDOW]; + let mut staging: [Hash; PARENT_WINDOW] = [[0u8; 32]; PARENT_WINDOW]; + let mut i = 0; + let mut win = 0; + for run in child.runs(start_parent * N, full * N) { + for group in run.chunks_exact(N) { + refs[win] = group; + win += 1; + if win == PARENT_WINDOW { + hasher.hash_many_into(&refs, &mut staging); + for (k, &h) in staging.iter().enumerate() { + out[i + k].write(h); + } + i += win; + win = 0; } - filled += 1; } - if filled == 0 { - // Could not borrow even one group contiguously; scalar. - let p = hasher.hash_children(&Self::_copy_group(child, parent_idx)); - parent.set_preallocated(parent_idx, p); - parent_idx += 1; - continue; - } - hasher.hash_many_into(&refs[..filled], &mut out[..filled]); - for (i, &h) in out[..filled].iter().enumerate() { - parent.set_preallocated(parent_idx + i, h); + } + if win > 0 { + hasher.hash_many_into(&refs[..win], &mut staging[..win]); + for (k, &h) in staging[..win].iter().enumerate() { + out[i + k].write(h); } - parent_idx += filled; + i += win; } - } - /// Copy a full N-child group into a stack buffer (boundary-straddling - /// fallback for `_batch_full_groups`). - #[inline] - fn _copy_group(child: &ChunkedLevel, parent_idx: usize) -> [Hash; N] { - let mut buf = [[0u8; 32]; N]; - child.get_group(parent_idx * N, N, &mut buf); - buf + // remainder + while i < out.len() { + let p = Self::_compute_parent(child, start_parent + i, level_len, hasher) + .expect("ensure_len guarantees valid indices"); + out[i].write(p); + i += 1; + } } pub(crate) fn _insert_many( @@ -543,18 +545,6 @@ impl LeanIMT inner.levels[0].extend(leaves)?; - // allocate upfront - { - let mut level_len = inner.levels[0].len(); - for level in 0..depth { - let num_parents = level_len.div_ceil(N); - if level + 1 < MAX_DEPTH { - inner.levels[level + 1].ensure_len(num_parents)?; - } - level_len = num_parents; - } - } - let old_size_usize = u64_to_usize(inner.size)?; let mut start_parent = old_size_usize / N; @@ -564,8 +554,6 @@ impl LeanIMT [0u8; 32] }; - #[cfg(feature = "parallel")] - let mut par_buf: std::vec::Vec = std::vec::Vec::new(); #[cfg(feature = "parallel")] let par_threshold = parallel_threshold(); @@ -577,47 +565,39 @@ impl LeanIMT #[cfg(feature = "parallel")] { let work = num_parents - start_parent; - if work >= par_threshold { - use rayon::prelude::*; - + if work >= par_threshold && level + 1 < MAX_DEPTH { let split_at = level + 1; let (child_levels, parent_levels) = inner.levels.split_at_mut(split_at); let child_level = &child_levels[level]; + let parent = &mut parent_levels[0]; + + let overlap = usize::from(start_parent < parent.len()); + if overlap == 1 { + parent.write_with(start_parent, 1, |_, out| { + Self::_compute_parents_batched( + child_level, + start_parent, + level_len, + hasher, + out, + ); + }); + } - par_buf.clear(); - par_buf.reserve(work); - - let spare = &mut par_buf.spare_capacity_mut()[..work]; - spare.par_chunks_mut(PAR_CHUNK_SIZE).enumerate().for_each( - |(ci, chunk)| { - let base = start_parent + ci * PAR_CHUNK_SIZE; - for (i, slot) in chunk.iter_mut().enumerate() { - slot.write( - Self::_compute_parent( - child_level, - base + i, - level_len, - hasher, - ) - .expect("ensure_len guarantees valid indices"), - ); - } - }, - ); - // SAFETY: the parallel loop above initialised every - // element in `spare[..work]` via `MaybeUninit::write`. - unsafe { par_buf.set_len(work) }; - - let parent_level = &mut parent_levels[0]; - for (i, &parent) in par_buf.iter().enumerate() { - let parent_idx = start_parent + i; - if split_at < MAX_DEPTH { - parent_level.set_preallocated(parent_idx, parent); - } - if is_root_level { - root = parent; - } + let appended = start_parent + overlap; + parent.extend_with(work - overlap, |off, out| { + Self::_compute_parents_batched( + child_level, + appended + off, + level_len, + hasher, + out, + ); + })?; + + if is_root_level { + root = parent.get(num_parents - 1)?; } } else { Self::_insert_many_level_seq( @@ -727,20 +707,23 @@ mod tests { } #[test] - fn chunked_level_promotes_at_chunk_size() { + fn chunked_level_pads_partial_chunk_with_zeros() { let mut level = ChunkedLevel::new(); for i in 0..CHUNK_SIZE { level.push(leaf(i as u32)).unwrap(); } assert_eq!(level.chunk_count(), 1); - assert_eq!(level.tail_len(), 0); assert_eq!(level.len(), CHUNK_SIZE); - // One more goes into the new tail. + // One more opens a second, partially filled chunk. level.push(leaf(0xFF)).unwrap(); - assert_eq!(level.chunk_count(), 1); - assert_eq!(level.tail_len(), 1); + assert_eq!(level.chunk_count(), 2); assert_eq!(level.len(), CHUNK_SIZE + 1); + assert!( + level.chunk(1).as_slice()[1..] + .iter() + .all(|h| *h == [0u8; 32]) + ); } #[test] @@ -752,7 +735,7 @@ mod tests { let snap = level.clone(); assert_eq!(snap.len(), level.len()); // The completed chunk Arc is shared. - assert!(Chunk::ptr_eq(level.get_chunk(0), snap.get_chunk(0))); + assert!(Chunk::ptr_eq(level.chunk(0), snap.chunk(0))); // Data matches. for i in 0..level.len() { assert_eq!(level.get(i).unwrap(), snap.get(i).unwrap()); diff --git a/crates/rotortree/tests/storage_integration.rs b/crates/rotortree/tests/storage_integration.rs index 2c172ba..7b31236 100644 --- a/crates/rotortree/tests/storage_integration.rs +++ b/crates/rotortree/tests/storage_integration.rs @@ -454,7 +454,6 @@ fn checkpoint_round_trip() { assert!(dir.path().join("data").join("header.bin").exists()); assert!(dir.path().join("data").join("checkpoint.meta").exists()); - assert!(dir.path().join("data").join("tails.bin").exists()); tree.close().unwrap(); From 6dbf3932a2782c719f6847405374103d723f39d2 Mon Sep 17 00:00:00 2001 From: rymnc <43716372+rymnc@users.noreply.github.com> Date: Fri, 14 Aug 2026 03:23:35 +0530 Subject: [PATCH 2/2] fix: fmt --- crates/rotortree/src/chunked_level.rs | 16 +++++++++++----- crates/rotortree/src/storage/mod.rs | 3 +-- 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/crates/rotortree/src/chunked_level.rs b/crates/rotortree/src/chunked_level.rs index 8a8f010..32ba04f 100644 --- a/crates/rotortree/src/chunked_level.rs +++ b/crates/rotortree/src/chunked_level.rs @@ -217,8 +217,7 @@ impl ChunkedLevel { pub(crate) fn chunk(&self, chunk_idx: usize) -> &Chunk { let committed = self.segments.len() * CHUNKS_PER_SEGMENT; if chunk_idx < committed { - &self.segments[chunk_idx / CHUNKS_PER_SEGMENT] - [chunk_idx % CHUNKS_PER_SEGMENT] + &self.segments[chunk_idx / CHUNKS_PER_SEGMENT][chunk_idx % CHUNKS_PER_SEGMENT] } else { &self.pending[chunk_idx - committed] } @@ -253,7 +252,11 @@ impl ChunkedLevel { /// Borrow `[start, start + count)` as contiguous runs, one per chunk. #[inline] - pub(crate) fn runs(&self, start: usize, count: usize) -> impl Iterator { + pub(crate) fn runs( + &self, + start: usize, + count: usize, + ) -> impl Iterator { debug_assert!( start + count <= self.len, "runs: {start}+{count} > len {}", @@ -331,7 +334,11 @@ impl ChunkedLevel { /// Append `count` hashes produced by `fill`, writing each chunk straight /// into its final allocation. - pub(crate) fn extend_with(&mut self, count: usize, fill: F) -> Result<(), TreeError> + pub(crate) fn extend_with( + &mut self, + count: usize, + fill: F, + ) -> Result<(), TreeError> where F: Fn(usize, &mut [MaybeUninit]) + Sync, { @@ -460,5 +467,4 @@ impl ChunkedLevel { .chain(unmapped) .for_each(|chunk| self.push_chunk(chunk)); } - } diff --git a/crates/rotortree/src/storage/mod.rs b/crates/rotortree/src/storage/mod.rs index 3589d92..5299d80 100644 --- a/crates/rotortree/src/storage/mod.rs +++ b/crates/rotortree/src/storage/mod.rs @@ -224,8 +224,7 @@ impl Shared for level_idx in 0..active_levels { // Count only whole chunks as done. The last chunk may be // partially filled and still growing - let total_chunks = - state.inner.levels[level_idx].len() / CHUNK_SIZE; + let total_chunks = state.inner.levels[level_idx].len() / CHUNK_SIZE; let already = if level_idx < state.checkpointed_chunks.len() { state.checkpointed_chunks[level_idx] } else {