Skip to main content

prolly/prolly/
write.rs

1//! Deterministic mutation stream and resynchronizing tree writer.
2
3use std::collections::HashSet;
4use std::ops::Range;
5
6use rayon::prelude::*;
7
8use super::boundary::entry_count_boundary;
9use super::builder::{BatchBuilder, LevelEmitter, NodeSummary};
10use super::cid::Cid;
11use super::config::Config;
12use super::error::{Error, Mutation};
13use super::format::{BoundaryInput, ChunkMeasure, NodeLayoutSpec};
14use super::node::Node;
15use super::parallel::{ExecutionPolicy, ParallelConfig};
16use super::store::Store;
17use super::{Prolly, Tree};
18
19const LOCAL_WRITE_CACHE_LIMIT: usize = 8;
20
21/// Store-neutral work performed by a tree write.
22#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
23pub struct WriteStats {
24    pub input_mutations: u64,
25    pub effective_mutations: u64,
26    pub entries_streamed: u64,
27    pub nodes_read: u64,
28    pub nodes_written: u64,
29    pub nodes_reused: u64,
30    pub bytes_read: u64,
31    pub bytes_written: u64,
32    pub resync_distance_entries: u64,
33    pub resync_distance_nodes: u64,
34    pub used_key_stable_fast_path: bool,
35    pub used_batched_value_update_path: bool,
36    pub parallel_width: u64,
37    pub parallel_tasks: u64,
38    pub structural_islands: u64,
39    pub coalesced_islands: u64,
40}
41
42pub(crate) struct EmittedLeaf {
43    pub(crate) summary: NodeSummary,
44    pub(crate) bytes: Vec<u8>,
45    pub(crate) node: Node,
46}
47
48struct EmittedInternal {
49    cid: Cid,
50    bytes: Vec<u8>,
51    node: Node,
52}
53
54const MUTATION_ISLAND_GUARD_LEAVES: usize = 8;
55
56#[derive(Clone, Debug, PartialEq, Eq)]
57struct MutationIsland {
58    /// Existing leaf range owned by this replay. The end advances to the
59    /// matched anchor after a successful replay.
60    leaf_range: Range<usize>,
61    /// Normalized mutations owned exclusively by this replay.
62    mutation_range: Range<usize>,
63    /// Exclusive old-leaf guard that replay must not cross.
64    protected_end: usize,
65}
66
67struct IslandReplay {
68    island: MutationIsland,
69    summaries: Vec<NodeSummary>,
70    emitted: Vec<EmittedLeaf>,
71    resynced_at: Option<usize>,
72    entries_streamed: u64,
73    nodes_read: u64,
74    bytes_read: u64,
75}
76
77impl IslandReplay {
78    fn proved_independent(&self) -> bool {
79        self.resynced_at
80            .map(|index| index < self.island.protected_end)
81            .unwrap_or(false)
82    }
83}
84
85#[derive(Clone, Copy)]
86struct MutationIslandCandidate {
87    first_leaf: usize,
88    last_leaf: usize,
89    first_mutation: usize,
90    mutation_end: usize,
91}
92
93struct StructuralIslandReplay {
94    summaries: Vec<NodeSummary>,
95    emitted: Vec<EmittedLeaf>,
96}
97
98pub(crate) struct LeafEmitter {
99    emitter: LevelEmitter,
100    pub(crate) emitted: Vec<EmittedLeaf>,
101}
102
103impl LeafEmitter {
104    pub(crate) fn new(config: &super::config::Config) -> Result<Self, Error> {
105        Ok(Self {
106            emitter: LevelEmitter::new(config.clone(), true, 0)?,
107            emitted: Vec::new(),
108        })
109    }
110
111    pub(crate) fn push(&mut self, key: Vec<u8>, value: Vec<u8>) -> Result<(), Error> {
112        let output = &mut self.emitted;
113        self.emitter.push_leaf_with(key, value, |emitted| {
114            output.push(EmittedLeaf {
115                summary: emitted.summary,
116                bytes: emitted.bytes,
117                node: emitted.node,
118            });
119        })
120    }
121
122    pub(crate) fn flush(&mut self) -> Result<(), Error> {
123        if let Some(emitted) = self.emitter.finish()? {
124            self.emitted.push(EmittedLeaf {
125                summary: emitted.summary,
126                bytes: emitted.bytes,
127                node: emitted.node,
128            });
129        }
130        Ok(())
131    }
132
133    pub(crate) fn is_aligned_with(&self, old: &NodeSummary) -> bool {
134        self.emitter.is_empty()
135            && self
136                .emitted
137                .last()
138                .map(|leaf| leaf.summary.cid == old.cid)
139                .unwrap_or(false)
140    }
141}
142
143/// Apply last-write-wins mutations and emit the unique deterministic tree.
144pub(crate) fn apply<S: Store>(
145    manager: &Prolly<S>,
146    tree: &Tree,
147    mutations: Vec<Mutation>,
148) -> Result<(Tree, WriteStats), Error> {
149    apply_impl(manager, tree, mutations, true, None)
150}
151
152pub(crate) fn apply_configured<S: Store>(
153    manager: &Prolly<S>,
154    tree: &Tree,
155    mutations: Vec<Mutation>,
156    config: &ParallelConfig,
157) -> Result<(Tree, WriteStats), Error> {
158    apply_impl(manager, tree, mutations, true, Some(config))
159}
160
161pub(crate) fn apply_tree<S: Store>(
162    manager: &Prolly<S>,
163    tree: &Tree,
164    mutations: Vec<Mutation>,
165) -> Result<Tree, Error> {
166    Ok(apply_impl(manager, tree, mutations, false, None)?.0)
167}
168
169pub(crate) fn apply_tree_configured<S: Store>(
170    manager: &Prolly<S>,
171    tree: &Tree,
172    mutations: Vec<Mutation>,
173    config: &ParallelConfig,
174) -> Result<Tree, Error> {
175    Ok(apply_impl(manager, tree, mutations, false, Some(config))?.0)
176}
177
178fn apply_impl<S: Store>(
179    manager: &Prolly<S>,
180    tree: &Tree,
181    mutations: Vec<Mutation>,
182    measure_read_bytes: bool,
183    parallel_config: Option<&ParallelConfig>,
184) -> Result<(Tree, WriteStats), Error> {
185    let parallel_threshold = parallel_config
186        .map(|config| config.parallelism_threshold)
187        .unwrap_or_else(|| ParallelConfig::default().parallelism_threshold);
188    let mut stats = WriteStats {
189        input_mutations: mutations.len() as u64,
190        ..WriteStats::default()
191    };
192    if mutations.is_empty() {
193        return Ok((tree.clone(), stats));
194    }
195    if let Some(root) = &tree.root {
196        let node = manager.load_arc(root)?;
197        if node.format != tree.config.format {
198            return Err(Error::FormatMismatch {
199                expected: tree.config.format.digest()?,
200                actual: node.format.digest()?,
201            });
202        }
203    }
204
205    let mut mutations = normalize(mutations);
206    stats.effective_mutations = mutations.len() as u64;
207    let _concurrency_guard = (mutations.len() >= parallel_threshold
208        && rayon::current_num_threads() > 1)
209        .then(super::parallel::CanonicalWriteConcurrencyGuard::enter);
210    let policy = parallel_config.map_or_else(
211        || ExecutionPolicy::automatic(mutations.len(), mutations.len()),
212        |config| ExecutionPolicy::from_config(config, mutations.len(), mutations.len()),
213    );
214    // Report the maximum width actually used, not merely the width admitted
215    // by policy. Canonical fallbacks that execute no independent work remain
216    // honest width-one operations in public telemetry.
217    stats.parallel_width = 1;
218    if tree.root.is_none() {
219        return build_empty_base(manager, tree, mutations, stats);
220    }
221    if let Some(result) = try_append(
222        manager,
223        tree,
224        &mut mutations,
225        &mut stats,
226        measure_read_bytes,
227    )? {
228        return Ok(result);
229    }
230    mutations = match try_direct_value_updates(
231        manager,
232        tree,
233        mutations,
234        &mut stats,
235        measure_read_bytes,
236        policy,
237    )? {
238        DirectValueUpdateAttempt::Applied(result) => return Ok(*result),
239        DirectValueUpdateAttempt::Fallback(mutations) => mutations,
240    };
241    if let Some(result) =
242        try_direct_single_delete(manager, tree, &mutations, &mut stats, measure_read_bytes)?
243    {
244        return Ok(result);
245    }
246    if let Some(result) =
247        try_localized_height_two_deletes(manager, tree, &mutations, &mut stats, measure_read_bytes)?
248    {
249        return Ok(result);
250    }
251    let (old_leaves, old_internal_cids) =
252        collect_leaf_summaries(manager, tree, &mut stats, measure_read_bytes)?;
253    if old_leaves.is_empty() {
254        return build_empty_base(manager, tree, mutations, stats);
255    }
256    let parallel_replay = try_parallel_structural_islands(
257        manager,
258        tree,
259        &old_leaves,
260        &mutations,
261        policy,
262        &mut stats,
263        measure_read_bytes,
264    )?;
265    let (summaries, emitted) = if let Some(replay) = parallel_replay {
266        (replay.summaries, replay.emitted)
267    } else {
268        let mut mutation_index = 0;
269        let mut old_cursor = 0usize;
270        let mut summaries = Vec::with_capacity(old_leaves.len());
271        let mut emitted = Vec::<EmittedLeaf>::new();
272
273        while mutation_index < mutations.len() {
274            let start = old_leaves
275                .partition_point(|leaf| {
276                    leaf.first_key.as_slice() <= mutations[mutation_index].0.as_slice()
277                })
278                // A hard-cap split occurs immediately before the entry that did
279                // not fit. If that entry shrinks, the canonical split can move
280                // into the preceding leaf, so replay one predecessor as context.
281                .saturating_sub(2)
282                .max(old_cursor);
283            summaries.extend_from_slice(&old_leaves[old_cursor..start]);
284            stats.nodes_reused += start.saturating_sub(old_cursor) as u64;
285
286            let mut emitter = LeafEmitter::new(&tree.config)?;
287            let mut resynced_at = None;
288            let first_pending_mutation = mutation_index;
289            for leaf_index in start..old_leaves.len() {
290                let leaf = manager.load_arc(&old_leaves[leaf_index].cid)?;
291                stats.nodes_read += 1;
292                if measure_read_bytes {
293                    stats.bytes_read += leaf.encoded_len() as u64;
294                }
295                if !leaf.leaf || leaf.keys.len() != leaf.vals.len() {
296                    return Err(Error::InvalidNode);
297                }
298
299                for (key, value) in leaf.keys.iter().cloned().zip(leaf.vals.iter().cloned()) {
300                    while mutation_index < mutations.len() && mutations[mutation_index].0 < key {
301                        let (mutation_key, mutation_value) =
302                            take_mutation(&mut mutations[mutation_index]);
303                        if let Some(value) = mutation_value {
304                            emitter.push(mutation_key, value)?;
305                            stats.entries_streamed += 1;
306                        }
307                        mutation_index += 1;
308                    }
309                    if mutation_index < mutations.len() && mutations[mutation_index].0 == key {
310                        let (_, mutation_value) = take_mutation(&mut mutations[mutation_index]);
311                        if let Some(value) = mutation_value {
312                            emitter.push(key, value)?;
313                            stats.entries_streamed += 1;
314                        }
315                        mutation_index += 1;
316                    } else {
317                        emitter.push(key, value)?;
318                        stats.entries_streamed += 1;
319                    }
320                }
321
322                let next_first = old_leaves.get(leaf_index + 1).map(|leaf| &leaf.first_key);
323                while mutation_index < mutations.len()
324                    && next_first
325                        .map(|next| mutations[mutation_index].0 < *next)
326                        .unwrap_or(true)
327                {
328                    let (mutation_key, mutation_value) =
329                        take_mutation(&mut mutations[mutation_index]);
330                    if let Some(value) = mutation_value {
331                        emitter.push(mutation_key, value)?;
332                        stats.entries_streamed += 1;
333                    }
334                    mutation_index += 1;
335                }
336
337                stats.resync_distance_nodes += 1;
338                if mutation_index > first_pending_mutation
339                    && emitter.is_aligned_with(&old_leaves[leaf_index])
340                {
341                    resynced_at = Some(leaf_index);
342                    break;
343                }
344            }
345
346            while mutation_index < mutations.len() && resynced_at.is_none() {
347                let (mutation_key, mutation_value) = take_mutation(&mut mutations[mutation_index]);
348                if let Some(value) = mutation_value {
349                    emitter.push(mutation_key, value)?;
350                    stats.entries_streamed += 1;
351                }
352                mutation_index += 1;
353            }
354            emitter.flush()?;
355            summaries.extend(emitter.emitted.iter().map(|leaf| leaf.summary.clone()));
356            emitted.extend(emitter.emitted);
357            old_cursor = resynced_at.map_or(old_leaves.len(), |index| index + 1);
358            if resynced_at.is_none() {
359                break;
360            }
361        }
362        summaries.extend_from_slice(&old_leaves[old_cursor..]);
363        stats.nodes_reused += old_leaves.len().saturating_sub(old_cursor) as u64;
364        (summaries, emitted)
365    };
366    stats.resync_distance_entries = stats.entries_streamed;
367
368    if summaries
369        .iter()
370        .map(|leaf| &leaf.cid)
371        .eq(old_leaves.iter().map(|leaf| &leaf.cid))
372    {
373        return Ok((tree.clone(), stats));
374    }
375
376    let old_cids = old_leaves
377        .iter()
378        .map(|leaf| leaf.cid.clone())
379        .collect::<HashSet<_>>();
380    let mut changed_cids = HashSet::new();
381    let changed_leaves = emitted
382        .iter()
383        .filter(|leaf| {
384            !old_cids.contains(&leaf.summary.cid) && changed_cids.insert(leaf.summary.cid.clone())
385        })
386        .collect::<Vec<_>>();
387    let fixed_separators = tree.config.format.chunking.measure == ChunkMeasure::EntryCount
388        && tree.config.format.chunking.input == BoundaryInput::Key
389        && !matches!(
390            tree.config.format.node_layout,
391            NodeLayoutSpec::Custom { .. }
392        )
393        && summaries.len() == old_leaves.len()
394        && summaries
395            .iter()
396            .zip(&old_leaves)
397            .all(|(new, old)| new.first_key == old.first_key);
398    if fixed_separators {
399        let changes = changed_leaves
400            .iter()
401            .map(|leaf| leaf.summary.clone())
402            .collect::<Vec<_>>();
403        let (written, internal_nodes) = rewrite_fixed_separator_paths(manager, tree, &changes)?;
404        let writes = changed_leaves
405            .iter()
406            .map(|leaf| (leaf.summary.cid.as_bytes(), leaf.bytes.as_slice()))
407            .chain(
408                internal_nodes
409                    .iter()
410                    .filter(|node| !old_internal_cids.contains(&node.cid))
411                    .map(|node| (node.cid.as_bytes(), node.bytes.as_slice())),
412            )
413            .collect::<Vec<_>>();
414        manager
415            .store()
416            .batch_put(&writes)
417            .map_err(|error| Error::Store(Box::new(error)))?;
418        let bytes_written = writes.iter().map(|(_, bytes)| bytes.len()).sum::<usize>();
419        manager.record_batch_write_metrics(writes.len(), bytes_written);
420        stats.nodes_written += writes.len() as u64;
421        stats.bytes_written += bytes_written as u64;
422        if writes.len() <= LOCAL_WRITE_CACHE_LIMIT {
423            for leaf in &changed_leaves {
424                manager.cache_node(leaf.summary.cid.clone(), leaf.node.clone());
425            }
426            for node in internal_nodes
427                .iter()
428                .filter(|node| !old_internal_cids.contains(&node.cid))
429            {
430                manager.cache_node(node.cid.clone(), node.node.clone());
431            }
432        }
433        return Ok((written, stats));
434    }
435    let builder = BatchBuilder::new(manager.store(), tree.config.clone());
436    let (written, internal_nodes) = builder.build_from_chunks_serial_deferred(summaries)?;
437    let writes = changed_leaves
438        .iter()
439        .map(|leaf| (leaf.summary.cid.as_bytes(), leaf.bytes.as_slice()))
440        .chain(
441            internal_nodes
442                .iter()
443                .filter(|node| !old_internal_cids.contains(&node.cid))
444                .map(|node| (node.cid.as_bytes(), node.bytes.as_slice())),
445        )
446        .collect::<Vec<_>>();
447    if !writes.is_empty() {
448        manager
449            .store()
450            .batch_put(&writes)
451            .map_err(|error| Error::Store(Box::new(error)))?;
452        let bytes_written = writes.iter().map(|(_, bytes)| bytes.len()).sum::<usize>();
453        stats.nodes_written += writes.len() as u64;
454        stats.bytes_written += bytes_written as u64;
455        manager.record_batch_write_metrics(writes.len(), bytes_written);
456        if writes.len() <= LOCAL_WRITE_CACHE_LIMIT {
457            for leaf in &changed_leaves {
458                manager.cache_node(leaf.summary.cid.clone(), leaf.node.clone());
459            }
460            for node in internal_nodes
461                .iter()
462                .filter(|node| !old_internal_cids.contains(&node.cid))
463            {
464                manager.cache_node(node.cid.clone(), node.node.clone());
465            }
466        }
467    }
468    if let Some(root) = &written.root {
469        let _ = manager.load_arc(root)?;
470    }
471    Ok((written, stats))
472}
473
474/// Plan conservative, non-overlapping mutation regions at canonical leaf
475/// boundaries. Close candidates are coalesced before replay so each remaining
476/// candidate has an untouched, guarded old-leaf interval before its neighbor.
477fn plan_mutation_islands(
478    old_leaves: &[NodeSummary],
479    mutations: &[(Vec<u8>, Option<Vec<u8>>)],
480) -> Vec<MutationIsland> {
481    if old_leaves.is_empty() || mutations.is_empty() {
482        return Vec::new();
483    }
484
485    let candidates = mutation_island_candidates(old_leaves, mutations);
486    let mut islands = Vec::<MutationIsland>::with_capacity(candidates.len());
487    for candidate in candidates {
488        let leaf_end = candidate.last_leaf.saturating_add(1).min(old_leaves.len());
489        let island = MutationIsland {
490            // Exact byte caps can move a split into either predecessor, so
491            // retain the same two-leaf context as the canonical writer.
492            leaf_range: candidate.first_leaf.saturating_sub(2)..leaf_end,
493            mutation_range: candidate.first_mutation..candidate.mutation_end,
494            protected_end: leaf_end
495                .saturating_add(MUTATION_ISLAND_GUARD_LEAVES)
496                .min(old_leaves.len()),
497        };
498
499        if let Some(previous) = islands.last_mut() {
500            if previous.protected_end > island.leaf_range.start {
501                previous.leaf_range.end = island.leaf_range.end;
502                previous.mutation_range.end = island.mutation_range.end;
503                previous.protected_end = previous.protected_end.max(island.protected_end);
504                continue;
505            }
506        }
507        islands.push(island);
508    }
509    islands
510}
511
512fn mutation_island_candidates(
513    old_leaves: &[NodeSummary],
514    mutations: &[(Vec<u8>, Option<Vec<u8>>)],
515) -> Vec<MutationIslandCandidate> {
516    let target_leaf = |key: &[u8]| {
517        old_leaves
518            .partition_point(|leaf| leaf.first_key.as_slice() <= key)
519            .saturating_sub(1)
520    };
521    let mut candidates = Vec::<MutationIslandCandidate>::new();
522    for (mutation_index, (key, _)) in mutations.iter().enumerate() {
523        let leaf_index = target_leaf(key);
524        match candidates.last_mut() {
525            Some(candidate) if leaf_index <= candidate.last_leaf.saturating_add(1) => {
526                candidate.last_leaf = candidate.last_leaf.max(leaf_index);
527                candidate.mutation_end = mutation_index + 1;
528            }
529            _ => candidates.push(MutationIslandCandidate {
530                first_leaf: leaf_index,
531                last_leaf: leaf_index,
532                first_mutation: mutation_index,
533                mutation_end: mutation_index + 1,
534            }),
535        }
536    }
537    candidates
538}
539
540/// Replay one candidate into a private emitter. No bytes are persisted and no
541/// cache entries are published until the caller validates CID resynchronizing.
542fn replay_mutation_island<S: Store>(
543    manager: &Prolly<S>,
544    old_leaves: &[NodeSummary],
545    mutations: &[(Vec<u8>, Option<Vec<u8>>)],
546    mut island: MutationIsland,
547    config: &Config,
548    measure_read_bytes: bool,
549) -> Result<IslandReplay, Error> {
550    if island.leaf_range.start >= island.leaf_range.end
551        || island.leaf_range.end > island.protected_end
552        || island.protected_end > old_leaves.len()
553        || island.mutation_range.start >= island.mutation_range.end
554        || island.mutation_range.end > mutations.len()
555    {
556        return Err(Error::InvalidNode);
557    }
558
559    let mutation_end = island.mutation_range.end;
560    let first_pending_mutation = island.mutation_range.start;
561    let mut mutation_index = first_pending_mutation;
562    let mut emitter = LeafEmitter::new(config)?;
563    let mut resynced_at = None;
564    let mut entries_streamed = 0u64;
565    let mut nodes_read = 0u64;
566    let mut bytes_read = 0u64;
567    let mut processed_end = island.leaf_range.start;
568
569    for leaf_index in island.leaf_range.start..island.protected_end {
570        let leaf = manager.load_arc(&old_leaves[leaf_index].cid)?;
571        nodes_read += 1;
572        if measure_read_bytes {
573            bytes_read += leaf.encoded_len() as u64;
574        }
575        if !leaf.leaf || leaf.keys.len() != leaf.vals.len() || leaf.format != config.format {
576            return Err(Error::InvalidNode);
577        }
578
579        for (key, value) in leaf.keys.iter().cloned().zip(leaf.vals.iter().cloned()) {
580            while mutation_index < mutation_end && mutations[mutation_index].0 < key {
581                let (mutation_key, mutation_value) = &mutations[mutation_index];
582                if let Some(value) = mutation_value {
583                    emitter.push(mutation_key.clone(), value.clone())?;
584                    entries_streamed += 1;
585                }
586                mutation_index += 1;
587            }
588            if mutation_index < mutation_end && mutations[mutation_index].0 == key {
589                if let Some(value) = &mutations[mutation_index].1 {
590                    emitter.push(key, value.clone())?;
591                    entries_streamed += 1;
592                }
593                mutation_index += 1;
594            } else {
595                emitter.push(key, value)?;
596                entries_streamed += 1;
597            }
598        }
599
600        let next_first = old_leaves.get(leaf_index + 1).map(|leaf| &leaf.first_key);
601        while mutation_index < mutation_end
602            && next_first
603                .map(|next| mutations[mutation_index].0 < *next)
604                .unwrap_or(true)
605        {
606            let (mutation_key, mutation_value) = &mutations[mutation_index];
607            if let Some(value) = mutation_value {
608                emitter.push(mutation_key.clone(), value.clone())?;
609                entries_streamed += 1;
610            }
611            mutation_index += 1;
612        }
613
614        processed_end = leaf_index + 1;
615        if mutation_index == mutation_end
616            && mutation_index > first_pending_mutation
617            && emitter.is_aligned_with(&old_leaves[leaf_index])
618        {
619            resynced_at = Some(leaf_index);
620            break;
621        }
622    }
623
624    emitter.flush()?;
625    island.leaf_range.end = processed_end;
626    let summaries = emitter
627        .emitted
628        .iter()
629        .map(|leaf| leaf.summary.clone())
630        .collect();
631    Ok(IslandReplay {
632        island,
633        summaries,
634        emitted: emitter.emitted,
635        resynced_at,
636        entries_streamed,
637        nodes_read,
638        bytes_read,
639    })
640}
641
642fn execute_mutation_islands<S: Store>(
643    manager: &Prolly<S>,
644    old_leaves: &[NodeSummary],
645    mutations: &[(Vec<u8>, Option<Vec<u8>>)],
646    config: &Config,
647    islands: Vec<MutationIsland>,
648    policy: ExecutionPolicy,
649    measure_read_bytes: bool,
650) -> Result<(Vec<IslandReplay>, usize), Error> {
651    let policy = policy.limit_to(islands.len());
652    if !policy.enabled() {
653        let mut replayed = Vec::with_capacity(islands.len());
654        for island in islands {
655            let replay = replay_mutation_island(
656                manager,
657                old_leaves,
658                mutations,
659                island,
660                config,
661                measure_read_bytes,
662            )?;
663            let proved_independent = replay.proved_independent();
664            replayed.push(replay);
665            if !proved_independent {
666                break;
667            }
668        }
669        return Ok((replayed, 0));
670    }
671
672    let mut replayed = Vec::with_capacity(islands.len());
673    let mut parallel_tasks = 0usize;
674    let mut islands = islands.into_iter();
675    loop {
676        let wave = islands
677            .by_ref()
678            .take(policy.wave_size())
679            .collect::<Vec<_>>();
680        if wave.is_empty() {
681            break;
682        }
683
684        let ranges = policy.ranges(wave.len());
685        parallel_tasks = parallel_tasks.saturating_add(ranges.len());
686        let mut wave = wave.into_iter();
687        let partitions = ranges
688            .into_iter()
689            .map(|range| wave.by_ref().take(range.len()).collect::<Vec<_>>())
690            .collect::<Vec<_>>();
691        let partition_results = partitions
692            .into_par_iter()
693            .map(|partition| {
694                partition
695                    .into_iter()
696                    .map(|island| {
697                        replay_mutation_island(
698                            manager,
699                            old_leaves,
700                            mutations,
701                            island,
702                            config,
703                            measure_read_bytes,
704                        )
705                    })
706                    .collect::<Result<Vec<_>, Error>>()
707            })
708            .collect::<Vec<_>>();
709
710        let mut wave_proved_independent = true;
711        // Partition collection is indexed, so checking results left-to-right
712        // returns the first error in canonical island order.
713        for result in partition_results {
714            let partition = result?;
715            wave_proved_independent &= partition.iter().all(IslandReplay::proved_independent);
716            replayed.extend(partition);
717        }
718        if !wave_proved_independent {
719            // Finish the bounded wave already in flight, then fall back before
720            // any later speculative wave starts.
721            break;
722        }
723    }
724    Ok((replayed, parallel_tasks))
725}
726
727fn merge_proved_island_replays(
728    old_leaves: &[NodeSummary],
729    replays: Vec<IslandReplay>,
730    stats: &mut WriteStats,
731) -> Result<StructuralIslandReplay, Error> {
732    let emitted_capacity = replays.iter().map(|replay| replay.emitted.len()).sum();
733    let mut summaries = Vec::with_capacity(old_leaves.len());
734    let mut emitted = Vec::with_capacity(emitted_capacity);
735    let mut old_cursor = 0usize;
736
737    for replay in replays {
738        if !replay.proved_independent()
739            || replay.island.leaf_range.start < old_cursor
740            || replay.island.leaf_range.end > old_leaves.len()
741        {
742            return Err(Error::InvalidNode);
743        }
744        summaries.extend_from_slice(&old_leaves[old_cursor..replay.island.leaf_range.start]);
745        stats.nodes_reused += replay.island.leaf_range.start.saturating_sub(old_cursor) as u64;
746        summaries.extend(replay.summaries);
747        emitted.extend(replay.emitted);
748        old_cursor = replay.island.leaf_range.end;
749    }
750    summaries.extend_from_slice(&old_leaves[old_cursor..]);
751    stats.nodes_reused += old_leaves.len().saturating_sub(old_cursor) as u64;
752    Ok(StructuralIslandReplay { summaries, emitted })
753}
754
755fn try_parallel_structural_islands<S: Store>(
756    manager: &Prolly<S>,
757    tree: &Tree,
758    old_leaves: &[NodeSummary],
759    mutations: &[(Vec<u8>, Option<Vec<u8>>)],
760    policy: ExecutionPolicy,
761    stats: &mut WriteStats,
762    measure_read_bytes: bool,
763) -> Result<Option<StructuralIslandReplay>, Error> {
764    if !policy.enabled() {
765        return Ok(None);
766    }
767    if !structural_mutations_can_form_distant_islands(old_leaves, mutations) {
768        return Ok(None);
769    }
770
771    let initial_count = mutation_island_candidates(old_leaves, mutations).len();
772    let islands = plan_mutation_islands(old_leaves, mutations);
773    stats.structural_islands += initial_count as u64;
774    stats.coalesced_islands += initial_count.saturating_sub(islands.len()) as u64;
775    let island_policy = policy.limit_to(islands.len());
776    if !island_policy.enabled()
777        || !structural_islands_worth_speculating(
778            initial_count,
779            &islands,
780            old_leaves.len(),
781            island_policy,
782        )
783    {
784        return Ok(None);
785    }
786    stats.parallel_width = island_policy.width() as u64;
787
788    let execution_policy = policy.limit_to(islands.len());
789    let (replays, parallel_tasks) = execute_mutation_islands(
790        manager,
791        old_leaves,
792        mutations,
793        &tree.config,
794        islands,
795        execution_policy,
796        measure_read_bytes,
797    )?;
798    stats.parallel_tasks += parallel_tasks as u64;
799    for replay in &replays {
800        stats.entries_streamed += replay.entries_streamed;
801        stats.nodes_read += replay.nodes_read;
802        stats.bytes_read += replay.bytes_read;
803        stats.resync_distance_nodes += replay.nodes_read;
804    }
805
806    if replays.iter().all(IslandReplay::proved_independent) {
807        return Ok(Some(merge_proved_island_replays(
808            old_leaves, replays, stats,
809        )?));
810    }
811
812    // One failed CID proof is enough to invalidate independence. Retrying
813    // progressively larger speculative regions can turn a bounded miss into
814    // O(n log n) replay work, so fall back immediately.
815    Ok(None)
816}
817
818/// Fast rejection for a dense mutation span. At least two guarded islands
819/// require gaps wider than the predecessor context plus the proof guard. If
820/// the average leaf spacing cannot provide that gap, candidate planning can
821/// only coalesce the work into one local region.
822fn structural_mutations_can_form_distant_islands(
823    old_leaves: &[NodeSummary],
824    mutations: &[(Vec<u8>, Option<Vec<u8>>)],
825) -> bool {
826    let Some((first, last)) = mutations.first().zip(mutations.last()) else {
827        return false;
828    };
829    let target_leaf = |key: &[u8]| {
830        old_leaves
831            .partition_point(|leaf| leaf.first_key.as_slice() <= key)
832            .saturating_sub(1)
833    };
834    let first_leaf = target_leaf(&first.0);
835    let last_leaf = target_leaf(&last.0);
836    let leaf_span = last_leaf.saturating_sub(first_leaf).saturating_add(1);
837    let minimum_separation = MUTATION_ISLAND_GUARD_LEAVES.saturating_add(3);
838    leaf_span > 1 && mutations.len().saturating_mul(minimum_separation) < leaf_span
839}
840
841/// Reject island layouts whose up-front guard coalescing already proves that
842/// speculative replay would duplicate most of the canonical sequential scan.
843/// This is a scheduling decision only; rejected layouts immediately use the
844/// same canonical fallback that would run after a failed proof.
845fn structural_islands_worth_speculating(
846    candidate_count: usize,
847    islands: &[MutationIsland],
848    leaf_count: usize,
849    policy: ExecutionPolicy,
850) -> bool {
851    if islands.len() < 2 || leaf_count == 0 || !policy.enabled() {
852        return false;
853    }
854
855    // Many small candidates collapsing into very few guarded regions is the
856    // dense/random case. Replaying those regions and then falling back can
857    // read the tree twice without exposing enough independent CPU work.
858    if candidate_count > islands.len().saturating_mul(4) {
859        return false;
860    }
861
862    let guarded_leaves = islands.iter().fold(0usize, |total, island| {
863        total.saturating_add(island.protected_end.saturating_sub(island.leaf_range.start))
864    });
865    guarded_leaves.saturating_mul(4) <= leaf_count
866}
867
868fn try_localized_height_two_deletes<S: Store>(
869    manager: &Prolly<S>,
870    tree: &Tree,
871    mutations: &[(Vec<u8>, Option<Vec<u8>>)],
872    stats: &mut WriteStats,
873    measure_read_bytes: bool,
874) -> Result<Option<(Tree, WriteStats)>, Error> {
875    if mutations.len() < 2
876        || mutations.iter().any(|(_, value)| value.is_some())
877        || matches!(
878            tree.config.format.node_layout,
879            NodeLayoutSpec::Custom { .. }
880        )
881    {
882        return Ok(None);
883    }
884    let Some(root_cid) = &tree.root else {
885        return Ok(None);
886    };
887    let root = manager.load_arc(root_cid)?;
888    stats.nodes_read += 1;
889    if measure_read_bytes {
890        stats.bytes_read += root.encoded_len() as u64;
891    }
892    if root.leaf
893        || root.level != 2
894        || root.keys.is_empty()
895        || root.keys.len() != root.vals.len()
896        || root.child_counts.len() != root.len()
897    {
898        return Ok(None);
899    }
900
901    let first_key = &mutations[0].0;
902    let last_key = &mutations[mutations.len() - 1].0;
903    let first_child = root
904        .keys
905        .partition_point(|separator| separator.as_slice() <= first_key.as_slice())
906        .saturating_sub(1);
907    let last_child = root
908        .keys
909        .partition_point(|separator| separator.as_slice() <= last_key.as_slice())
910        .saturating_sub(1);
911    let window_start = first_child.saturating_sub(1);
912    let window_end = last_child.saturating_add(3).min(root.len());
913    if window_start >= window_end {
914        return Ok(None);
915    }
916
917    let window_cids = root.vals[window_start..window_end]
918        .iter()
919        .map(|value| child_cid(value))
920        .collect::<Result<Vec<_>, _>>()?;
921    let window_nodes = if window_cids.len() > 1 {
922        manager.load_many_ordered(&window_cids)?
923    } else {
924        window_cids
925            .iter()
926            .map(|cid| manager.load_arc(cid))
927            .collect::<Result<Vec<_>, _>>()?
928    };
929    let mut old_leaves = Vec::new();
930    for node in &window_nodes {
931        stats.nodes_read += 1;
932        if measure_read_bytes {
933            stats.bytes_read += node.encoded_len() as u64;
934        }
935        if node.leaf
936            || node.level != 1
937            || node.format != tree.config.format
938            || node.keys.len() != node.vals.len()
939            || node.child_counts.len() != node.len()
940            || node.child_counts.contains(&0)
941        {
942            return Ok(None);
943        }
944        for index in 0..node.len() {
945            old_leaves.push(NodeSummary {
946                cid: child_cid(&node.vals[index])?,
947                first_key: node.keys[index].clone(),
948                count: node.child_counts[index],
949            });
950        }
951    }
952    if old_leaves.is_empty() {
953        return Ok(None);
954    }
955
956    let replay_start = old_leaves
957        .partition_point(|leaf| leaf.first_key.as_slice() <= first_key.as_slice())
958        .saturating_sub(2);
959    {
960        let last_mutation_leaf = old_leaves
961            .partition_point(|leaf| leaf.first_key.as_slice() <= last_key.as_slice())
962            .saturating_sub(1);
963        let prefetch_end = last_mutation_leaf.saturating_add(2).min(old_leaves.len());
964        if prefetch_end.saturating_sub(replay_start) > 1 {
965            let leaf_cids = old_leaves[replay_start..prefetch_end]
966                .iter()
967                .map(|leaf| leaf.cid.clone())
968                .collect::<Vec<_>>();
969            let _ = manager.load_many_ordered(&leaf_cids)?;
970        }
971    }
972    let mut mutation_index = 0usize;
973    let mut emitter = LeafEmitter::new(&tree.config)?;
974    let mut resynced_at = None;
975    for leaf_index in replay_start..old_leaves.len() {
976        let leaf = manager.load_arc(&old_leaves[leaf_index].cid)?;
977        stats.nodes_read += 1;
978        if measure_read_bytes {
979            stats.bytes_read += leaf.encoded_len() as u64;
980        }
981        if !leaf.leaf || leaf.keys.len() != leaf.vals.len() {
982            return Err(Error::InvalidNode);
983        }
984
985        for (key, value) in leaf.keys.iter().cloned().zip(leaf.vals.iter().cloned()) {
986            while mutation_index < mutations.len() && mutations[mutation_index].0 < key {
987                mutation_index += 1;
988            }
989            if mutation_index < mutations.len() && mutations[mutation_index].0 == key {
990                mutation_index += 1;
991            } else {
992                emitter.push(key, value)?;
993                stats.entries_streamed += 1;
994            }
995        }
996
997        let next_first = old_leaves.get(leaf_index + 1).map(|leaf| &leaf.first_key);
998        while mutation_index < mutations.len()
999            && next_first
1000                .map(|next| mutations[mutation_index].0 < *next)
1001                .unwrap_or(true)
1002        {
1003            mutation_index += 1;
1004        }
1005
1006        stats.resync_distance_nodes += 1;
1007        if mutation_index > 0 && emitter.is_aligned_with(&old_leaves[leaf_index]) {
1008            resynced_at = Some(leaf_index);
1009            break;
1010        }
1011    }
1012    if mutation_index != mutations.len() {
1013        return Ok(None);
1014    }
1015    if resynced_at.is_none() && window_end < root.len() {
1016        return Ok(None);
1017    }
1018    emitter.flush()?;
1019
1020    let old_cursor = resynced_at.map_or(old_leaves.len(), |index| index + 1);
1021    let mut leaf_summaries = Vec::with_capacity(old_leaves.len());
1022    leaf_summaries.extend_from_slice(&old_leaves[..replay_start]);
1023    leaf_summaries.extend(emitter.emitted.iter().map(|leaf| leaf.summary.clone()));
1024    leaf_summaries.extend_from_slice(&old_leaves[old_cursor..]);
1025    stats.nodes_reused += replay_start.saturating_add(old_leaves.len() - old_cursor) as u64;
1026    stats.resync_distance_entries = stats.entries_streamed;
1027
1028    let builder = BatchBuilder::new(manager.store(), tree.config.clone());
1029    let (replacement_summaries, internal_nodes) =
1030        builder.build_level_serial_deferred(leaf_summaries, 1)?;
1031    if window_end < root.len()
1032        && replacement_summaries.last().map(|summary| &summary.cid) != window_cids.last()
1033    {
1034        return Ok(None);
1035    }
1036    if replacement_summaries.is_empty() {
1037        return Ok(None);
1038    }
1039
1040    let mut updated_root = (*root).clone();
1041    updated_root.keys.splice(
1042        window_start..window_end,
1043        replacement_summaries
1044            .iter()
1045            .map(|summary| summary.first_key.clone()),
1046    );
1047    updated_root.vals.splice(
1048        window_start..window_end,
1049        replacement_summaries
1050            .iter()
1051            .map(|summary| summary.cid.0.to_vec()),
1052    );
1053    updated_root.child_counts.splice(
1054        window_start..window_end,
1055        replacement_summaries.iter().map(|summary| summary.count),
1056    );
1057    if updated_root.keys.windows(2).any(|pair| pair[0] >= pair[1])
1058        || updated_root.len() > updated_root.max_chunk_size()
1059        || updated_root.encoded_len() as u64 > tree.config.format.chunking.hard_max_node_bytes
1060    {
1061        return Ok(None);
1062    }
1063
1064    let rebuilt_root = if updated_root.len() == 1 {
1065        None
1066    } else {
1067        let candidate_children = updated_root
1068            .keys
1069            .iter()
1070            .zip(&updated_root.vals)
1071            .zip(&updated_root.child_counts)
1072            .map(|((key, value), count)| {
1073                Ok(NodeSummary {
1074                    cid: child_cid(value)?,
1075                    first_key: key.clone(),
1076                    count: *count,
1077                })
1078            })
1079            .collect::<Result<Vec<_>, Error>>()?;
1080        let (root_summaries, mut root_nodes) =
1081            builder.build_level_serial_deferred(candidate_children, 2)?;
1082        if root_summaries.len() != 1 || root_nodes.len() != 1 {
1083            return Ok(None);
1084        }
1085        let rebuilt_root = root_nodes.pop().ok_or(Error::InvalidNode)?;
1086        if root_summaries[0].cid != rebuilt_root.cid || rebuilt_root.node != updated_root {
1087            return Ok(None);
1088        }
1089        Some(rebuilt_root)
1090    };
1091
1092    let old_leaf_cids = old_leaves
1093        .iter()
1094        .map(|summary| summary.cid.clone())
1095        .collect::<HashSet<_>>();
1096    let old_internal_cids = window_cids.iter().cloned().collect::<HashSet<_>>();
1097    let mut writes = Vec::<(Cid, Vec<u8>, Node)>::new();
1098    for leaf in emitter.emitted {
1099        if !old_leaf_cids.contains(&leaf.summary.cid) {
1100            writes.push((leaf.summary.cid, leaf.bytes, leaf.node));
1101        }
1102    }
1103    for node in internal_nodes {
1104        if !old_internal_cids.contains(&node.cid) {
1105            writes.push((node.cid, node.bytes, node.node));
1106        }
1107    }
1108
1109    let root = if updated_root.len() == 1 {
1110        child_cid(&updated_root.vals[0])?
1111    } else {
1112        let rebuilt_root = rebuilt_root.ok_or(Error::InvalidNode)?;
1113        let cid = rebuilt_root.cid.clone();
1114        if cid != *root_cid {
1115            writes.push((cid.clone(), rebuilt_root.bytes, rebuilt_root.node));
1116        }
1117        cid
1118    };
1119    if !writes.is_empty() {
1120        let entries = writes
1121            .iter()
1122            .map(|(cid, bytes, _)| (cid.as_bytes(), bytes.as_slice()))
1123            .collect::<Vec<_>>();
1124        manager
1125            .store()
1126            .batch_put(&entries)
1127            .map_err(|error| Error::Store(Box::new(error)))?;
1128        let bytes_written = entries.iter().map(|(_, bytes)| bytes.len()).sum::<usize>();
1129        manager.record_batch_write_metrics(entries.len(), bytes_written);
1130        stats.nodes_written += entries.len() as u64;
1131        stats.bytes_written += bytes_written as u64;
1132        if entries.len() <= LOCAL_WRITE_CACHE_LIMIT {
1133            drop(entries);
1134            for (cid, _, node) in writes {
1135                manager.cache_node(cid, node);
1136            }
1137        }
1138    }
1139
1140    Ok(Some((
1141        Tree {
1142            root: Some(root),
1143            config: tree.config.clone(),
1144        },
1145        *stats,
1146    )))
1147}
1148
1149fn try_append<S: Store>(
1150    manager: &Prolly<S>,
1151    tree: &Tree,
1152    mutations: &mut Vec<(Vec<u8>, Option<Vec<u8>>)>,
1153    stats: &mut WriteStats,
1154    measure_read_bytes: bool,
1155) -> Result<Option<(Tree, WriteStats)>, Error> {
1156    if mutations.iter().any(|(_, value)| value.is_none()) {
1157        return Ok(None);
1158    }
1159    let path = rightmost_internal_path(manager, tree)?;
1160    let last_cid = match path.last() {
1161        Some((_, node)) => child_cid(node.vals.last().ok_or(Error::InvalidNode)?)?,
1162        None => tree.root.clone().ok_or(Error::InvalidNode)?,
1163    };
1164    let last_leaf = manager.load_arc(&last_cid)?;
1165    stats.nodes_read += 1;
1166    if measure_read_bytes {
1167        stats.bytes_read += last_leaf.encoded_len() as u64;
1168    }
1169    if !last_leaf.leaf || last_leaf.keys.len() != last_leaf.vals.len() {
1170        return Err(Error::InvalidNode);
1171    }
1172    let Some(max_key) = last_leaf.keys.last() else {
1173        return Err(Error::InvalidNode);
1174    };
1175    if mutations
1176        .first()
1177        .map_or(true, |(key, _)| key.as_slice() <= max_key.as_slice())
1178    {
1179        return Ok(None);
1180    }
1181
1182    let mut emitter = LeafEmitter::new(&tree.config)?;
1183    for (key, value) in last_leaf
1184        .keys
1185        .iter()
1186        .cloned()
1187        .zip(last_leaf.vals.iter().cloned())
1188    {
1189        emitter.push(key, value)?;
1190        stats.entries_streamed += 1;
1191    }
1192    for (key, value) in mutations.drain(..) {
1193        emitter.push(
1194            key,
1195            value.expect("append path rejects deletes before streaming"),
1196        )?;
1197        stats.entries_streamed += 1;
1198    }
1199    emitter.flush()?;
1200
1201    let builder = BatchBuilder::new(manager.store(), tree.config.clone());
1202    let mut current = emitter
1203        .emitted
1204        .iter()
1205        .map(|leaf| leaf.summary.clone())
1206        .collect::<Vec<_>>();
1207    let mut internal_nodes = Vec::new();
1208    let mut current_level = 0u8;
1209
1210    for (_, node) in path.iter().rev() {
1211        let mut children = internal_child_summaries(node)?;
1212        children.pop().ok_or(Error::InvalidNode)?;
1213        children.extend(current);
1214        let (summaries, pending) = builder.build_level_serial_deferred(children, node.level)?;
1215        current = summaries;
1216        internal_nodes.extend(pending);
1217        current_level = node.level;
1218    }
1219    while current.len() > 1 {
1220        current_level = current_level.checked_add(1).ok_or(Error::InvalidNode)?;
1221        let (summaries, pending) = builder.build_level_serial_deferred(current, current_level)?;
1222        current = summaries;
1223        internal_nodes.extend(pending);
1224    }
1225    let root = current.into_iter().next().ok_or(Error::InvalidNode)?.cid;
1226
1227    let old_internal_cids = path
1228        .iter()
1229        .map(|(cid, _)| cid.clone())
1230        .collect::<HashSet<_>>();
1231    let changed_leaves = emitter
1232        .emitted
1233        .iter()
1234        .filter(|leaf| leaf.summary.cid != last_cid)
1235        .collect::<Vec<_>>();
1236    let writes = changed_leaves
1237        .iter()
1238        .map(|leaf| (leaf.summary.cid.as_bytes(), leaf.bytes.as_slice()))
1239        .chain(
1240            internal_nodes
1241                .iter()
1242                .filter(|node| !old_internal_cids.contains(&node.cid))
1243                .map(|node| (node.cid.as_bytes(), node.bytes.as_slice())),
1244        )
1245        .collect::<Vec<_>>();
1246    if !writes.is_empty() {
1247        manager
1248            .store()
1249            .batch_put(&writes)
1250            .map_err(|error| Error::Store(Box::new(error)))?;
1251        let bytes_written = writes.iter().map(|(_, bytes)| bytes.len()).sum::<usize>();
1252        manager.record_batch_write_metrics(writes.len(), bytes_written);
1253        stats.nodes_written += writes.len() as u64;
1254        stats.bytes_written += bytes_written as u64;
1255        if writes.len() <= LOCAL_WRITE_CACHE_LIMIT {
1256            for leaf in changed_leaves {
1257                manager.cache_node(leaf.summary.cid.clone(), leaf.node.clone());
1258            }
1259            for node in internal_nodes
1260                .iter()
1261                .filter(|node| !old_internal_cids.contains(&node.cid))
1262            {
1263                manager.cache_node(node.cid.clone(), node.node.clone());
1264            }
1265        }
1266    }
1267    stats.resync_distance_entries = stats.entries_streamed;
1268    stats.resync_distance_nodes = 1;
1269    Ok(Some((
1270        Tree {
1271            root: Some(root),
1272            config: tree.config.clone(),
1273        },
1274        *stats,
1275    )))
1276}
1277
1278fn rightmost_internal_path<S: Store>(
1279    manager: &Prolly<S>,
1280    tree: &Tree,
1281) -> Result<Vec<(Cid, std::sync::Arc<Node>)>, Error> {
1282    let Some(mut cid) = tree.root.clone() else {
1283        return Ok(Vec::new());
1284    };
1285    let mut path = Vec::new();
1286    loop {
1287        let node = manager.load_arc(&cid)?;
1288        if node.leaf {
1289            return Ok(path);
1290        }
1291        let child = node.vals.last().ok_or(Error::InvalidNode)?;
1292        let next = child_cid(child)?;
1293        path.push((cid, node));
1294        cid = next;
1295    }
1296}
1297
1298fn internal_child_summaries(node: &Node) -> Result<Vec<NodeSummary>, Error> {
1299    if node.leaf || node.keys.len() != node.vals.len() || node.child_counts.len() != node.len() {
1300        return Err(Error::InvalidNode);
1301    }
1302    node.keys
1303        .iter()
1304        .zip(&node.vals)
1305        .zip(&node.child_counts)
1306        .map(|((key, value), count)| {
1307            Ok(NodeSummary {
1308                cid: child_cid(value)?,
1309                first_key: key.clone(),
1310                count: *count,
1311            })
1312        })
1313        .collect()
1314}
1315
1316enum DirectValueUpdateAttempt {
1317    Applied(Box<(Tree, WriteStats)>),
1318    Fallback(Vec<(Vec<u8>, Option<Vec<u8>>)>),
1319}
1320
1321fn sampled_value_updates_are_likely_key_stable<S: Store>(
1322    manager: &Prolly<S>,
1323    tree: &Tree,
1324    mutations: &[(Vec<u8>, Option<Vec<u8>>)],
1325    policy: ExecutionPolicy,
1326) -> Result<bool, Error> {
1327    // A single representative route catches pure insert batches without
1328    // duplicating meaningful work for the dominant key-stable update case.
1329    // The full route still proves every mutation before applying anything.
1330    const PREFLIGHT_SAMPLES: usize = 1;
1331
1332    let sample_count = mutations.len().min(PREFLIGHT_SAMPLES);
1333    if sample_count == 0 {
1334        return Ok(false);
1335    }
1336    let indexes = (0..sample_count)
1337        .map(|sample| {
1338            if sample_count == 1 {
1339                0
1340            } else {
1341                sample * (mutations.len() - 1) / (sample_count - 1)
1342            }
1343        })
1344        .collect::<Vec<_>>();
1345    let samples = indexes
1346        .into_iter()
1347        .map(|index| Mutation::Upsert {
1348            key: mutations[index].0.clone(),
1349            val: mutations[index]
1350                .1
1351                .clone()
1352                .expect("direct value path rejects deletes before routing"),
1353        })
1354        .collect::<Vec<_>>();
1355    super::batch::sampled_value_updates_are_likely_key_stable(manager, tree, samples, policy)
1356}
1357
1358fn add_write_read_metric_delta(
1359    stats: &mut WriteStats,
1360    before: super::ProllyMetricsSnapshot,
1361    after: super::ProllyMetricsSnapshot,
1362    measure_read_bytes: bool,
1363) {
1364    stats.nodes_read += after.nodes_read.saturating_sub(before.nodes_read);
1365    if measure_read_bytes {
1366        stats.bytes_read += after.bytes_read.saturating_sub(before.bytes_read);
1367    }
1368}
1369
1370fn try_direct_value_updates<S: Store>(
1371    manager: &Prolly<S>,
1372    tree: &Tree,
1373    mutations: Vec<(Vec<u8>, Option<Vec<u8>>)>,
1374    stats: &mut WriteStats,
1375    measure_read_bytes: bool,
1376    policy: ExecutionPolicy,
1377) -> Result<DirectValueUpdateAttempt, Error> {
1378    let chunking = &tree.config.format.chunking;
1379    if chunking.measure != ChunkMeasure::EntryCount
1380        || chunking.input != BoundaryInput::Key
1381        || matches!(
1382            tree.config.format.node_layout,
1383            NodeLayoutSpec::Custom { .. }
1384        )
1385        || mutations.iter().any(|(_, value)| value.is_none())
1386    {
1387        return Ok(DirectValueUpdateAttempt::Fallback(mutations));
1388    }
1389    let Some(root) = &tree.root else {
1390        return Ok(DirectValueUpdateAttempt::Fallback(mutations));
1391    };
1392
1393    let mutations =
1394        if super::batch::should_try_batched_value_updates(manager, tree, mutations.len(), policy) {
1395            let metrics_before = manager.metrics();
1396            if sampled_value_updates_are_likely_key_stable(manager, tree, &mutations, policy)? {
1397                let batched_mutations = mutations
1398                    .into_iter()
1399                    .map(|(key, value)| Mutation::Upsert {
1400                        key,
1401                        val: value.expect("direct value path rejects deletes before routing"),
1402                    })
1403                    .collect::<Vec<_>>();
1404                let attempt = super::batch::try_apply_batched_value_updates(
1405                    manager,
1406                    tree,
1407                    batched_mutations,
1408                    policy,
1409                )?;
1410                add_write_read_metric_delta(
1411                    stats,
1412                    metrics_before,
1413                    manager.metrics(),
1414                    measure_read_bytes,
1415                );
1416                match attempt {
1417                    super::batch::KeyStableBatchAttempt::Applied(result) => {
1418                        stats.nodes_written += result.written_nodes as u64;
1419                        stats.bytes_written += result.written_bytes as u64;
1420                        stats.entries_streamed += result.entries_streamed as u64;
1421                        stats.resync_distance_entries = result.entries_streamed as u64;
1422                        stats.resync_distance_nodes = result.affected_leaves as u64;
1423                        stats.used_key_stable_fast_path = true;
1424                        stats.used_batched_value_update_path = true;
1425                        stats.parallel_width = result.parallel_width as u64;
1426                        stats.parallel_tasks += result.parallel_tasks as u64;
1427                        return Ok(DirectValueUpdateAttempt::Applied(Box::new((
1428                            result.tree,
1429                            *stats,
1430                        ))));
1431                    }
1432                    super::batch::KeyStableBatchAttempt::Fallback {
1433                        mutations,
1434                        parallel_width,
1435                        parallel_tasks,
1436                    } => {
1437                        stats.parallel_width = stats.parallel_width.max(parallel_width as u64);
1438                        stats.parallel_tasks += parallel_tasks as u64;
1439                        mutations
1440                            .into_iter()
1441                            .map(mutation_parts)
1442                            .collect::<Vec<_>>()
1443                    }
1444                }
1445            } else {
1446                add_write_read_metric_delta(
1447                    stats,
1448                    metrics_before,
1449                    manager.metrics(),
1450                    measure_read_bytes,
1451                );
1452                mutations
1453            }
1454        } else {
1455            mutations
1456        };
1457
1458    let mut leaves = Vec::new();
1459    let mut internals = Vec::new();
1460    let root_summary = {
1461        let mut context = DirectRewriteContext {
1462            leaves: &mut leaves,
1463            internals: &mut internals,
1464            stats,
1465            measure_read_bytes,
1466        };
1467        rewrite_value_update_subtree(manager, root, &mutations, true, &mut context)?
1468    };
1469    let Some(root_summary) = root_summary else {
1470        return Ok(DirectValueUpdateAttempt::Fallback(mutations));
1471    };
1472
1473    let writes = leaves
1474        .iter()
1475        .map(|leaf: &EmittedLeaf| (leaf.summary.cid.as_bytes(), leaf.bytes.as_slice()))
1476        .chain(
1477            internals
1478                .iter()
1479                .map(|node: &EmittedInternal| (node.cid.as_bytes(), node.bytes.as_slice())),
1480        )
1481        .collect::<Vec<_>>();
1482    let changed_leaf_count = leaves.len();
1483    if !writes.is_empty() {
1484        manager
1485            .store()
1486            .batch_put(&writes)
1487            .map_err(|error| Error::Store(Box::new(error)))?;
1488        let bytes_written = writes.iter().map(|(_, bytes)| bytes.len()).sum::<usize>();
1489        manager.record_batch_write_metrics(writes.len(), bytes_written);
1490        stats.nodes_written += writes.len() as u64;
1491        stats.bytes_written += bytes_written as u64;
1492        if writes.len() <= LOCAL_WRITE_CACHE_LIMIT {
1493            drop(writes);
1494            for leaf in leaves {
1495                manager.cache_node(leaf.summary.cid, leaf.node);
1496            }
1497            for node in internals {
1498                manager.cache_node(node.cid, node.node);
1499            }
1500        }
1501    }
1502    stats.resync_distance_entries = stats.entries_streamed;
1503    stats.resync_distance_nodes = changed_leaf_count as u64;
1504    stats.used_key_stable_fast_path = true;
1505    Ok(DirectValueUpdateAttempt::Applied(Box::new((
1506        Tree {
1507            root: Some(root_summary.cid),
1508            config: tree.config.clone(),
1509        },
1510        *stats,
1511    ))))
1512}
1513
1514enum DirectDelete {
1515    Applied(NodeSummary),
1516    Unchanged,
1517    Fallback,
1518}
1519
1520fn try_direct_single_delete<S: Store>(
1521    manager: &Prolly<S>,
1522    tree: &Tree,
1523    mutations: &[(Vec<u8>, Option<Vec<u8>>)],
1524    stats: &mut WriteStats,
1525    measure_read_bytes: bool,
1526) -> Result<Option<(Tree, WriteStats)>, Error> {
1527    let chunking = &tree.config.format.chunking;
1528    if mutations.len() != 1
1529        || mutations[0].1.is_some()
1530        || chunking.measure != ChunkMeasure::EntryCount
1531        || chunking.input != BoundaryInput::Key
1532        || !matches!(
1533            chunking.rule,
1534            super::format::BoundaryRule::HashThreshold { .. }
1535        )
1536        || matches!(
1537            tree.config.format.node_layout,
1538            NodeLayoutSpec::Custom { .. }
1539        )
1540    {
1541        return Ok(None);
1542    }
1543    let Some(root) = &tree.root else {
1544        return Ok(Some((tree.clone(), *stats)));
1545    };
1546    let mut leaves = Vec::new();
1547    let mut internals = Vec::new();
1548    let result = rewrite_single_delete_subtree(
1549        manager,
1550        root,
1551        &mutations[0].0,
1552        &tree.config,
1553        &mut leaves,
1554        &mut internals,
1555        stats,
1556        measure_read_bytes,
1557    )?;
1558    let root = match result {
1559        DirectDelete::Applied(summary) => summary.cid,
1560        DirectDelete::Unchanged => return Ok(Some((tree.clone(), *stats))),
1561        DirectDelete::Fallback => return Ok(None),
1562    };
1563
1564    let writes = leaves
1565        .iter()
1566        .map(|leaf: &EmittedLeaf| (leaf.summary.cid.as_bytes(), leaf.bytes.as_slice()))
1567        .chain(
1568            internals
1569                .iter()
1570                .map(|node: &EmittedInternal| (node.cid.as_bytes(), node.bytes.as_slice())),
1571        )
1572        .collect::<Vec<_>>();
1573    manager
1574        .store()
1575        .batch_put(&writes)
1576        .map_err(|error| Error::Store(Box::new(error)))?;
1577    let bytes_written = writes.iter().map(|(_, bytes)| bytes.len()).sum::<usize>();
1578    manager.record_batch_write_metrics(writes.len(), bytes_written);
1579    stats.nodes_written += writes.len() as u64;
1580    stats.bytes_written += bytes_written as u64;
1581    stats.resync_distance_nodes = 1;
1582    stats.resync_distance_entries = stats.entries_streamed;
1583    drop(writes);
1584    for leaf in leaves {
1585        manager.cache_node(leaf.summary.cid, leaf.node);
1586    }
1587    for node in internals {
1588        manager.cache_node(node.cid, node.node);
1589    }
1590    Ok(Some((
1591        Tree {
1592            root: Some(root),
1593            config: tree.config.clone(),
1594        },
1595        *stats,
1596    )))
1597}
1598
1599#[allow(clippy::too_many_arguments)]
1600fn rewrite_single_delete_subtree<S: Store>(
1601    manager: &Prolly<S>,
1602    cid: &Cid,
1603    key: &[u8],
1604    config: &super::config::Config,
1605    leaves: &mut Vec<EmittedLeaf>,
1606    internals: &mut Vec<EmittedInternal>,
1607    stats: &mut WriteStats,
1608    measure_read_bytes: bool,
1609) -> Result<DirectDelete, Error> {
1610    let node = manager.load_arc(cid)?;
1611    stats.nodes_read += 1;
1612    if measure_read_bytes {
1613        stats.bytes_read += node.encoded_len() as u64;
1614    }
1615    if node.keys.is_empty() || node.keys.len() != node.vals.len() {
1616        return Err(Error::InvalidNode);
1617    }
1618
1619    if node.leaf {
1620        let Ok(index) = node
1621            .keys
1622            .binary_search_by(|candidate| candidate.as_slice().cmp(key))
1623        else {
1624            return Ok(DirectDelete::Unchanged);
1625        };
1626        if index == 0 || node.len() == 1 {
1627            return Ok(DirectDelete::Fallback);
1628        }
1629
1630        let old_closed = entry_count_boundary(
1631            &config.format.chunking,
1632            0,
1633            node.len(),
1634            node.keys.last().ok_or(Error::InvalidNode)?,
1635        )?;
1636        let mut updated = (*node).clone();
1637        updated.keys.remove(index);
1638        updated.vals.remove(index);
1639        let new_closed = entry_count_boundary(
1640            &config.format.chunking,
1641            0,
1642            updated.len(),
1643            updated.keys.last().ok_or(Error::InvalidNode)?,
1644        )?;
1645        if new_closed != old_closed {
1646            return Ok(DirectDelete::Fallback);
1647        }
1648        stats.entries_streamed += updated.len() as u64;
1649        let bytes = updated.to_bytes();
1650        let new_cid = Cid::from_bytes(&bytes);
1651        let summary = NodeSummary {
1652            cid: new_cid.clone(),
1653            first_key: updated.keys[0].clone(),
1654            count: updated.len() as u64,
1655        };
1656        leaves.push(EmittedLeaf {
1657            summary: summary.clone(),
1658            bytes,
1659            node: updated,
1660        });
1661        return Ok(DirectDelete::Applied(summary));
1662    }
1663    if node.child_counts.len() != node.len() {
1664        return Err(Error::InvalidNode);
1665    }
1666
1667    let child_index = node
1668        .keys
1669        .partition_point(|separator| separator.as_slice() <= key)
1670        .saturating_sub(1);
1671    let child = child_cid(&node.vals[child_index])?;
1672    let replacement = match rewrite_single_delete_subtree(
1673        manager,
1674        &child,
1675        key,
1676        config,
1677        leaves,
1678        internals,
1679        stats,
1680        measure_read_bytes,
1681    )? {
1682        DirectDelete::Applied(summary) => summary,
1683        DirectDelete::Unchanged => return Ok(DirectDelete::Unchanged),
1684        DirectDelete::Fallback => return Ok(DirectDelete::Fallback),
1685    };
1686
1687    let old_child_count = node.child_counts[child_index];
1688    let mut updated = (*node).clone();
1689    updated.vals[child_index] = replacement.cid.0.to_vec();
1690    updated.child_counts[child_index] = replacement.count;
1691    let first_key = updated.keys[0].clone();
1692    let count = node
1693        .child_counts
1694        .iter()
1695        .copied()
1696        .sum::<u64>()
1697        .saturating_sub(old_child_count)
1698        .saturating_add(replacement.count);
1699    let bytes = updated.to_bytes();
1700    let new_cid = Cid::from_bytes(&bytes);
1701    internals.push(EmittedInternal {
1702        cid: new_cid.clone(),
1703        bytes,
1704        node: updated,
1705    });
1706    Ok(DirectDelete::Applied(NodeSummary {
1707        cid: new_cid,
1708        first_key,
1709        count,
1710    }))
1711}
1712
1713struct DirectRewriteContext<'a> {
1714    leaves: &'a mut Vec<EmittedLeaf>,
1715    internals: &'a mut Vec<EmittedInternal>,
1716    stats: &'a mut WriteStats,
1717    measure_read_bytes: bool,
1718}
1719
1720fn rewrite_value_update_subtree<S: Store>(
1721    manager: &Prolly<S>,
1722    cid: &Cid,
1723    mutations: &[(Vec<u8>, Option<Vec<u8>>)],
1724    rightmost: bool,
1725    context: &mut DirectRewriteContext<'_>,
1726) -> Result<Option<NodeSummary>, Error> {
1727    let node = manager.load_arc(cid)?;
1728    context.stats.nodes_read += 1;
1729    if context.measure_read_bytes {
1730        context.stats.bytes_read += node.encoded_len() as u64;
1731    }
1732    if node.keys.is_empty() || node.keys.len() != node.vals.len() {
1733        return Err(Error::InvalidNode);
1734    }
1735
1736    if node.leaf {
1737        if !rightmost {
1738            let Some(last_key) = node.keys.last() else {
1739                return Err(Error::InvalidNode);
1740            };
1741            if !super::boundary::entry_count_boundary(
1742                &node.format.chunking,
1743                u16::from(node.level),
1744                node.len(),
1745                last_key,
1746            )? {
1747                return Ok(None);
1748            }
1749        }
1750        let mut updated = (*node).clone();
1751        for (key, value) in mutations {
1752            let Ok(index) = updated
1753                .keys
1754                .binary_search_by(|candidate| candidate.as_slice().cmp(key.as_slice()))
1755            else {
1756                return Ok(None);
1757            };
1758            updated.vals[index] = value
1759                .clone()
1760                .expect("direct value path rejects deletes before routing");
1761        }
1762        context.stats.entries_streamed += updated.len() as u64;
1763        let bytes = updated.to_bytes();
1764        let hard_max =
1765            usize::try_from(updated.format.chunking.hard_max_node_bytes).unwrap_or(usize::MAX);
1766        if node.encoded_len() >= hard_max || bytes.len() >= hard_max {
1767            return Ok(None);
1768        }
1769        let new_cid = Cid::from_bytes(&bytes);
1770        let summary = NodeSummary {
1771            cid: new_cid.clone(),
1772            first_key: updated.keys[0].clone(),
1773            count: updated.len() as u64,
1774        };
1775        if new_cid != *cid {
1776            context.leaves.push(EmittedLeaf {
1777                summary: summary.clone(),
1778                bytes,
1779                node: updated,
1780            });
1781        }
1782        return Ok(Some(summary));
1783    }
1784    if node.child_counts.len() != node.len() {
1785        return Err(Error::InvalidNode);
1786    }
1787
1788    let mut updated = (*node).clone();
1789    let mut start = 0usize;
1790    let mut touched_children = 0usize;
1791    while start < mutations.len() {
1792        let child_index = updated
1793            .keys
1794            .partition_point(|key| key.as_slice() <= mutations[start].0.as_slice())
1795            .saturating_sub(1);
1796        let end = updated
1797            .keys
1798            .get(child_index + 1)
1799            .map_or(mutations.len(), |boundary| {
1800                start
1801                    + mutations[start..]
1802                        .partition_point(|mutation| mutation.0.as_slice() < boundary.as_slice())
1803            });
1804        let child = child_cid(&updated.vals[child_index])?;
1805        let Some(replacement) = rewrite_value_update_subtree(
1806            manager,
1807            &child,
1808            &mutations[start..end],
1809            rightmost && child_index.saturating_add(1) == updated.len(),
1810            context,
1811        )?
1812        else {
1813            return Ok(None);
1814        };
1815        updated.keys[child_index] = replacement.first_key;
1816        updated.vals[child_index] = replacement.cid.0.to_vec();
1817        updated.child_counts[child_index] = replacement.count;
1818        touched_children += 1;
1819        start = end;
1820    }
1821    context.stats.nodes_reused += updated.len().saturating_sub(touched_children) as u64;
1822
1823    let first_key = updated.keys[0].clone();
1824    let count = updated.child_counts.iter().copied().sum();
1825    let bytes = updated.to_bytes();
1826    let new_cid = Cid::from_bytes(&bytes);
1827    if new_cid != *cid {
1828        context.internals.push(EmittedInternal {
1829            cid: new_cid.clone(),
1830            bytes,
1831            node: updated,
1832        });
1833    }
1834    Ok(Some(NodeSummary {
1835        cid: new_cid,
1836        first_key,
1837        count,
1838    }))
1839}
1840
1841fn rewrite_fixed_separator_paths<S: Store>(
1842    manager: &Prolly<S>,
1843    tree: &Tree,
1844    changes: &[NodeSummary],
1845) -> Result<(Tree, Vec<EmittedInternal>), Error> {
1846    let Some(root) = &tree.root else {
1847        return Err(Error::InvalidNode);
1848    };
1849    let root_node = manager.load_arc(root)?;
1850    if root_node.leaf {
1851        let replacement = changes.first().ok_or(Error::InvalidNode)?;
1852        return Ok((
1853            Tree {
1854                root: Some(replacement.cid.clone()),
1855                config: tree.config.clone(),
1856            },
1857            Vec::new(),
1858        ));
1859    }
1860
1861    let mut pending = Vec::new();
1862    let root = rewrite_internal_node(manager, root, changes, &mut pending)?;
1863    Ok((
1864        Tree {
1865            root: Some(root.cid),
1866            config: tree.config.clone(),
1867        },
1868        pending,
1869    ))
1870}
1871
1872fn rewrite_internal_node<S: Store>(
1873    manager: &Prolly<S>,
1874    cid: &Cid,
1875    changes: &[NodeSummary],
1876    pending: &mut Vec<EmittedInternal>,
1877) -> Result<NodeSummary, Error> {
1878    let node = manager.load_arc(cid)?;
1879    if node.leaf
1880        || node.keys.is_empty()
1881        || node.keys.len() != node.vals.len()
1882        || node.child_counts.len() != node.len()
1883    {
1884        return Err(Error::InvalidNode);
1885    }
1886
1887    let mut updated = (*node).clone();
1888    let mut change_start = 0usize;
1889    while change_start < changes.len() {
1890        let child_index = updated
1891            .keys
1892            .partition_point(|key| key.as_slice() <= changes[change_start].first_key.as_slice())
1893            .saturating_sub(1);
1894        let mut change_end = change_start + 1;
1895        while change_end < changes.len()
1896            && updated
1897                .keys
1898                .partition_point(|key| key.as_slice() <= changes[change_end].first_key.as_slice())
1899                .saturating_sub(1)
1900                == child_index
1901        {
1902            change_end += 1;
1903        }
1904
1905        let replacement = if updated.level == 1 {
1906            if change_end != change_start + 1
1907                || updated.keys[child_index] != changes[change_start].first_key
1908            {
1909                return Err(Error::InvalidNode);
1910            }
1911            changes[change_start].clone()
1912        } else {
1913            let child = child_cid(&updated.vals[child_index])?;
1914            rewrite_internal_node(manager, &child, &changes[change_start..change_end], pending)?
1915        };
1916        updated.keys[child_index] = replacement.first_key;
1917        updated.vals[child_index] = replacement.cid.0.to_vec();
1918        updated.child_counts[child_index] = replacement.count;
1919        change_start = change_end;
1920    }
1921
1922    let first_key = updated.keys[0].clone();
1923    let count = updated.child_counts.iter().copied().sum();
1924    let bytes = updated.to_bytes();
1925    let new_cid = Cid::from_bytes(&bytes);
1926    if new_cid != *cid {
1927        pending.push(EmittedInternal {
1928            cid: new_cid.clone(),
1929            bytes,
1930            node: updated,
1931        });
1932    }
1933    Ok(NodeSummary {
1934        cid: new_cid,
1935        first_key,
1936        count,
1937    })
1938}
1939
1940fn normalize(mutations: Vec<Mutation>) -> Vec<(Vec<u8>, Option<Vec<u8>>)> {
1941    if mutations
1942        .windows(2)
1943        .all(|pair| pair[0].key() <= pair[1].key())
1944    {
1945        let mut normalized = Vec::<(Vec<u8>, Option<Vec<u8>>)>::with_capacity(mutations.len());
1946        for mutation in mutations {
1947            let (key, value) = mutation_parts(mutation);
1948            match normalized.last_mut() {
1949                Some((previous_key, previous_value)) if *previous_key == key => {
1950                    *previous_value = value;
1951                }
1952                _ => normalized.push((key, value)),
1953            }
1954        }
1955        return normalized;
1956    }
1957
1958    let mut sorted = mutations
1959        .into_iter()
1960        .map(mutation_parts)
1961        .collect::<Vec<_>>();
1962    sorted.sort_by(|left, right| left.0.cmp(&right.0));
1963    let mut normalized = Vec::<(Vec<u8>, Option<Vec<u8>>)>::with_capacity(sorted.len());
1964    for (key, value) in sorted {
1965        match normalized.last_mut() {
1966            Some((previous_key, previous_value)) if *previous_key == key => {
1967                *previous_value = value;
1968            }
1969            _ => normalized.push((key, value)),
1970        }
1971    }
1972    normalized
1973}
1974
1975fn mutation_parts(mutation: Mutation) -> (Vec<u8>, Option<Vec<u8>>) {
1976    match mutation {
1977        Mutation::Upsert { key, val } => (key, Some(val)),
1978        Mutation::Delete { key } => (key, None),
1979    }
1980}
1981
1982fn take_mutation(mutation: &mut (Vec<u8>, Option<Vec<u8>>)) -> (Vec<u8>, Option<Vec<u8>>) {
1983    (std::mem::take(&mut mutation.0), mutation.1.take())
1984}
1985
1986fn build_empty_base<S: Store>(
1987    manager: &Prolly<S>,
1988    tree: &Tree,
1989    mutations: Vec<(Vec<u8>, Option<Vec<u8>>)>,
1990    mut stats: WriteStats,
1991) -> Result<(Tree, WriteStats), Error> {
1992    let mut writer = super::builder::SortedBatchBuilder::new(manager.store(), tree.config.clone());
1993    for (key, value) in mutations {
1994        if let Some(value) = value {
1995            writer.add(key, value)?;
1996            stats.entries_streamed += 1;
1997        }
1998    }
1999    let tree = writer.build()?;
2000    if let Some(root) = &tree.root {
2001        let node = manager.load_arc(root)?;
2002        manager.record_batch_write_metrics(1, node.encoded_len());
2003        stats.nodes_written = 1;
2004        stats.bytes_written = node.encoded_len() as u64;
2005        stats.resync_distance_nodes = 1;
2006    }
2007    Ok((tree, stats))
2008}
2009
2010fn collect_leaf_summaries<S: Store>(
2011    manager: &Prolly<S>,
2012    tree: &Tree,
2013    stats: &mut WriteStats,
2014    measure_read_bytes: bool,
2015) -> Result<(Vec<NodeSummary>, HashSet<Cid>), Error> {
2016    let Some(root) = &tree.root else {
2017        return Ok((Vec::new(), HashSet::new()));
2018    };
2019    let mut leaves = Vec::new();
2020    let mut internals = HashSet::new();
2021    collect_from_node(
2022        manager,
2023        root,
2024        &tree.config.format,
2025        stats,
2026        &mut leaves,
2027        &mut internals,
2028        measure_read_bytes,
2029    )?;
2030    Ok((leaves, internals))
2031}
2032
2033fn collect_from_node<S: Store>(
2034    manager: &Prolly<S>,
2035    cid: &Cid,
2036    expected_format: &super::format::TreeFormat,
2037    stats: &mut WriteStats,
2038    leaves: &mut Vec<NodeSummary>,
2039    internals: &mut HashSet<Cid>,
2040    measure_read_bytes: bool,
2041) -> Result<(), Error> {
2042    let node = manager.load_arc(cid)?;
2043    stats.nodes_read += 1;
2044    if measure_read_bytes {
2045        stats.bytes_read += node.encoded_len() as u64;
2046    }
2047    if node.format != *expected_format {
2048        return Err(Error::FormatMismatch {
2049            expected: expected_format.digest()?,
2050            actual: node.format.digest()?,
2051        });
2052    }
2053    if node.leaf {
2054        leaves.push(NodeSummary {
2055            cid: cid.clone(),
2056            first_key: node.keys.first().cloned().unwrap_or_default(),
2057            count: node.keys.len() as u64,
2058        });
2059        return Ok(());
2060    }
2061    internals.insert(cid.clone());
2062    if node.keys.len() != node.vals.len() {
2063        return Err(Error::InvalidNode);
2064    }
2065
2066    if node.level == 1 {
2067        for index in 0..node.len() {
2068            let child = child_cid(&node.vals[index])?;
2069            let count = node.child_counts.get(index).copied().unwrap_or(0);
2070            let count = if count == 0 {
2071                let leaf = manager.load_arc(&child)?;
2072                stats.nodes_read += 1;
2073                if measure_read_bytes {
2074                    stats.bytes_read += leaf.encoded_len() as u64;
2075                }
2076                leaf.keys.len() as u64
2077            } else {
2078                count
2079            };
2080            leaves.push(NodeSummary {
2081                cid: child,
2082                first_key: node.keys[index].clone(),
2083                count,
2084            });
2085        }
2086        return Ok(());
2087    }
2088
2089    for value in &node.vals {
2090        collect_from_node(
2091            manager,
2092            &child_cid(value)?,
2093            expected_format,
2094            stats,
2095            leaves,
2096            internals,
2097            measure_read_bytes,
2098        )?;
2099    }
2100    Ok(())
2101}
2102
2103fn child_cid(bytes: &[u8]) -> Result<Cid, Error> {
2104    let bytes: [u8; 32] = bytes.try_into().map_err(|_| Error::InvalidNode)?;
2105    Ok(Cid(bytes))
2106}
2107
2108#[cfg(test)]
2109mod tests {
2110    use super::*;
2111    use crate::prolly::chunking;
2112    use crate::prolly::format::NodeLayoutSpec;
2113    use crate::prolly::store::MemStore;
2114    use std::sync::Arc;
2115
2116    fn synthetic_leaf_summaries(count: usize) -> Vec<NodeSummary> {
2117        (0..count)
2118            .map(|index| NodeSummary {
2119                cid: Cid::from_bytes(format!("leaf-{index:04}").as_bytes()),
2120                first_key: format!("k{index:04}").into_bytes(),
2121                count: 1,
2122            })
2123            .collect()
2124    }
2125
2126    fn populated_tree(config: Config, count: usize) -> (Prolly<Arc<MemStore>>, Tree) {
2127        let manager = Prolly::new(Arc::new(MemStore::new()), config);
2128        let mutations = (0..count)
2129            .map(|index| Mutation::Upsert {
2130                key: format!("k{index:04}").into_bytes(),
2131                val: vec![b'v'; 32],
2132            })
2133            .collect();
2134        let tree = manager.batch(&manager.create(), mutations).unwrap();
2135        (manager, tree)
2136    }
2137
2138    fn old_leaf_summaries<S: Store>(manager: &Prolly<S>, tree: &Tree) -> Vec<NodeSummary> {
2139        collect_leaf_summaries(manager, tree, &mut WriteStats::default(), false)
2140            .unwrap()
2141            .0
2142    }
2143
2144    #[test]
2145    fn mutation_island_planner_separates_distant_clusters() {
2146        let leaves = synthetic_leaf_summaries(64);
2147        let mutations = vec![
2148            (b"k0010".to_vec(), None),
2149            (b"k0040".to_vec(), Some(b"changed".to_vec())),
2150        ];
2151
2152        let islands = plan_mutation_islands(&leaves, &mutations);
2153
2154        assert_eq!(islands.len(), 2);
2155        assert_eq!(islands[0].mutation_range, 0..1);
2156        assert_eq!(islands[1].mutation_range, 1..2);
2157        assert!(islands[0].protected_end <= islands[1].leaf_range.start);
2158    }
2159
2160    #[test]
2161    fn mutation_island_planner_coalesces_adjacent_guards_and_covers_mutations_once() {
2162        let leaves = synthetic_leaf_summaries(96);
2163        let mutations = vec![
2164            (b"k0010".to_vec(), None),
2165            (b"k0012".to_vec(), Some(b"a".to_vec())),
2166            (b"k0050".to_vec(), Some(b"b".to_vec())),
2167            (b"k0080".to_vec(), None),
2168        ];
2169
2170        let islands = plan_mutation_islands(&leaves, &mutations);
2171        let covered = islands
2172            .iter()
2173            .flat_map(|island| island.mutation_range.clone())
2174            .collect::<Vec<_>>();
2175
2176        assert_eq!(covered, (0..mutations.len()).collect::<Vec<_>>());
2177        assert_eq!(islands[0].mutation_range, 0..2);
2178        assert!(islands.windows(2).all(|pair| {
2179            pair[0].protected_end <= pair[1].leaf_range.start
2180                && pair[0].mutation_range.end == pair[1].mutation_range.start
2181        }));
2182    }
2183
2184    #[test]
2185    fn structural_island_admission_rejects_dense_guards_but_keeps_sparse_work() {
2186        let leaves = synthetic_leaf_summaries(128);
2187        let dense_mutations = (0..32)
2188            .map(|index| (format!("k{:04}", index * 4).into_bytes(), None))
2189            .collect::<Vec<_>>();
2190        let dense_candidates = mutation_island_candidates(&leaves, &dense_mutations).len();
2191        let dense_islands = plan_mutation_islands(&leaves, &dense_mutations);
2192        let policy = ExecutionPolicy::from_config(
2193            &ParallelConfig::new(4, 1),
2194            dense_mutations.len(),
2195            dense_islands.len(),
2196        );
2197
2198        assert!(!structural_islands_worth_speculating(
2199            dense_candidates,
2200            &dense_islands,
2201            leaves.len(),
2202            policy,
2203        ));
2204        assert!(!structural_mutations_can_form_distant_islands(
2205            &leaves,
2206            &dense_mutations,
2207        ));
2208
2209        let sparse_mutations = vec![
2210            (b"k0010".to_vec(), None),
2211            (b"k0100".to_vec(), Some(b"changed".to_vec())),
2212        ];
2213        let sparse_candidates = mutation_island_candidates(&leaves, &sparse_mutations).len();
2214        let sparse_islands = plan_mutation_islands(&leaves, &sparse_mutations);
2215        let policy = ExecutionPolicy::from_config(
2216            &ParallelConfig::new(2, 1),
2217            sparse_mutations.len(),
2218            sparse_islands.len(),
2219        );
2220
2221        assert!(structural_islands_worth_speculating(
2222            sparse_candidates,
2223            &sparse_islands,
2224            leaves.len(),
2225            policy,
2226        ));
2227        assert!(structural_mutations_can_form_distant_islands(
2228            &leaves,
2229            &sparse_mutations,
2230        ));
2231    }
2232
2233    #[test]
2234    fn mutation_island_replay_proves_an_unchanged_anchor_without_writes() {
2235        let config = Config::builder()
2236            .min_chunk_size(2)
2237            .max_chunk_size(4)
2238            .chunking_factor(1)
2239            .build();
2240        let (manager, tree) = populated_tree(config, 128);
2241        let leaves = old_leaf_summaries(&manager, &tree);
2242        let mutations = vec![(b"k0020".to_vec(), Some(vec![b'x'; 32]))];
2243        let island = plan_mutation_islands(&leaves, &mutations)
2244            .into_iter()
2245            .next()
2246            .unwrap();
2247        let metrics_before = manager.metrics();
2248
2249        let replay =
2250            replay_mutation_island(&manager, &leaves, &mutations, island, &tree.config, true)
2251                .unwrap();
2252
2253        assert!(replay.proved_independent());
2254        assert!(replay.entries_streamed > 0);
2255        assert!(replay.nodes_read > 0);
2256        assert!(replay.bytes_read > 0);
2257        assert!(!replay.summaries.is_empty());
2258        assert!(!replay.emitted.is_empty());
2259        assert_eq!(
2260            manager.metrics().nodes_written,
2261            metrics_before.nodes_written
2262        );
2263    }
2264
2265    #[test]
2266    fn mutation_island_replay_rejects_a_cascade_that_reaches_its_guard() {
2267        let config = Config::builder()
2268            .min_chunk_size(2)
2269            .max_chunk_size(4)
2270            .chunking_factor(u32::MAX)
2271            .build();
2272        let (manager, tree) = populated_tree(config, 256);
2273        let leaves = old_leaf_summaries(&manager, &tree);
2274        let mutations = vec![(b"k0020a".to_vec(), Some(vec![b'x'; 32]))];
2275        let island = plan_mutation_islands(&leaves, &mutations)
2276            .into_iter()
2277            .next()
2278            .unwrap();
2279
2280        let replay =
2281            replay_mutation_island(&manager, &leaves, &mutations, island, &tree.config, false)
2282                .unwrap();
2283
2284        assert!(!replay.proved_independent());
2285        assert_eq!(replay.resynced_at, None);
2286        assert_eq!(replay.island.leaf_range.end, replay.island.protected_end);
2287    }
2288
2289    #[test]
2290    fn mutation_island_executor_stops_after_the_first_failed_wave() {
2291        rayon::ThreadPoolBuilder::new()
2292            .num_threads(4)
2293            .build()
2294            .unwrap()
2295            .install(|| {
2296                let config = Config::builder()
2297                    .min_chunk_size(2)
2298                    .max_chunk_size(4)
2299                    .chunking_factor(u32::MAX)
2300                    .build();
2301                let (manager, tree) = populated_tree(config, 4_096);
2302                let leaves = old_leaf_summaries(&manager, &tree);
2303                let mutations = (0..10)
2304                    .map(|island| {
2305                        (
2306                            format!("k{:04}a", 256 + island * 350).into_bytes(),
2307                            Some(vec![b'a' + island as u8; 32]),
2308                        )
2309                    })
2310                    .collect::<Vec<_>>();
2311                let islands = plan_mutation_islands(&leaves, &mutations);
2312                assert_eq!(islands.len(), 10);
2313                let policy = ExecutionPolicy::from_config(
2314                    &ParallelConfig::new(2, 1),
2315                    mutations.len(),
2316                    islands.len(),
2317                );
2318
2319                let (replays, parallel_tasks) = execute_mutation_islands(
2320                    &manager,
2321                    &leaves,
2322                    &mutations,
2323                    &tree.config,
2324                    islands,
2325                    policy,
2326                    false,
2327                )
2328                .unwrap();
2329
2330                assert_eq!(parallel_tasks, 2);
2331                assert_eq!(replays.len(), policy.wave_size());
2332                assert!(replays.iter().any(|replay| !replay.proved_independent()));
2333            });
2334    }
2335
2336    #[test]
2337    fn mutation_island_replay_for_rolling_and_weibull_never_publishes_speculation() {
2338        for mut policy in [
2339            chunking::logical_bytes_rolling_hash(),
2340            chunking::logical_bytes_key_weibull(),
2341        ] {
2342            policy.min = 96;
2343            policy.target = 192;
2344            policy.max = 384;
2345            policy.hard_max_node_bytes = 512;
2346            let config = Config::builder()
2347                .chunking(policy)
2348                .node_layout(NodeLayoutSpec::PrefixCompressed)
2349                .build();
2350            let (manager, tree) = populated_tree(config, 256);
2351            let leaves = old_leaf_summaries(&manager, &tree);
2352            let mutations = vec![(b"k0064a".to_vec(), Some(vec![b'z'; 32]))];
2353            let island = plan_mutation_islands(&leaves, &mutations)
2354                .into_iter()
2355                .next()
2356                .unwrap();
2357            let metrics_before = manager.metrics();
2358
2359            let replay =
2360                replay_mutation_island(&manager, &leaves, &mutations, island, &tree.config, true)
2361                    .unwrap();
2362
2363            if replay.proved_independent() {
2364                let anchor = replay.resynced_at.unwrap();
2365                assert_eq!(replay.summaries.last().unwrap().cid, leaves[anchor].cid);
2366            }
2367            assert_eq!(
2368                manager.metrics().nodes_written,
2369                metrics_before.nodes_written
2370            );
2371            assert_eq!(tree.root, manager.export_snapshot(&tree).unwrap().tree.root);
2372        }
2373    }
2374}