Skip to main content

bazaar/groupcompress/
gcvf.rs

1//! Pure-logic core of `GroupCompressVersionedFiles`.
2//!
3//! This module ports the orchestration that the Python
4//! `bzrformats.groupcompress.GroupCompressVersionedFiles` class performs on
5//! top of the already-Rust groupcompress block/compressor/manager code. The
6//! pyo3 layer (`crates/bazaar-py/src/groupcompress.rs`) wraps a Python index
7//! and access object and drives these helpers.
8
9use crate::groupcompress::block::GroupCompressBlock;
10use crate::knit::FileRef;
11
12/// A versioned-file key: a tuple of byte segments, the last being the
13/// version id. Groupcompress shares the keyspace type with knit.
14pub type GcKey = crate::versionedfile::Key;
15
16/// Number of bytes a fetch batch accumulates before it is flushed.
17///
18/// Mirrors `bzrformats.groupcompress.BATCH_SIZE`.
19pub const BATCH_SIZE: u64 = 1 << 16;
20
21/// Default cap on the bytes a `GroupCompressor` indexes for delta matching.
22///
23/// Mirrors `GroupCompressVersionedFiles._DEFAULT_MAX_BYTES_TO_INDEX`.
24pub const DEFAULT_MAX_BYTES_TO_INDEX: usize = 1024 * 1024;
25
26/// Identifies a single groupcompress block within a store, plus the byte
27/// range it occupies.
28///
29/// This is the cache key for the block cache and the unit `_get_blocks`
30/// fetches. Mirrors the Python `read_memo = index_memo[0:3]` triple
31/// `(index, start, stop)`. `index` is abstracted via [`FileRef`] so the
32/// pure crate does not depend on the Python graph-index object.
33#[derive(Debug, Clone, PartialEq, Eq, Hash)]
34pub struct ReadMemo<F: FileRef = String> {
35    /// Identifies which backing index/shard the block lives in.
36    pub index: F,
37    /// Byte offset of the block's start in the backing file.
38    pub start: u64,
39    /// Byte offset one past the block's end.
40    pub stop: u64,
41}
42
43impl<F: FileRef> ReadMemo<F> {
44    pub fn new(index: F, start: u64, stop: u64) -> Self {
45        ReadMemo { index, start, stop }
46    }
47
48    /// The on-disk byte length of the block.
49    pub fn byte_length(&self) -> u64 {
50        self.stop.saturating_sub(self.start)
51    }
52}
53
54/// Locates a single record: which block it lives in (`read_memo`) and the
55/// `[entry_start, entry_end)` slice of the decompressed block that holds it.
56///
57/// Mirrors the Python `index_memo = (index, start, stop, basis_end,
58/// delta_end)` 5-tuple; the trailing pair becomes [`Self::entry_start`] /
59/// [`Self::entry_end`].
60#[derive(Debug, Clone, PartialEq, Eq, Hash)]
61pub struct IndexMemo<F: FileRef = String> {
62    /// The block this record lives in.
63    pub read_memo: ReadMemo<F>,
64    /// Offset of the record inside the decompressed block.
65    pub entry_start: u64,
66    /// Offset one past the record inside the decompressed block.
67    pub entry_end: u64,
68}
69
70impl<F: FileRef> IndexMemo<F> {
71    pub fn new(read_memo: ReadMemo<F>, entry_start: u64, entry_end: u64) -> Self {
72        IndexMemo {
73            read_memo,
74            entry_start,
75            entry_end,
76        }
77    }
78}
79
80/// A fetched groupcompress block paired with the memo it was fetched for.
81///
82/// `_get_blocks` yields these in the order the read-memos were requested.
83pub struct FetchedBlock<F: FileRef = String> {
84    pub read_memo: ReadMemo<F>,
85    pub block: GroupCompressBlock,
86}
87
88/// What a [`GcIndex`] knows about one stored key.
89///
90/// Mirrors the Python `_GCGraphIndex.get_build_details` result (a
91/// `GCBuildDetails`): where the record's bytes live (`index_memo`) and the
92/// key's graph parents. groupcompress records are always whole entries
93/// inside a block, so there is no separate compression-parent field.
94#[derive(Debug, Clone, PartialEq, Eq)]
95pub struct GcBuildDetails<F: FileRef = String> {
96    /// Locates the record: its block and slice within the block.
97    pub index_memo: IndexMemo<F>,
98    /// The key's graph parents, or `None` if the index stores no graph.
99    pub parents: Option<Vec<GcKey>>,
100}
101
102/// The index half of a groupcompress store.
103///
104/// Resolves keys to build details and graph parents, and accepts new
105/// records. The pyo3 layer implements this by wrapping a Python
106/// `_GCGraphIndex`; pure-Rust callers implement it directly. Mirrors the
107/// `KnitIndex` trait.
108pub trait GcIndex {
109    /// Identifies which backing block-file a record lives in.
110    type F: FileRef;
111
112    /// Build details for `keys`. Missing keys are absent from the result.
113    fn get_build_details(
114        &self,
115        keys: &[GcKey],
116    ) -> Result<std::collections::HashMap<GcKey, GcBuildDetails<Self::F>>, crate::knit::KnitError>;
117
118    /// Graph parents for `keys`. Missing keys are absent from the result.
119    fn get_parent_map(
120        &self,
121        keys: &[GcKey],
122    ) -> Result<std::collections::HashMap<GcKey, Vec<GcKey>>, crate::knit::KnitError>;
123
124    /// All keys present in this index.
125    fn keys(&self) -> Result<Vec<GcKey>, crate::knit::KnitError>;
126
127    /// Whether this index stores graph parents.
128    fn has_graph(&self) -> bool;
129
130    /// Assert that a write is permitted, erroring otherwise.
131    fn check_write_ok(&self) -> Result<(), crate::knit::KnitError>;
132
133    /// Add records to the index.
134    ///
135    /// Each record is `(key, index_memo, parents)`. `random_id` skips the
136    /// duplicate check.
137    fn add_records(
138        &self,
139        records: &[(GcKey, IndexMemo<Self::F>, Option<Vec<GcKey>>)],
140        random_id: bool,
141    ) -> Result<(), crate::knit::KnitError>;
142}
143
144/// The raw-data half of a groupcompress store.
145///
146/// Fetches and appends whole compressed blocks. Mirrors `KnitAccess`.
147pub trait GcAccess {
148    /// Identifies which backing block-file a block lives in.
149    type F: FileRef;
150
151    /// Fetch the raw (compressed) bytes for each read-memo, in order.
152    fn get_raw_records(
153        &self,
154        memos: &[ReadMemo<Self::F>],
155    ) -> Result<Vec<Vec<u8>>, crate::knit::KnitError>;
156
157    /// Append a compressed block and return the read-memo locating it.
158    fn add_raw_record(
159        &self,
160        size: usize,
161        chunks: Vec<Vec<u8>>,
162    ) -> Result<ReadMemo<Self::F>, crate::knit::KnitError>;
163
164    /// Call the reload hook after a stale-pack error, or re-raise it.
165    ///
166    /// Returns `Ok(())` if the caller should retry, `Err` if unrecoverable.
167    fn reload_or_raise(&self, err: crate::knit::KnitError) -> Result<(), crate::knit::KnitError> {
168        Err(err)
169    }
170}
171
172/// A shared, mutable reference to a decoded block.
173///
174/// `Arc<Mutex<..>>` rather than `Rc<RefCell<..>>` so the cache and the store
175/// are `Send + Sync`-compatible — the pyo3 layer can hold the store in a
176/// regular `#[pyclass]` without resorting to `unsendable`.
177pub type SharedBlock = std::sync::Arc<std::sync::Mutex<GroupCompressBlock>>;
178
179/// The decoded-block cache the store consults before fetching.
180///
181/// Decoupled into a trait so the pyo3 layer can back it with the Python
182/// `LRUSizeCache` (size-bounded) while pure-Rust callers use a plain map.
183pub trait BlockCache<F: FileRef>: Send + Sync {
184    /// Fetch a cached block by read-memo.
185    fn get(&self, memo: &ReadMemo<F>) -> Option<SharedBlock>;
186
187    /// Store a freshly decoded block.
188    fn insert(&self, memo: ReadMemo<F>, block: SharedBlock);
189
190    /// Whether a block is in the cache.
191    fn contains(&self, memo: &ReadMemo<F>) -> bool {
192        self.get(memo).is_some()
193    }
194
195    /// Drop every cached block.
196    fn clear(&self);
197
198    /// Number of cached blocks.
199    fn len(&self) -> usize;
200
201    /// Whether the cache currently holds nothing.
202    fn is_empty(&self) -> bool {
203        self.len() == 0
204    }
205}
206
207/// Unbounded in-memory `BlockCache` for pure-Rust callers.
208pub struct MapBlockCache<F: FileRef> {
209    inner: std::sync::Mutex<std::collections::HashMap<ReadMemo<F>, SharedBlock>>,
210}
211
212impl<F: FileRef> Default for MapBlockCache<F> {
213    fn default() -> Self {
214        MapBlockCache {
215            inner: std::sync::Mutex::new(std::collections::HashMap::new()),
216        }
217    }
218}
219
220impl<F: FileRef> BlockCache<F> for MapBlockCache<F> {
221    fn get(&self, memo: &ReadMemo<F>) -> Option<SharedBlock> {
222        self.inner.lock().unwrap().get(memo).cloned()
223    }
224    fn insert(&self, memo: ReadMemo<F>, block: SharedBlock) {
225        self.inner.lock().unwrap().insert(memo, block);
226    }
227    fn clear(&self) {
228        self.inner.lock().unwrap().clear();
229    }
230    fn len(&self) -> usize {
231        self.inner.lock().unwrap().len()
232    }
233}
234
235/// Given the read-memos a batch wants and which of them are already
236/// cached, return the de-duplicated, order-preserving list of memos that
237/// still need to be fetched.
238///
239/// Mirrors the partitioning loop in `GroupCompressVersionedFiles._get_blocks`:
240/// a memo that is cached is skipped, and a memo already queued for fetch is
241/// not queued twice. The first-seen request order is preserved so the
242/// fetched raw records line up with the consume order. `is_cached` lets the
243/// caller plug in whatever block cache it holds.
244pub fn memos_to_fetch<F: FileRef>(
245    read_memos: &[ReadMemo<F>],
246    is_cached: impl Fn(&ReadMemo<F>) -> bool,
247) -> Vec<ReadMemo<F>> {
248    let mut out: Vec<ReadMemo<F>> = Vec::new();
249    let mut seen: std::collections::HashSet<ReadMemo<F>> = std::collections::HashSet::new();
250    for memo in read_memos {
251        if is_cached(memo) {
252            continue;
253        }
254        if seen.insert(memo.clone()) {
255            out.push(memo.clone());
256        }
257    }
258    out
259}
260
261/// Which store a record-stream key is served from.
262///
263/// The Python code carries the `GroupCompressVersionedFiles` object itself
264/// (`self`) or a fallback VF object. The pure crate cannot hold Python
265/// objects, so a fallback is identified by its index into the ordered
266/// fallback list.
267#[derive(Debug, Clone, Copy, PartialEq, Eq)]
268pub enum Source {
269    /// This versioned-files store (the Python `self`).
270    Local,
271    /// The fallback at this index in the immediate-fallback list.
272    Fallback(usize),
273}
274
275/// Group an ordered key sequence into `(source, [keys])` runs.
276///
277/// `source_of` maps each key to its [`Source`]; consecutive keys from the
278/// same source are collected into one run. Mirrors the "Now group by
279/// source" loops shared by the three Python ordering helpers.
280fn group_by_source(
281    keys: impl IntoIterator<Item = GcKey>,
282    source_of: impl Fn(&GcKey) -> Source,
283) -> Vec<(Source, Vec<GcKey>)> {
284    let mut runs: Vec<(Source, Vec<GcKey>)> = Vec::new();
285    for key in keys {
286        let source = source_of(&key);
287        match runs.last_mut() {
288            Some((s, run)) if *s == source => run.push(key),
289            _ => runs.push((source, vec![key])),
290        }
291    }
292    runs
293}
294
295/// Order keys topologically (or in groupcompress order) and group by source.
296///
297/// Mirrors `GroupCompressVersionedFiles._get_ordered_source_keys`. `ordering`
298/// is `"topological"` or `"groupcompress"`; any key absent from
299/// `key_to_source` is served locally.
300pub fn ordered_source_keys(
301    ordering: &str,
302    parent_map: &[(GcKey, Vec<GcKey>)],
303    key_to_source: &std::collections::HashMap<GcKey, Source>,
304) -> Vec<(Source, Vec<GcKey>)> {
305    let raw: Vec<(Vec<Vec<u8>>, Vec<Vec<Vec<u8>>>)> = parent_map
306        .iter()
307        .map(|(k, ps)| {
308            (
309                k.segments().to_vec(),
310                ps.iter().map(|p| p.segments().to_vec()).collect(),
311            )
312        })
313        .collect();
314    let present: Vec<Vec<Vec<u8>>> = if ordering == "topological" {
315        let mut sorter = vcs_graph::tsort::TopoSorter::new(raw.into_iter());
316        sorter
317            .sorted()
318            .expect("groupcompress parent_map should not contain cycles")
319    } else {
320        crate::groupcompress::sort::sort_gc_optimal(raw)
321    };
322    let keys = present.into_iter().map(GcKey::fixed);
323    group_by_source(keys, |k| {
324        key_to_source.get(k).copied().unwrap_or(Source::Local)
325    })
326}
327
328/// Keep the caller's requested order, grouping by source and dropping keys
329/// that are absent from every store.
330///
331/// Mirrors `GroupCompressVersionedFiles._get_as_requested_source_keys`. A
332/// key present in `locations` or `unadded` is local; otherwise its
333/// `key_to_source` entry is used; a key in none of them is skipped.
334pub fn as_requested_source_keys(
335    orig_keys: &[GcKey],
336    locations: &std::collections::HashSet<GcKey>,
337    unadded: &std::collections::HashSet<GcKey>,
338    key_to_source: &std::collections::HashMap<GcKey, Source>,
339) -> Vec<(Source, Vec<GcKey>)> {
340    let present: Vec<GcKey> = orig_keys
341        .iter()
342        .filter(|k| {
343            locations.contains(*k) || unadded.contains(*k) || key_to_source.contains_key(*k)
344        })
345        .cloned()
346        .collect();
347    group_by_source(present, |k| {
348        if locations.contains(k) || unadded.contains(k) {
349            Source::Local
350        } else {
351            key_to_source.get(k).copied().unwrap_or(Source::Local)
352        }
353    })
354}
355
356/// Accumulates keys into a fetch batch, tracking the read-memos that batch
357/// touches and a running byte estimate.
358///
359/// Ports the state that `_BatchingBlockFetcher.add_key` maintains. The
360/// block cache is a Python `LRUSizeCache`, so the cache lookup is passed in
361/// as a predicate; the actual fetch happens later in the pyo3 layer using
362/// [`Self::memos_to_get`].
363#[derive(Debug, Default)]
364pub struct BatchAccumulator<F: FileRef = String> {
365    keys: Vec<GcKey>,
366    /// Read-memos seen in this batch, in first-seen order.
367    batch_memos: Vec<ReadMemo<F>>,
368    /// Read-memos in this batch that were not cached and must be fetched.
369    memos_to_get: Vec<ReadMemo<F>>,
370    total_bytes: u64,
371}
372
373impl<F: FileRef> BatchAccumulator<F> {
374    pub fn new() -> Self {
375        BatchAccumulator {
376            keys: Vec::new(),
377            batch_memos: Vec::new(),
378            memos_to_get: Vec::new(),
379            total_bytes: 0,
380        }
381    }
382
383    /// Add a key to the current batch and return the running byte estimate.
384    ///
385    /// `read_memo` is the key's block memo (`index_memo[0:3]`). `is_cached`
386    /// reports whether that block is already in the Python block cache.
387    /// Mirrors `_BatchingBlockFetcher.add_key`: a memo already in the batch
388    /// is not re-counted; a new uncached memo is queued for fetch and its
389    /// `stop` offset is added to the estimate (Python adds `read_memo[2]`,
390    /// the absolute stop, not the byte length — preserved here so the
391    /// `BATCH_SIZE` threshold behaves identically).
392    pub fn add_key(
393        &mut self,
394        key: GcKey,
395        read_memo: ReadMemo<F>,
396        is_cached: impl Fn(&ReadMemo<F>) -> bool,
397    ) -> u64 {
398        self.keys.push(key);
399        if self.batch_memos.contains(&read_memo) {
400            return self.total_bytes;
401        }
402        if !is_cached(&read_memo) {
403            self.total_bytes += read_memo.stop;
404            self.memos_to_get.push(read_memo.clone());
405        }
406        self.batch_memos.push(read_memo);
407        self.total_bytes
408    }
409
410    /// Keys added to this batch, in insertion order.
411    pub fn keys(&self) -> &[GcKey] {
412        &self.keys
413    }
414
415    /// Uncached read-memos this batch must fetch, in first-seen order.
416    pub fn memos_to_get(&self) -> &[ReadMemo<F>] {
417        &self.memos_to_get
418    }
419
420    /// Running byte estimate for the batch.
421    pub fn total_bytes(&self) -> u64 {
422        self.total_bytes
423    }
424
425    /// Clear all batch state, ready for the next batch.
426    pub fn reset(&mut self) {
427        self.keys.clear();
428        self.batch_memos.clear();
429        self.memos_to_get.clear();
430        self.total_bytes = 0;
431    }
432}
433
434/// Order keys for I/O efficiency: in-memory (unadded) keys first, then
435/// located keys grouped by the block they live in, then fallback runs.
436///
437/// Mirrors `GroupCompressVersionedFiles._get_io_ordered_source_keys`.
438/// `located_keys` is the located keys in the caller's order; each must have
439/// an entry in `locations`. They are stably sorted by their block index so
440/// keys in one group stay together while keeping their relative order, as
441/// Python's `sorted(locations, key=get_group)` does over an insertion-
442/// ordered dict. `fallback_runs` is the already-grouped `(source, keys)`
443/// list for keys served by fallbacks.
444pub fn io_ordered_source_keys<F: FileRef>(
445    located_keys: &[GcKey],
446    locations: &std::collections::HashMap<GcKey, IndexMemo<F>>,
447    unadded: &[GcKey],
448    fallback_runs: Vec<(Source, Vec<GcKey>)>,
449) -> Vec<(Source, Vec<GcKey>)> {
450    let mut local: Vec<GcKey> = unadded.to_vec();
451    let mut located: Vec<GcKey> = located_keys.to_vec();
452    // Python sorts located keys by the group object alone (index_memo[0]);
453    // the sort is stable, so keys within one group keep their relative order.
454    located.sort_by(|a, b| {
455        locations[a]
456            .read_memo
457            .index
458            .cmp(&locations[b].read_memo.index)
459    });
460    local.extend(located);
461    let mut runs = vec![(Source::Local, local)];
462    runs.extend(fallback_runs);
463    runs
464}
465
466/// A groupcompress versioned-file store.
467///
468/// The pure-Rust equivalent of Python's `GroupCompressVersionedFiles`,
469/// generic over a [`GcIndex`] and [`GcAccess`] exactly as
470/// [`crate::knit::KnitVersionedFiles`] is over `KnitIndex` / `KnitAccess`.
471/// A pure-Rust caller implements those two traits; the pyo3 layer wraps a
472/// Python index / access object.
473pub struct GroupCompressVersionedFiles<I, A, C = MapBlockCache<<I as GcIndex>::F>>
474where
475    I: GcIndex,
476    A: GcAccess<F = I::F>,
477    C: BlockCache<I::F>,
478{
479    index: I,
480    access: A,
481    /// Whether to delta-compress; carried for parity with the Python store.
482    delta: bool,
483    /// Decoded-block cache. Decoupled via [`BlockCache`] so the pyo3 layer
484    /// can back it with the Python `LRUSizeCache` (size-bounded) while
485    /// pure-Rust callers get an unbounded [`MapBlockCache`] by default.
486    block_cache: C,
487    /// Fallback stores consulted for keys absent from this one.
488    fallbacks: Vec<Box<dyn crate::versionedfile::VersionedFiles>>,
489}
490
491impl<I, A> GroupCompressVersionedFiles<I, A, MapBlockCache<I::F>>
492where
493    I: GcIndex,
494    A: GcAccess<F = I::F>,
495{
496    /// Create a store with the default in-memory block cache.
497    pub fn new(index: I, access: A, delta: bool) -> Self {
498        GroupCompressVersionedFiles::with_cache(index, access, delta, MapBlockCache::default())
499    }
500}
501
502impl<I, A, C> GroupCompressVersionedFiles<I, A, C>
503where
504    I: GcIndex,
505    A: GcAccess<F = I::F>,
506    C: BlockCache<I::F>,
507{
508    /// Create a store with a caller-supplied block cache.
509    pub fn with_cache(index: I, access: A, delta: bool, block_cache: C) -> Self {
510        GroupCompressVersionedFiles {
511            index,
512            access,
513            delta,
514            block_cache,
515            fallbacks: Vec::new(),
516        }
517    }
518
519    /// The block cache.
520    pub fn block_cache(&self) -> &C {
521        &self.block_cache
522    }
523
524    /// Add a fallback store for keys not present in this one.
525    ///
526    /// Mirrors `GroupCompressVersionedFiles.add_fallback_versioned_files`.
527    pub fn add_fallback_versioned_files(
528        &mut self,
529        fallback: Box<dyn crate::versionedfile::VersionedFiles>,
530    ) {
531        self.fallbacks.push(fallback);
532    }
533
534    /// The backing index.
535    pub fn index(&self) -> &I {
536        &self.index
537    }
538
539    /// The backing access object.
540    pub fn access(&self) -> &A {
541        &self.access
542    }
543
544    /// Whether this store delta-compresses.
545    pub fn delta(&self) -> bool {
546        self.delta
547    }
548
549    /// Fetch decoded blocks for `read_memos`, in request order.
550    ///
551    /// Mirrors `GroupCompressVersionedFiles._get_blocks`: cached blocks are
552    /// reused, uncached read-memos are de-duplicated and fetched in one
553    /// `get_raw_records` call, then decoded and cached.
554    #[allow(clippy::type_complexity)]
555    pub fn get_blocks(
556        &self,
557        read_memos: &[ReadMemo<I::F>],
558    ) -> Result<Vec<(ReadMemo<I::F>, SharedBlock)>, crate::knit::KnitError> {
559        let to_fetch = memos_to_fetch(read_memos, |m| self.block_cache.contains(m));
560        let raw = self.access.get_raw_records(&to_fetch)?;
561        if raw.len() != to_fetch.len() {
562            return Err(crate::knit::KnitError::Corrupt(
563                "get_raw_records returned the wrong number of records".to_string(),
564            ));
565        }
566        // Decode and cache the freshly fetched blocks.
567        for (memo, zdata) in to_fetch.iter().zip(raw) {
568            let block = GroupCompressBlock::from_bytes(&zdata[..])
569                .map_err(|e| crate::knit::KnitError::Corrupt(e.to_string()))?;
570            self.block_cache.insert(
571                memo.clone(),
572                std::sync::Arc::new(std::sync::Mutex::new(block)),
573            );
574        }
575        // Now every requested read-memo is in the cache; yield in order.
576        let mut out = Vec::with_capacity(read_memos.len());
577        for memo in read_memos {
578            let block = self.block_cache.get(memo).ok_or_else(|| {
579                crate::knit::KnitError::Corrupt("block missing after fetch".to_string())
580            })?;
581            out.push((memo.clone(), block));
582        }
583        Ok(out)
584    }
585
586    /// Graph parents for `keys`; absent keys are omitted.
587    ///
588    /// Mirrors `GroupCompressVersionedFiles.get_parent_map`: the local
589    /// index is consulted first, then each fallback in turn for keys still
590    /// missing.
591    pub fn get_parent_map(
592        &self,
593        keys: &[GcKey],
594    ) -> Result<std::collections::HashMap<GcKey, Vec<GcKey>>, crate::knit::KnitError> {
595        let mut result = self.index.get_parent_map(keys)?;
596        let mut missing: Vec<GcKey> = keys
597            .iter()
598            .filter(|k| !result.contains_key(*k))
599            .cloned()
600            .collect();
601        for fb in &self.fallbacks {
602            if missing.is_empty() {
603                break;
604            }
605            let found = fb.get_parent_map(&missing)?;
606            missing.retain(|k| !found.contains_key(k));
607            result.extend(found);
608        }
609        Ok(result)
610    }
611
612    /// All keys present in this store or any fallback.
613    pub fn keys(&self) -> Result<Vec<GcKey>, crate::knit::KnitError> {
614        let mut seen: std::collections::HashSet<GcKey> = self.index.keys()?.into_iter().collect();
615        for fb in &self.fallbacks {
616            for k in fb.keys()? {
617                seen.insert(k);
618            }
619        }
620        Ok(seen.into_iter().collect())
621    }
622
623    /// Get a stream of records for `keys`.
624    ///
625    /// Mirrors `GroupCompressVersionedFiles.get_record_stream` for a store
626    /// with no fallbacks: locate the keys, order them per `ordering`, fetch
627    /// the blocks they live in, and extract each record into a
628    /// `ChunkedContentFactory`. Keys absent from the index yield an
629    /// `AbsentContentFactory`. Records are returned eagerly as a `Vec`, as
630    /// the knit pure store does.
631    ///
632    /// `ordering` is `"unordered"`, `"topological"`, `"groupcompress"`, or
633    /// `"as-requested"`.
634    pub fn get_record_stream(
635        &self,
636        keys: &[GcKey],
637        ordering: &str,
638    ) -> Result<Vec<Box<dyn crate::versionedfile::ContentFactory>>, crate::knit::KnitError> {
639        let locations = self.index.get_build_details(keys)?;
640        let mut out: Vec<Box<dyn crate::versionedfile::ContentFactory>> = Vec::new();
641        // Keys not in the local index are tried against fallbacks; only keys
642        // absent from every store yield an AbsentContentFactory.
643        let mut nonlocal: Vec<GcKey> = keys
644            .iter()
645            .filter(|k| !locations.contains_key(*k))
646            .cloned()
647            .collect();
648        for fb in &self.fallbacks {
649            if nonlocal.is_empty() {
650                break;
651            }
652            let fb_records = fb.get_record_stream(&nonlocal, ordering, true)?;
653            let mut still_missing: Vec<GcKey> = Vec::new();
654            let mut found: std::collections::HashSet<GcKey> = std::collections::HashSet::new();
655            for rec in fb_records {
656                let rec = rec?;
657                if rec.storage_kind() == "absent" {
658                    continue;
659                }
660                found.insert(rec.key());
661                out.push(rec);
662            }
663            for k in nonlocal {
664                if !found.contains(&k) {
665                    still_missing.push(k);
666                }
667            }
668            nonlocal = still_missing;
669        }
670        for key in nonlocal {
671            out.push(Box::new(crate::versionedfile::AbsentContentFactory::new(
672                key,
673            )));
674        }
675        // Order the located keys.
676        let located: Vec<GcKey> = keys
677            .iter()
678            .filter(|k| locations.contains_key(*k))
679            .cloned()
680            .collect();
681        let ordered: Vec<GcKey> = match ordering {
682            "topological" | "groupcompress" => {
683                let parent_map: Vec<(GcKey, Vec<GcKey>)> = located
684                    .iter()
685                    .map(|k| (k.clone(), locations[k].parents.clone().unwrap_or_default()))
686                    .collect();
687                let empty = std::collections::HashMap::new();
688                ordered_source_keys(ordering, &parent_map, &empty)
689                    .into_iter()
690                    .flat_map(|(_, ks)| ks)
691                    .collect()
692            }
693            "as-requested" => located,
694            // "unordered" and anything else: I/O order, grouped by block.
695            _ => {
696                let mut io = located;
697                io.sort_by(|a, b| {
698                    let ka = &locations[a].index_memo;
699                    let kb = &locations[b].index_memo;
700                    (
701                        &ka.read_memo.index,
702                        ka.read_memo.start,
703                        ka.read_memo.stop,
704                        ka.entry_start,
705                        ka.entry_end,
706                    )
707                        .cmp(&(
708                            &kb.read_memo.index,
709                            kb.read_memo.start,
710                            kb.read_memo.stop,
711                            kb.entry_start,
712                            kb.entry_end,
713                        ))
714                });
715                io
716            }
717        };
718        // Fetch every block the ordered keys touch, then extract records.
719        let read_memos: Vec<ReadMemo<I::F>> = ordered
720            .iter()
721            .map(|k| locations[k].index_memo.read_memo.clone())
722            .collect();
723        let blocks = self.get_blocks(&read_memos)?;
724        let block_of: std::collections::HashMap<ReadMemo<I::F>, SharedBlock> =
725            blocks.into_iter().collect();
726        for key in &ordered {
727            let memo = &locations[key].index_memo;
728            let block = block_of.get(&memo.read_memo).ok_or_else(|| {
729                crate::knit::KnitError::Corrupt("block missing for located key".to_string())
730            })?;
731            let chunks = block
732                .lock()
733                .unwrap()
734                .extract(memo.entry_start as usize, memo.entry_end as usize)
735                .map_err(|e| crate::knit::KnitError::Corrupt(format!("{:?}", e)))?;
736            out.push(Box::new(crate::versionedfile::ChunkedContentFactory::new(
737                None,
738                key.clone(),
739                locations[key].parents.clone(),
740                chunks,
741            )));
742        }
743        Ok(out)
744    }
745
746    /// Insert a stream of records into this store.
747    ///
748    /// Mirrors `GroupCompressVersionedFiles._insert_record_stream` for the
749    /// ordinary (non-block-reuse) path: each record's text is compressed
750    /// into a `RabinGroupCompressor`; when the start-new-block heuristic
751    /// fires the current block is flushed (written via `GcAccess` and
752    /// indexed via `GcIndex`) and a fresh compressor started. Returns the
753    /// `(sha1, length)` of each inserted record.
754    ///
755    /// `nostore_sha`, when set on a record's content, causes
756    /// `KnitError::ExistingContent` if the text is already stored.
757    pub fn insert_record_stream(
758        &self,
759        stream: impl IntoIterator<Item = Box<dyn crate::versionedfile::ContentFactory>>,
760        random_id: bool,
761    ) -> Result<Vec<(Vec<u8>, usize)>, crate::knit::KnitError> {
762        use crate::groupcompress::compressor::{GroupCompressor, RabinGroupCompressor};
763
764        self.index.check_write_ok()?;
765        let mut results: Vec<(Vec<u8>, usize)> = Vec::new();
766        let mut compressor = RabinGroupCompressor::new(None);
767        // Records compressed into the block not yet flushed: their key, the
768        // (entry_start, entry_end) within the block, and graph parents.
769        let mut pending: Vec<(GcKey, usize, usize, Option<Vec<GcKey>>)> = Vec::new();
770        let mut inserted: std::collections::HashSet<GcKey> = std::collections::HashSet::new();
771
772        let mut last_prefix: Option<Vec<u8>> = None;
773        let mut max_fulltext_len: usize = 0;
774        let mut max_fulltext_prefix: Option<Vec<u8>> = None;
775
776        for record in stream {
777            let key = record.key();
778            if record.storage_kind() == "absent" {
779                return Err(crate::knit::KnitError::RevisionNotPresent(
780                    key.segments().to_vec(),
781                ));
782            }
783            if random_id && !inserted.insert(key.clone()) {
784                // Same key offered twice under random_id: skip the dup.
785                continue;
786            }
787            // The record's text, as chunks.
788            let chunks: Vec<Vec<u8>> = record.to_chunks().map(|c| c.into_owned()).collect();
789            let chunks_len: usize = record
790                .size()
791                .unwrap_or_else(|| chunks.iter().map(|c| c.len()).sum());
792            let chunk_refs: Vec<&[u8]> = chunks.iter().map(|c| c.as_slice()).collect();
793
794            // The prefix is the key's first segment for multi-segment keys.
795            let prefix: Option<Vec<u8>> = if key.segments().len() > 1 {
796                Some(key.segments()[0].clone())
797            } else {
798                None
799            };
800            let soft = prefix.is_some() && prefix == last_prefix;
801            if max_fulltext_len < chunks_len {
802                max_fulltext_len = chunks_len;
803                max_fulltext_prefix = prefix.clone();
804            }
805            let (mut sha1, mut start_point, mut end_point) = compress_record(
806                &mut compressor,
807                &key,
808                &chunk_refs,
809                chunks_len,
810                record.sha1(),
811                soft,
812            )?;
813            // Start-new-block heuristic (mirrors the Python conditions).
814            let same_prefix = prefix == max_fulltext_prefix;
815            let start_new_block = if same_prefix && end_point < 2 * max_fulltext_len {
816                false
817            } else if end_point > 4 * 1024 * 1024 {
818                true
819            } else {
820                prefix.is_some() && prefix != last_prefix && end_point > 2 * 1024 * 1024
821            };
822            last_prefix = prefix;
823            if start_new_block {
824                let (content_chunks, content_len) = compressor.flush_without_last();
825                self.flush_block(content_chunks, content_len, &mut pending, random_id)?;
826                compressor = RabinGroupCompressor::new(None);
827                max_fulltext_len = chunks_len;
828                let recompressed = compress_record(
829                    &mut compressor,
830                    &key,
831                    &chunk_refs,
832                    chunks_len,
833                    record.sha1(),
834                    false,
835                )?;
836                sha1 = recompressed.0;
837                start_point = recompressed.1;
838                end_point = recompressed.2;
839            }
840            // A content-addressed key (None version id) gets the sha1 filled in.
841            let stored_key = if key.version_id().is_empty() {
842                GcKey::from_prefix_and_suffix(key.prefix(), {
843                    let mut seg = b"sha1:".to_vec();
844                    seg.extend_from_slice(&sha1);
845                    seg
846                })
847            } else {
848                key.clone()
849            };
850            results.push((sha1, chunks_len));
851            pending.push((stored_key, start_point, end_point, record.parents()));
852        }
853        if !pending.is_empty() {
854            let (content_chunks, content_len) = compressor.flush();
855            self.flush_block(content_chunks, content_len, &mut pending, random_id)?;
856        }
857        Ok(results)
858    }
859
860    /// Wrap the compressor's flushed content into a block, write it via the
861    /// access object, and index every pending record against it.
862    ///
863    /// Mirrors the `flush` closure inside `_insert_record_stream`.
864    fn flush_block(
865        &self,
866        content_chunks: Vec<Vec<u8>>,
867        content_len: usize,
868        pending: &mut Vec<(GcKey, usize, usize, Option<Vec<GcKey>>)>,
869        random_id: bool,
870    ) -> Result<(), crate::knit::KnitError> {
871        let mut block = GroupCompressBlock::new();
872        block.set_chunked_content(&content_chunks, content_len);
873        let on_disk = block.to_bytes();
874        let size = on_disk.len();
875        let read_memo = self.access.add_raw_record(size, vec![on_disk])?;
876        let records: Vec<(GcKey, IndexMemo<I::F>, Option<Vec<GcKey>>)> = pending
877            .drain(..)
878            .map(|(key, start, end, parents)| {
879                (
880                    key,
881                    IndexMemo::new(read_memo.clone(), start as u64, end as u64),
882                    parents,
883                )
884            })
885            .collect();
886        self.index.add_records(&records, random_id)?;
887        Ok(())
888    }
889
890    /// Add one text given as a list of lines.
891    ///
892    /// Mirrors `GroupCompressVersionedFiles.add_lines` / `add_content`:
893    /// wraps the lines in a content factory and inserts it. Returns the
894    /// text's `(sha1, length)`.
895    pub fn add_lines(
896        &self,
897        key: GcKey,
898        parents: Option<Vec<GcKey>>,
899        lines: Vec<Vec<u8>>,
900    ) -> Result<(Vec<u8>, usize), crate::knit::KnitError> {
901        // _check_add: a version id with embedded whitespace is rejected.
902        let version_id = key.version_id();
903        if version_id
904            .iter()
905            .any(|b| matches!(b, b' ' | b'\t' | b'\n' | b'\r' | 0x0b | 0x0c))
906        {
907            return Err(crate::knit::KnitError::Corrupt(format!(
908                "invalid revision id: {:?}",
909                version_id
910            )));
911        }
912        let sha1 = crate::weave::sha_strings(&lines);
913        let factory: Box<dyn crate::versionedfile::ContentFactory> = Box::new(
914            crate::versionedfile::ChunkedContentFactory::new(Some(sha1), key, parents, lines),
915        );
916        let mut inserted = self.insert_record_stream(std::iter::once(factory), false)?;
917        inserted
918            .pop()
919            .ok_or_else(|| crate::knit::KnitError::Corrupt("add_lines inserted nothing".into()))
920    }
921
922    /// SHA-1 of every requested key; absent keys are omitted.
923    ///
924    /// Mirrors `GroupCompressVersionedFiles.get_sha1s`.
925    pub fn get_sha1s(
926        &self,
927        keys: &[GcKey],
928    ) -> Result<std::collections::HashMap<GcKey, Vec<u8>>, crate::knit::KnitError> {
929        let mut out = std::collections::HashMap::new();
930        for record in self.get_record_stream(keys, "unordered")? {
931            if record.storage_kind() == "absent" {
932                continue;
933            }
934            let digest = match record.sha1() {
935                Some(s) => s,
936                None => {
937                    let chunks: Vec<Vec<u8>> = record.to_chunks().map(|c| c.into_owned()).collect();
938                    crate::weave::sha_strings(&chunks)
939                }
940            };
941            out.insert(record.key(), digest);
942        }
943        Ok(out)
944    }
945
946    /// Keys of missing compression parents.
947    ///
948    /// Mirrors `get_missing_compression_parent_keys`: groupcompress cannot
949    /// reference texts outside the group, so this is always empty.
950    pub fn get_missing_compression_parent_keys(&self) -> Vec<GcKey> {
951        Vec::new()
952    }
953
954    /// Drop the decoded-block cache.
955    ///
956    /// Mirrors `GroupCompressVersionedFiles.clear_cache`.
957    pub fn clear_cache(&self) {
958        self.block_cache.clear();
959    }
960
961    /// Walk the lines of every requested key.
962    ///
963    /// Mirrors `iter_lines_added_or_present_in_keys`: each key's text is
964    /// read and split into lines, returned as `(line, key)` pairs. A key
965    /// absent from every store is an error.
966    /// Yield `(line, key)` for every line added by or present in
967    /// `keys`. The record stream is fetched up front (the index lookup
968    /// is inherently batched), but each record is only decoded into
969    /// lines when the iterator reaches it. Yields `Err` for an absent
970    /// key or a line that fails to decode.
971    pub fn iter_lines_added_or_present_in_keys(
972        &self,
973        keys: &[GcKey],
974    ) -> Result<
975        impl Iterator<Item = Result<(Vec<u8>, GcKey), crate::knit::KnitError>>,
976        crate::knit::KnitError,
977    > {
978        let records = self.get_record_stream(keys, "unordered")?;
979        Ok(records.into_iter().flat_map(|record| {
980            if record.storage_kind() == "absent" {
981                let err =
982                    crate::knit::KnitError::RevisionNotPresent(record.key().segments().to_vec());
983                return vec![Err(err)].into_iter();
984            }
985            let key = record.key();
986            let chunks: Vec<Vec<u8>> = record.to_chunks().map(|c| c.into_owned()).collect();
987            crate::osutils::chunks_to_lines(chunks.into_iter().map(Ok::<_, std::io::Error>))
988                .map(|line| {
989                    line.map(|l| (l.into_owned(), key.clone()))
990                        .map_err(|e| crate::knit::KnitError::Corrupt(e.to_string()))
991                })
992                .collect::<Vec<_>>()
993                .into_iter()
994        }))
995    }
996
997    /// Check the store reads back: every key's text is fetched and decoded.
998    ///
999    /// Mirrors `GroupCompressVersionedFiles.check` with `keys=None`.
1000    pub fn check(&self) -> Result<(), crate::knit::KnitError> {
1001        let all_keys = self.keys()?;
1002        for record in self.get_record_stream(&all_keys, "unordered")? {
1003            if record.storage_kind() == "absent" {
1004                return Err(crate::knit::KnitError::RevisionNotPresent(
1005                    record.key().segments().to_vec(),
1006                ));
1007            }
1008            // Force decoding of the content.
1009            let _ = record.to_fulltext();
1010        }
1011        Ok(())
1012    }
1013}
1014
1015/// `GroupCompressVersionedFiles` is a full [`VersionedFiles`] backend, so
1016/// it can be used anywhere the trait is expected (as a CHK store for
1017/// `CHKInventory`, as a fallback store, etc.). The trait methods delegate
1018/// to the inherent ones, adapting the few signatures that differ (the
1019/// trait borrows keys/lines and boxes the record stream).
1020impl<I, A, C> crate::versionedfile::VersionedFiles for GroupCompressVersionedFiles<I, A, C>
1021where
1022    I: GcIndex + Send + Sync,
1023    A: GcAccess<F = I::F> + Send + Sync,
1024    C: BlockCache<I::F> + Send + Sync,
1025{
1026    fn get_parent_map(
1027        &self,
1028        keys: &[GcKey],
1029    ) -> Result<std::collections::HashMap<GcKey, Vec<GcKey>>, crate::knit::KnitError> {
1030        GroupCompressVersionedFiles::get_parent_map(self, keys)
1031    }
1032
1033    fn get_record_stream(
1034        &self,
1035        keys: &[GcKey],
1036        ordering: &str,
1037        _include_delta_closure: bool,
1038    ) -> Result<
1039        Box<
1040            dyn Iterator<
1041                Item = Result<
1042                    Box<dyn crate::versionedfile::ContentFactory>,
1043                    crate::knit::KnitError,
1044                >,
1045            >,
1046        >,
1047        crate::knit::KnitError,
1048    > {
1049        let records = GroupCompressVersionedFiles::get_record_stream(self, keys, ordering)?;
1050        Ok(Box::new(records.into_iter().map(Ok)))
1051    }
1052
1053    fn get_sha1s(
1054        &self,
1055        keys: &[GcKey],
1056    ) -> Result<std::collections::HashMap<GcKey, Vec<u8>>, crate::knit::KnitError> {
1057        GroupCompressVersionedFiles::get_sha1s(self, keys)
1058    }
1059
1060    fn keys(&self) -> Result<Vec<GcKey>, crate::knit::KnitError> {
1061        GroupCompressVersionedFiles::keys(self)
1062    }
1063
1064    fn add_lines(
1065        &self,
1066        key: &GcKey,
1067        parents: Option<&[GcKey]>,
1068        lines: &[Vec<u8>],
1069    ) -> Result<(Vec<u8>, usize), crate::knit::KnitError> {
1070        GroupCompressVersionedFiles::add_lines(
1071            self,
1072            key.clone(),
1073            parents.map(|p| p.to_vec()),
1074            lines.to_vec(),
1075        )
1076    }
1077
1078    fn insert_record_stream(
1079        &self,
1080        stream: Box<dyn Iterator<Item = Box<dyn crate::versionedfile::ContentFactory>>>,
1081    ) -> Result<(), crate::knit::KnitError> {
1082        GroupCompressVersionedFiles::insert_record_stream(self, stream, false)?;
1083        Ok(())
1084    }
1085
1086    fn iter_lines_added_or_present_in_keys(
1087        &self,
1088        keys: &[GcKey],
1089    ) -> Result<Vec<(Vec<u8>, GcKey)>, crate::knit::KnitError> {
1090        GroupCompressVersionedFiles::iter_lines_added_or_present_in_keys(self, keys)?
1091            .collect::<Result<Vec<_>, _>>()
1092    }
1093
1094    fn annotate(&self, _key: &GcKey) -> Result<Vec<(GcKey, Vec<u8>)>, crate::knit::KnitError> {
1095        Err(crate::knit::KnitError::NotImplemented(
1096            "groupcompress annotate",
1097        ))
1098    }
1099
1100    fn get_missing_compression_parent_keys(&self) -> Result<Vec<GcKey>, crate::knit::KnitError> {
1101        Ok(GroupCompressVersionedFiles::get_missing_compression_parent_keys(self))
1102    }
1103
1104    fn clear_cache(&self) {
1105        GroupCompressVersionedFiles::clear_cache(self)
1106    }
1107
1108    fn check(&self) -> Result<(), crate::knit::KnitError> {
1109        GroupCompressVersionedFiles::check(self)
1110    }
1111}
1112
1113/// Compress one record into `compressor`, mapping the result to plain
1114/// types and `nostore_sha` rejection to `KnitError::ExistingContent`.
1115fn compress_record(
1116    compressor: &mut crate::groupcompress::compressor::RabinGroupCompressor,
1117    key: &GcKey,
1118    chunks: &[&[u8]],
1119    length: usize,
1120    expected_sha: Option<Vec<u8>>,
1121    soft: bool,
1122) -> Result<(Vec<u8>, usize, usize), crate::knit::KnitError> {
1123    use crate::groupcompress::compressor::GroupCompressor;
1124    let expected = expected_sha
1125        .map(|s| String::from_utf8(s))
1126        .transpose()
1127        .map_err(|e| crate::knit::KnitError::Corrupt(e.to_string()))?;
1128    match compressor.compress(key, chunks, length, expected, None, Some(soft)) {
1129        Ok((sha, start, end, _kind)) => Ok((sha.into_bytes(), start, end)),
1130        Err(crate::versionedfile::Error::ExistingContent(_)) => {
1131            Err(crate::knit::KnitError::ExistingContent(Vec::new()))
1132        }
1133        Err(e) => Err(crate::knit::KnitError::Corrupt(e.to_string())),
1134    }
1135}
1136
1137#[cfg(test)]
1138mod tests {
1139    use super::*;
1140
1141    #[test]
1142    fn read_memo_byte_length_is_stop_minus_start() {
1143        let memo = ReadMemo::new("idx".to_string(), 100, 350);
1144        assert_eq!(memo.byte_length(), 250);
1145    }
1146
1147    #[test]
1148    fn read_memo_byte_length_saturates_when_stop_below_start() {
1149        // A corrupt memo with stop < start yields 0 rather than underflowing.
1150        let memo = ReadMemo::new("idx".to_string(), 500, 100);
1151        assert_eq!(memo.byte_length(), 0);
1152    }
1153
1154    #[test]
1155    fn read_memos_compare_by_all_three_fields() {
1156        let base = ReadMemo::new("a".to_string(), 0, 10);
1157        assert_eq!(base, ReadMemo::new("a".to_string(), 0, 10));
1158        assert_ne!(base, ReadMemo::new("b".to_string(), 0, 10));
1159        assert_ne!(base, ReadMemo::new("a".to_string(), 1, 10));
1160        assert_ne!(base, ReadMemo::new("a".to_string(), 0, 11));
1161    }
1162
1163    #[test]
1164    fn index_memo_carries_read_memo_and_entry_range() {
1165        let rm = ReadMemo::new("idx".to_string(), 0, 1000);
1166        let im = IndexMemo::new(rm.clone(), 40, 120);
1167        assert_eq!(im.read_memo, rm);
1168        assert_eq!(im.entry_start, 40);
1169        assert_eq!(im.entry_end, 120);
1170    }
1171
1172    fn memo(idx: &str, start: u64) -> ReadMemo<String> {
1173        ReadMemo::new(idx.to_string(), start, start + 10)
1174    }
1175
1176    #[test]
1177    fn memos_to_fetch_skips_cached_and_preserves_order() {
1178        let req = vec![memo("a", 0), memo("b", 0), memo("c", 0)];
1179        let cached = [memo("b", 0)];
1180        let out = memos_to_fetch(&req, |m| cached.contains(m));
1181        assert_eq!(out, vec![memo("a", 0), memo("c", 0)]);
1182    }
1183
1184    #[test]
1185    fn memos_to_fetch_dedups_repeated_memos() {
1186        // The same block requested twice is fetched once, in first-seen order.
1187        let req = vec![memo("a", 0), memo("b", 0), memo("a", 0), memo("c", 0)];
1188        let out = memos_to_fetch(&req, |_| false);
1189        assert_eq!(out, vec![memo("a", 0), memo("b", 0), memo("c", 0)]);
1190    }
1191
1192    #[test]
1193    fn memos_to_fetch_empty_when_all_cached() {
1194        let req = vec![memo("a", 0), memo("b", 0)];
1195        let out = memos_to_fetch(&req, |_| true);
1196        assert!(out.is_empty());
1197    }
1198
1199    fn gckey(id: &[u8]) -> GcKey {
1200        GcKey::fixed(vec![id.to_vec()])
1201    }
1202
1203    #[test]
1204    fn ordered_source_keys_topological_groups_by_source() {
1205        // Chain a -> b -> c; b is served by fallback 0, the rest locally.
1206        let a = gckey(b"a");
1207        let b = gckey(b"b");
1208        let c = gckey(b"c");
1209        let parent_map = vec![
1210            (a.clone(), vec![]),
1211            (b.clone(), vec![a.clone()]),
1212            (c.clone(), vec![b.clone()]),
1213        ];
1214        let mut k2s = std::collections::HashMap::new();
1215        k2s.insert(b.clone(), Source::Fallback(0));
1216        let runs = ordered_source_keys("topological", &parent_map, &k2s);
1217        assert_eq!(
1218            runs,
1219            vec![
1220                (Source::Local, vec![a]),
1221                (Source::Fallback(0), vec![b]),
1222                (Source::Local, vec![c]),
1223            ]
1224        );
1225    }
1226
1227    #[test]
1228    fn as_requested_source_keys_keeps_order_and_drops_absent() {
1229        let a = gckey(b"a");
1230        let b = gckey(b"b");
1231        let absent = gckey(b"absent");
1232        let f = gckey(b"f");
1233        let locations: std::collections::HashSet<GcKey> = vec![a.clone()].into_iter().collect();
1234        let unadded: std::collections::HashSet<GcKey> = vec![b.clone()].into_iter().collect();
1235        let mut k2s = std::collections::HashMap::new();
1236        k2s.insert(f.clone(), Source::Fallback(1));
1237        let runs = as_requested_source_keys(
1238            &[a.clone(), absent, b.clone(), f.clone()],
1239            &locations,
1240            &unadded,
1241            &k2s,
1242        );
1243        // `absent` is dropped; a and b are both local and merge into one run.
1244        assert_eq!(
1245            runs,
1246            vec![(Source::Local, vec![a, b]), (Source::Fallback(1), vec![f])]
1247        );
1248    }
1249
1250    #[test]
1251    fn io_ordered_source_keys_unadded_first_then_grouped_then_fallbacks() {
1252        let u = gckey(b"u");
1253        let x = gckey(b"x");
1254        let y = gckey(b"y");
1255        let f = gckey(b"f");
1256        let mut locations = std::collections::HashMap::new();
1257        // x in block "g2", y in block "g1" — sort pulls y ahead of x.
1258        locations.insert(
1259            x.clone(),
1260            IndexMemo::new(ReadMemo::new("g2".to_string(), 0, 10), 0, 5),
1261        );
1262        locations.insert(
1263            y.clone(),
1264            IndexMemo::new(ReadMemo::new("g1".to_string(), 0, 10), 0, 5),
1265        );
1266        let runs = io_ordered_source_keys(
1267            &[x.clone(), y.clone()],
1268            &locations,
1269            &[u.clone()],
1270            vec![(Source::Fallback(0), vec![f.clone()])],
1271        );
1272        assert_eq!(
1273            runs,
1274            vec![
1275                (Source::Local, vec![u, y, x]),
1276                (Source::Fallback(0), vec![f]),
1277            ]
1278        );
1279    }
1280
1281    #[test]
1282    fn batch_accumulator_queues_uncached_memos_and_counts_stop() {
1283        let mut acc: BatchAccumulator<String> = BatchAccumulator::new();
1284        // Two keys in distinct uncached blocks.
1285        let t1 = acc.add_key(gckey(b"k1"), ReadMemo::new("g1".into(), 0, 30), |_| false);
1286        assert_eq!(t1, 30); // running estimate adds `stop`
1287        let t2 = acc.add_key(gckey(b"k2"), ReadMemo::new("g2".into(), 0, 50), |_| false);
1288        assert_eq!(t2, 80);
1289        assert_eq!(acc.keys(), &[gckey(b"k1"), gckey(b"k2")]);
1290        assert_eq!(
1291            acc.memos_to_get(),
1292            &[
1293                ReadMemo::new("g1".into(), 0, 30),
1294                ReadMemo::new("g2".into(), 0, 50),
1295            ]
1296        );
1297    }
1298
1299    #[test]
1300    fn batch_accumulator_does_not_recount_repeated_memo() {
1301        let mut acc: BatchAccumulator<String> = BatchAccumulator::new();
1302        let block = ReadMemo::new("g1".to_string(), 0, 30);
1303        acc.add_key(gckey(b"k1"), block.clone(), |_| false);
1304        // A second key in the same block adds the key but not the bytes.
1305        let total = acc.add_key(gckey(b"k2"), block.clone(), |_| false);
1306        assert_eq!(total, 30);
1307        assert_eq!(acc.keys().len(), 2);
1308        assert_eq!(acc.memos_to_get(), &[block]);
1309    }
1310
1311    #[test]
1312    fn batch_accumulator_skips_fetch_for_cached_memo() {
1313        let mut acc: BatchAccumulator<String> = BatchAccumulator::new();
1314        let block = ReadMemo::new("g1".to_string(), 0, 30);
1315        // Cached blocks are not queued and not counted.
1316        let total = acc.add_key(gckey(b"k1"), block, |_| true);
1317        assert_eq!(total, 0);
1318        assert!(acc.memos_to_get().is_empty());
1319        assert_eq!(acc.keys().len(), 1);
1320    }
1321
1322    #[test]
1323    fn batch_accumulator_reset_clears_state() {
1324        let mut acc: BatchAccumulator<String> = BatchAccumulator::new();
1325        acc.add_key(gckey(b"k1"), ReadMemo::new("g1".into(), 0, 30), |_| false);
1326        acc.reset();
1327        assert!(acc.keys().is_empty());
1328        assert!(acc.memos_to_get().is_empty());
1329        assert_eq!(acc.total_bytes(), 0);
1330    }
1331
1332    /// Serialize a one-record fulltext block to its on-disk bytes.
1333    fn block_bytes(body: &[u8]) -> Vec<u8> {
1334        let mut b = GroupCompressBlock::new();
1335        b.set_content(body);
1336        b.to_bytes()
1337    }
1338
1339    /// A `GcAccess` that serves canned block bytes and counts fetches.
1340    struct MockAccess {
1341        blocks: std::collections::HashMap<ReadMemo<String>, Vec<u8>>,
1342        fetched: std::cell::RefCell<Vec<ReadMemo<String>>>,
1343    }
1344
1345    impl GcAccess for MockAccess {
1346        type F = String;
1347        fn get_raw_records(
1348            &self,
1349            memos: &[ReadMemo<String>],
1350        ) -> Result<Vec<Vec<u8>>, crate::knit::KnitError> {
1351            self.fetched.borrow_mut().extend_from_slice(memos);
1352            memos
1353                .iter()
1354                .map(|m| {
1355                    self.blocks
1356                        .get(m)
1357                        .cloned()
1358                        .ok_or_else(|| crate::knit::KnitError::Corrupt("no such memo".into()))
1359                })
1360                .collect()
1361        }
1362        fn add_raw_record(
1363            &self,
1364            _size: usize,
1365            _chunks: Vec<Vec<u8>>,
1366        ) -> Result<ReadMemo<String>, crate::knit::KnitError> {
1367            unimplemented!("not needed for get_blocks tests")
1368        }
1369    }
1370
1371    /// A `GcIndex` stub; `get_blocks` does not touch the index.
1372    struct UnusedIndex;
1373    impl GcIndex for UnusedIndex {
1374        type F = String;
1375        fn get_build_details(
1376            &self,
1377            _keys: &[GcKey],
1378        ) -> Result<std::collections::HashMap<GcKey, GcBuildDetails<String>>, crate::knit::KnitError>
1379        {
1380            unimplemented!()
1381        }
1382        fn get_parent_map(
1383            &self,
1384            _keys: &[GcKey],
1385        ) -> Result<std::collections::HashMap<GcKey, Vec<GcKey>>, crate::knit::KnitError> {
1386            unimplemented!()
1387        }
1388        fn keys(&self) -> Result<Vec<GcKey>, crate::knit::KnitError> {
1389            unimplemented!()
1390        }
1391        fn has_graph(&self) -> bool {
1392            true
1393        }
1394        fn check_write_ok(&self) -> Result<(), crate::knit::KnitError> {
1395            unimplemented!()
1396        }
1397        fn add_records(
1398            &self,
1399            _records: &[(GcKey, IndexMemo<String>, Option<Vec<GcKey>>)],
1400            _random_id: bool,
1401        ) -> Result<(), crate::knit::KnitError> {
1402            unimplemented!()
1403        }
1404    }
1405
1406    #[test]
1407    fn get_blocks_fetches_decodes_and_caches() {
1408        let m1 = ReadMemo::new("g".to_string(), 0, 10);
1409        let m2 = ReadMemo::new("g".to_string(), 10, 20);
1410        let mut blocks = std::collections::HashMap::new();
1411        blocks.insert(m1.clone(), block_bytes(b"first block body\n"));
1412        blocks.insert(m2.clone(), block_bytes(b"second block body\n"));
1413        let access = MockAccess {
1414            blocks,
1415            fetched: std::cell::RefCell::new(Vec::new()),
1416        };
1417        let vf = GroupCompressVersionedFiles::new(UnusedIndex, access, true);
1418
1419        // First call fetches both, in request order.
1420        let out = vf.get_blocks(&[m1.clone(), m2.clone()]).unwrap();
1421        assert_eq!(out.len(), 2);
1422        assert_eq!(out[0].0, m1);
1423        assert_eq!(out[1].0, m2);
1424        assert_eq!(vf.access().fetched.borrow().len(), 2);
1425
1426        // Second call is served entirely from the cache: no new fetch.
1427        let out2 = vf.get_blocks(&[m1.clone(), m2.clone()]).unwrap();
1428        assert_eq!(out2.len(), 2);
1429        assert_eq!(vf.access().fetched.borrow().len(), 2);
1430    }
1431
1432    #[test]
1433    fn get_blocks_dedups_repeated_memos_in_one_call() {
1434        let m = ReadMemo::new("g".to_string(), 0, 10);
1435        let mut blocks = std::collections::HashMap::new();
1436        blocks.insert(m.clone(), block_bytes(b"body\n"));
1437        let access = MockAccess {
1438            blocks,
1439            fetched: std::cell::RefCell::new(Vec::new()),
1440        };
1441        let vf = GroupCompressVersionedFiles::new(UnusedIndex, access, true);
1442        // The same block requested twice is fetched once, yielded twice.
1443        let out = vf.get_blocks(&[m.clone(), m.clone()]).unwrap();
1444        assert_eq!(out.len(), 2);
1445        assert_eq!(vf.access().fetched.borrow().len(), 1);
1446    }
1447
1448    /// A `GcIndex` that serves a fixed build-details map.
1449    struct MapIndex {
1450        details: std::collections::HashMap<GcKey, GcBuildDetails<String>>,
1451    }
1452    impl GcIndex for MapIndex {
1453        type F = String;
1454        fn get_build_details(
1455            &self,
1456            keys: &[GcKey],
1457        ) -> Result<std::collections::HashMap<GcKey, GcBuildDetails<String>>, crate::knit::KnitError>
1458        {
1459            Ok(keys
1460                .iter()
1461                .filter_map(|k| self.details.get(k).map(|d| (k.clone(), d.clone())))
1462                .collect())
1463        }
1464        fn get_parent_map(
1465            &self,
1466            keys: &[GcKey],
1467        ) -> Result<std::collections::HashMap<GcKey, Vec<GcKey>>, crate::knit::KnitError> {
1468            Ok(keys
1469                .iter()
1470                .filter_map(|k| {
1471                    self.details
1472                        .get(k)
1473                        .map(|d| (k.clone(), d.parents.clone().unwrap_or_default()))
1474                })
1475                .collect())
1476        }
1477        fn keys(&self) -> Result<Vec<GcKey>, crate::knit::KnitError> {
1478            Ok(self.details.keys().cloned().collect())
1479        }
1480        fn has_graph(&self) -> bool {
1481            true
1482        }
1483        fn check_write_ok(&self) -> Result<(), crate::knit::KnitError> {
1484            Ok(())
1485        }
1486        fn add_records(
1487            &self,
1488            _records: &[(GcKey, IndexMemo<String>, Option<Vec<GcKey>>)],
1489            _random_id: bool,
1490        ) -> Result<(), crate::knit::KnitError> {
1491            unimplemented!()
1492        }
1493    }
1494
1495    #[test]
1496    fn get_record_stream_extracts_records_from_a_block() {
1497        use crate::groupcompress::compressor::{GroupCompressor, RabinGroupCompressor};
1498
1499        // Compress two records into one block, capturing their positions.
1500        let mut gc = RabinGroupCompressor::new(None);
1501        let body_a = b"record a text line one\nrecord a text line two\n";
1502        let body_b = b"record b text line one\nrecord b text line two\n";
1503        let (_sha, a_start, a_end, _k) = gc
1504            .compress(
1505                &GcKey::fixed(vec![b"a".to_vec()]),
1506                &[body_a.as_slice()],
1507                body_a.len(),
1508                None,
1509                None,
1510                None,
1511            )
1512            .unwrap();
1513        let (_sha, b_start, b_end, _k) = gc
1514            .compress(
1515                &GcKey::fixed(vec![b"b".to_vec()]),
1516                &[body_b.as_slice()],
1517                body_b.len(),
1518                None,
1519                None,
1520                None,
1521            )
1522            .unwrap();
1523        // flush() returns the raw concatenated content; wrap it in a block.
1524        let (content_chunks, content_len) = gc.flush();
1525        let mut gcb = GroupCompressBlock::new();
1526        gcb.set_chunked_content(&content_chunks, content_len);
1527        let block: Vec<u8> = gcb.to_bytes();
1528
1529        let rm = ReadMemo::new("g".to_string(), 0, block.len() as u64);
1530        let key_a = gckey(b"a");
1531        let key_b = gckey(b"b");
1532        let mut details = std::collections::HashMap::new();
1533        details.insert(
1534            key_a.clone(),
1535            GcBuildDetails {
1536                index_memo: IndexMemo::new(rm.clone(), a_start as u64, a_end as u64),
1537                parents: Some(vec![]),
1538            },
1539        );
1540        details.insert(
1541            key_b.clone(),
1542            GcBuildDetails {
1543                index_memo: IndexMemo::new(rm.clone(), b_start as u64, b_end as u64),
1544                parents: Some(vec![key_a.clone()]),
1545            },
1546        );
1547        let mut blocks = std::collections::HashMap::new();
1548        blocks.insert(rm, block);
1549        let access = MockAccess {
1550            blocks,
1551            fetched: std::cell::RefCell::new(Vec::new()),
1552        };
1553        let vf = GroupCompressVersionedFiles::new(MapIndex { details }, access, true);
1554
1555        let records = vf
1556            .get_record_stream(&[key_a.clone(), key_b.clone()], "as-requested")
1557            .unwrap();
1558        assert_eq!(records.len(), 2);
1559        assert_eq!(records[0].key(), key_a);
1560        assert_eq!(records[0].to_fulltext().as_ref(), body_a.as_slice());
1561        assert_eq!(records[1].key(), key_b);
1562        assert_eq!(records[1].to_fulltext().as_ref(), body_b.as_slice());
1563        assert_eq!(records[1].parents(), Some(vec![key_a]));
1564    }
1565
1566    #[test]
1567    fn get_record_stream_yields_absent_for_missing_keys() {
1568        let access = MockAccess {
1569            blocks: std::collections::HashMap::new(),
1570            fetched: std::cell::RefCell::new(Vec::new()),
1571        };
1572        let vf = GroupCompressVersionedFiles::new(
1573            MapIndex {
1574                details: std::collections::HashMap::new(),
1575            },
1576            access,
1577            true,
1578        );
1579        let records = vf
1580            .get_record_stream(&[gckey(b"nope")], "unordered")
1581            .unwrap();
1582        assert_eq!(records.len(), 1);
1583        assert_eq!(records[0].storage_kind(), "absent");
1584        assert_eq!(records[0].key(), gckey(b"nope"));
1585    }
1586
1587    /// A writable in-memory groupcompress store shared by a `MemIndex` and
1588    /// a `MemAccess`, so an insert/read round-trip can be exercised purely.
1589    #[derive(Default)]
1590    struct MemStore {
1591        /// Appended blocks, in write order; the read-memo index is the
1592        /// block's position in this vec, rendered as a string.
1593        blocks: Vec<Vec<u8>>,
1594        details: std::collections::HashMap<GcKey, GcBuildDetails<String>>,
1595    }
1596
1597    #[derive(Clone)]
1598    struct MemIndex(std::rc::Rc<std::cell::RefCell<MemStore>>);
1599    #[derive(Clone)]
1600    struct MemAccess(std::rc::Rc<std::cell::RefCell<MemStore>>);
1601
1602    impl GcIndex for MemIndex {
1603        type F = String;
1604        fn get_build_details(
1605            &self,
1606            keys: &[GcKey],
1607        ) -> Result<std::collections::HashMap<GcKey, GcBuildDetails<String>>, crate::knit::KnitError>
1608        {
1609            let store = self.0.borrow();
1610            Ok(keys
1611                .iter()
1612                .filter_map(|k| store.details.get(k).map(|d| (k.clone(), d.clone())))
1613                .collect())
1614        }
1615        fn get_parent_map(
1616            &self,
1617            keys: &[GcKey],
1618        ) -> Result<std::collections::HashMap<GcKey, Vec<GcKey>>, crate::knit::KnitError> {
1619            let store = self.0.borrow();
1620            Ok(keys
1621                .iter()
1622                .filter_map(|k| {
1623                    store
1624                        .details
1625                        .get(k)
1626                        .map(|d| (k.clone(), d.parents.clone().unwrap_or_default()))
1627                })
1628                .collect())
1629        }
1630        fn keys(&self) -> Result<Vec<GcKey>, crate::knit::KnitError> {
1631            Ok(self.0.borrow().details.keys().cloned().collect())
1632        }
1633        fn has_graph(&self) -> bool {
1634            true
1635        }
1636        fn check_write_ok(&self) -> Result<(), crate::knit::KnitError> {
1637            Ok(())
1638        }
1639        fn add_records(
1640            &self,
1641            records: &[(GcKey, IndexMemo<String>, Option<Vec<GcKey>>)],
1642            _random_id: bool,
1643        ) -> Result<(), crate::knit::KnitError> {
1644            let mut store = self.0.borrow_mut();
1645            for (key, memo, parents) in records {
1646                store.details.insert(
1647                    key.clone(),
1648                    GcBuildDetails {
1649                        index_memo: memo.clone(),
1650                        parents: parents.clone(),
1651                    },
1652                );
1653            }
1654            Ok(())
1655        }
1656    }
1657
1658    impl GcAccess for MemAccess {
1659        type F = String;
1660        fn get_raw_records(
1661            &self,
1662            memos: &[ReadMemo<String>],
1663        ) -> Result<Vec<Vec<u8>>, crate::knit::KnitError> {
1664            let store = self.0.borrow();
1665            memos
1666                .iter()
1667                .map(|m| {
1668                    let idx: usize = m
1669                        .index
1670                        .parse()
1671                        .map_err(|_| crate::knit::KnitError::Corrupt("bad block index".into()))?;
1672                    store
1673                        .blocks
1674                        .get(idx)
1675                        .cloned()
1676                        .ok_or_else(|| crate::knit::KnitError::Corrupt("no such block".into()))
1677                })
1678                .collect()
1679        }
1680        fn add_raw_record(
1681            &self,
1682            _size: usize,
1683            chunks: Vec<Vec<u8>>,
1684        ) -> Result<ReadMemo<String>, crate::knit::KnitError> {
1685            let mut store = self.0.borrow_mut();
1686            let idx = store.blocks.len();
1687            let bytes: Vec<u8> = chunks.concat();
1688            let len = bytes.len() as u64;
1689            store.blocks.push(bytes);
1690            Ok(ReadMemo::new(idx.to_string(), 0, len))
1691        }
1692    }
1693
1694    #[test]
1695    fn insert_record_stream_then_get_record_stream_round_trips() {
1696        use crate::versionedfile::{ChunkedContentFactory, ContentFactory};
1697
1698        let store = std::rc::Rc::new(std::cell::RefCell::new(MemStore::default()));
1699        let vf = GroupCompressVersionedFiles::new(
1700            MemIndex(store.clone()),
1701            MemAccess(store.clone()),
1702            true,
1703        );
1704
1705        let key_a = gckey(b"a");
1706        let key_b = gckey(b"b");
1707        let text_a = b"the quick brown fox\njumps over\n".to_vec();
1708        let text_b = b"the lazy dog\nsleeps all day\n".to_vec();
1709        let stream: Vec<Box<dyn ContentFactory>> = vec![
1710            Box::new(ChunkedContentFactory::new(
1711                None,
1712                key_a.clone(),
1713                Some(vec![]),
1714                vec![text_a.clone()],
1715            )),
1716            Box::new(ChunkedContentFactory::new(
1717                None,
1718                key_b.clone(),
1719                Some(vec![key_a.clone()]),
1720                vec![text_b.clone()],
1721            )),
1722        ];
1723        let inserted = vf.insert_record_stream(stream, false).unwrap();
1724        assert_eq!(inserted.len(), 2);
1725
1726        // Read the records back: the text must round-trip exactly.
1727        let records = vf
1728            .get_record_stream(&[key_a.clone(), key_b.clone()], "as-requested")
1729            .unwrap();
1730        assert_eq!(records.len(), 2);
1731        assert_eq!(records[0].key(), key_a);
1732        assert_eq!(records[0].to_fulltext().as_ref(), text_a.as_slice());
1733        assert_eq!(records[1].key(), key_b);
1734        assert_eq!(records[1].to_fulltext().as_ref(), text_b.as_slice());
1735        assert_eq!(records[1].parents(), Some(vec![key_a]));
1736    }
1737
1738    #[test]
1739    fn add_lines_then_get_sha1s_and_iter_lines() {
1740        let store = std::rc::Rc::new(std::cell::RefCell::new(MemStore::default()));
1741        let vf = GroupCompressVersionedFiles::new(
1742            MemIndex(store.clone()),
1743            MemAccess(store.clone()),
1744            true,
1745        );
1746        let key = gckey(b"v1");
1747        let lines = vec![b"alpha line\n".to_vec(), b"beta line\n".to_vec()];
1748        let (sha1, length) = vf
1749            .add_lines(key.clone(), Some(vec![]), lines.clone())
1750            .unwrap();
1751        assert_eq!(sha1, crate::weave::sha_strings(&lines));
1752        assert_eq!(length, lines.iter().map(|l| l.len()).sum::<usize>());
1753
1754        // get_sha1s returns the same digest.
1755        let sha1s = vf.get_sha1s(&[key.clone()]).unwrap();
1756        assert_eq!(sha1s.get(&key), Some(&sha1));
1757
1758        // iter_lines yields each line paired with the key.
1759        let iter_lines = vf
1760            .iter_lines_added_or_present_in_keys(&[key.clone()])
1761            .unwrap()
1762            .collect::<Result<Vec<_>, _>>()
1763            .unwrap();
1764        assert_eq!(
1765            iter_lines,
1766            vec![
1767                (b"alpha line\n".to_vec(), key.clone()),
1768                (b"beta line\n".to_vec(), key.clone()),
1769            ]
1770        );
1771
1772        // check passes over a sound store.
1773        vf.check().unwrap();
1774    }
1775}