Skip to main content

zeppelin_embed/graph/
refine.rs

1//! Copy-on-write post-consolidation graph refinement (Task 19-M9).
2
3use std::collections::VecDeque;
4use std::path::{Path, PathBuf};
5
6use xxhash_rust::xxh3::{xxh3_64, xxh3_64_with_seed};
7
8use crate::fts::index::{Document as LexicalDocument, SegmentIndex};
9use crate::fts::sealed::SealedSegment;
10use crate::fts::tokenizer::Analyzer;
11use crate::graph::GraphParams;
12use crate::graph::block::{
13    GraphNodeBlockBuild, GraphNodeBlockInput, GraphNodeBlocks, GraphNodeError, GraphNodeLayout,
14    encode_node_blocks_with_refinement_passes,
15};
16use crate::graph::build::{GraphBuildError, robust_prune_rows, write_segment_with_encoded_graph};
17use crate::lifecycle::durability::{DurabilityPolicy, SyncRequirement};
18use crate::lifecycle::stats::{AccountedCounter, AllocationComponent};
19use crate::lifecycle::{QueryCancellation, QueryControl, SnapshotLease, Store, StoreError};
20use crate::meta::{AliveSet, ColumnStoreBuilder};
21use crate::quant::{est_dot_bit4, prepare_bit4_query};
22use crate::scan::ScanError;
23use crate::segment::layout::RegionKind;
24use crate::segment::reader::SegmentReader;
25use crate::segment::writer::{
26    SegmentBuild, SegmentDocumentVersions, SegmentFactors, SegmentPayloads, SegmentPostings,
27    SegmentStoredMetadata, SegmentStoredText, write_segment, write_segment_with_documents_payloads,
28};
29use crate::segment::{SegmentError, SegmentId, SegmentMeta};
30use crate::vfs::Vfs;
31
32const ENTRY_FLAG: u8 = 0b0000_0001;
33const HUB_FLAG: u8 = 0b0000_0010;
34const ENTRY_POINT_COUNT: usize = 4;
35const INTERMEDIATE_ID_SEED: u64 = 0x0019_4d39_0000_0001;
36const CHECKPOINT_MAGIC: [u8; 8] = *b"ZEREFCP1";
37const CHECKPOINT_VERSION: u16 = 1;
38const CHECKPOINT_HEADER_BYTES: usize = 104;
39const CHECKPOINT_CHECKSUM_BYTES: usize = 8;
40
41/// One deterministic post-consolidation graph refinement pass.
42#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
43#[repr(u8)]
44pub enum RefinementPass {
45    /// Remap dense row ids into deterministic graph-traversal order.
46    Renumber = 0,
47    /// Re-prune finished neighborhoods with the profile's refinement alpha.
48    AlphaReprune = 1,
49    /// Replace sampled entry seeds with the exact medoid and spread seeds.
50    SeedRefit = 2,
51    /// Order each adjacency list by Bit4 distance to its owner.
52    NeighborReorder = 3,
53    /// Stitch unreachable components and densify the resulting hubs.
54    ConnectivityRepair = 4,
55}
56
57impl RefinementPass {
58    /// Catalog order for independently scheduled and measured passes.
59    pub const ALL: [Self; 4] = [
60        Self::Renumber,
61        Self::AlphaReprune,
62        Self::SeedRefit,
63        Self::NeighborReorder,
64    ];
65
66    const fn bit(self) -> u16 {
67        1_u16 << self as u8
68    }
69
70    /// Returns the stable command-line and evidence label.
71    #[must_use]
72    pub const fn label(self) -> &'static str {
73        match self {
74            Self::Renumber => "renumber",
75            Self::AlphaReprune => "alpha-reprune",
76            Self::SeedRefit => "seed-refit",
77            Self::NeighborReorder => "neighbor-reorder",
78            Self::ConnectivityRepair => "connectivity-repair",
79        }
80    }
81
82    /// Parses one stable command-line pass label.
83    pub fn named(label: &str) -> Result<Self, RefinementError> {
84        match label {
85            "renumber" => Ok(Self::Renumber),
86            "alpha-reprune" => Ok(Self::AlphaReprune),
87            "seed-refit" => Ok(Self::SeedRefit),
88            "neighbor-reorder" => Ok(Self::NeighborReorder),
89            "connectivity-repair" => Ok(Self::ConnectivityRepair),
90            _ => Err(RefinementError::Geometry(format!(
91                "unknown refinement pass {label:?}"
92            ))),
93        }
94    }
95}
96
97/// Authenticated node-trailer record of refinements already applied.
98#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
99pub struct RefinementPasses(u16);
100
101impl RefinementPasses {
102    pub(crate) const KNOWN_BITS: u16 = (1_u16 << 5) - 1;
103
104    /// Returns whether the named pass is stamped on this graph artifact.
105    #[must_use]
106    pub const fn contains(self, pass: RefinementPass) -> bool {
107        self.0 & pass.bit() != 0
108    }
109
110    pub(crate) const fn bits(self) -> u16 {
111        self.0
112    }
113
114    pub(crate) const fn from_bits(bits: u16) -> Option<Self> {
115        if bits & !Self::KNOWN_BITS == 0 {
116            Some(Self(bits))
117        } else {
118            None
119        }
120    }
121
122    pub(crate) const fn with(self, pass: RefinementPass) -> Self {
123        Self(self.0 | pass.bit())
124    }
125}
126
127impl FromIterator<RefinementPass> for RefinementPasses {
128    fn from_iter<T: IntoIterator<Item = RefinementPass>>(passes: T) -> Self {
129        passes
130            .into_iter()
131            .fold(Self::default(), |record, pass| record.with(pass))
132    }
133}
134
135/// Complete immutable graph replacement produced by one refinement pass.
136pub struct RefinementArtifact {
137    pass: RefinementPass,
138    encoded_region: Vec<u8>,
139    old_to_new: Option<Vec<u32>>,
140    new_to_old: Option<Vec<u32>>,
141    work_rows_completed: u64,
142}
143
144impl RefinementArtifact {
145    /// Returns the pass represented by this artifact.
146    #[must_use]
147    pub const fn pass(&self) -> RefinementPass {
148        self.pass
149    }
150
151    /// Returns the complete authenticated graph-region bytes.
152    #[must_use]
153    pub fn encoded_region(&self) -> &[u8] {
154        &self.encoded_region
155    }
156
157    /// Returns old-row to new-row ids for renumbering, and `None` otherwise.
158    #[must_use]
159    pub fn old_to_new_rows(&self) -> Option<&[u32]> {
160        self.old_to_new.as_deref()
161    }
162
163    /// Returns deterministic row work completed by this invocation.
164    #[must_use]
165    pub const fn work_rows_completed(&self) -> u64 {
166        self.work_rows_completed
167    }
168
169    /// Writes the replacement beside its immutable source without publishing it.
170    pub fn write_segment(
171        &self,
172        vfs: &dyn Vfs,
173        directory: &Path,
174        input: &SegmentReader,
175        output_id: SegmentId,
176        analyzer: &Analyzer,
177        policy: DurabilityPolicy,
178    ) -> Result<SegmentMeta, RefinementError> {
179        let Some(new_to_old) = self.new_to_old.as_deref() else {
180            return write_segment_with_encoded_graph(
181                &self.encoded_region,
182                vfs,
183                directory,
184                input,
185                output_id,
186                policy,
187            )
188            .map_err(RefinementError::Graph);
189        };
190        let intermediate_id = refinement_intermediate_id(output_id);
191        write_permuted_segment(
192            vfs,
193            directory,
194            input,
195            new_to_old,
196            intermediate_id,
197            analyzer,
198            policy,
199        )?;
200        let intermediate_path = directory.join(intermediate_id.file_name());
201        let intermediate = match SegmentReader::open(vfs, &intermediate_path, intermediate_id) {
202            Ok(intermediate) => intermediate,
203            Err(error) => {
204                let _ = vfs.delete(&intermediate_path);
205                return Err(error.into());
206            }
207        };
208        let written = write_segment_with_encoded_graph(
209            &self.encoded_region,
210            vfs,
211            directory,
212            &intermediate,
213            output_id,
214            policy,
215        )
216        .map_err(RefinementError::Graph);
217        drop(intermediate);
218        let removed = vfs.delete(&intermediate_path).map_err(|source| {
219            RefinementError::Segment(SegmentError::io(&intermediate_path, source))
220        });
221        match (written, removed) {
222            (Ok(written), Ok(())) => Ok(written),
223            (Err(error), _) | (Ok(_), Err(error)) => Err(error),
224        }
225    }
226}
227
228/// Typed failure before a refined artifact is published.
229#[derive(Debug)]
230pub enum RefinementError {
231    /// The graph already records the requested pass.
232    AlreadyApplied(RefinementPass),
233    /// Existing graph construction or graph-region writing failed.
234    Graph(GraphBuildError),
235    /// Fixed-stride graph bytes violated their persisted contract.
236    NodeBlock(GraphNodeError),
237    /// Reading or writing a sealed segment failed.
238    Segment(SegmentError),
239    /// Shared row-gather or accounting work failed.
240    Store(StoreError),
241    /// Checked row or pass geometry was invalid.
242    Geometry(String),
243}
244
245impl std::fmt::Display for RefinementError {
246    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
247        match self {
248            Self::AlreadyApplied(pass) => {
249                write!(
250                    formatter,
251                    "refinement pass {} is already applied",
252                    pass.label()
253                )
254            }
255            Self::Graph(error) => write!(formatter, "refinement graph: {error}"),
256            Self::NodeBlock(error) => write!(formatter, "refinement node blocks: {error}"),
257            Self::Segment(error) => write!(formatter, "refinement segment: {error}"),
258            Self::Store(error) => write!(formatter, "refinement store: {error}"),
259            Self::Geometry(detail) => write!(formatter, "refinement geometry: {detail}"),
260        }
261    }
262}
263
264impl std::error::Error for RefinementError {
265    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
266        match self {
267            Self::Graph(error) => Some(error),
268            Self::NodeBlock(error) => Some(error),
269            Self::Segment(error) => Some(error),
270            Self::Store(error) => Some(error),
271            Self::AlreadyApplied(_) | Self::Geometry(_) => None,
272        }
273    }
274}
275
276impl From<GraphNodeError> for RefinementError {
277    fn from(error: GraphNodeError) -> Self {
278        Self::NodeBlock(error)
279    }
280}
281
282impl From<SegmentError> for RefinementError {
283    fn from(error: SegmentError) -> Self {
284        Self::Segment(error)
285    }
286}
287
288impl From<StoreError> for RefinementError {
289    fn from(error: StoreError) -> Self {
290        Self::Store(error)
291    }
292}
293
294impl From<GraphBuildError> for RefinementError {
295    fn from(error: GraphBuildError) -> Self {
296        Self::Graph(error)
297    }
298}
299
300struct OwnedGraph {
301    layout: GraphNodeLayout,
302    passes: RefinementPasses,
303    flags: Vec<u8>,
304    neighbors: Vec<Vec<u32>>,
305}
306
307#[derive(Clone, Copy, Debug, Eq, PartialEq)]
308enum RefinementPhase {
309    Rows,
310    SeedMedoid,
311    SeedSpread,
312    ConnectivityTraverse,
313    ConnectivityDensify,
314    Complete,
315}
316
317struct RefinementState {
318    graph: OwnedGraph,
319    phase: RefinementPhase,
320    next_row: u32,
321    work_rows: u64,
322    permutation: Vec<u32>,
323    queue: VecDeque<u32>,
324    seen: Vec<u8>,
325    entries: Vec<u32>,
326    best_node: Option<u32>,
327    best_score: f64,
328    scratch_bytes: usize,
329    memory: Option<AccountedCounter>,
330}
331
332/// One resumable post-consolidation refinement request.
333pub struct CheckpointedRefinement<'a> {
334    pass: RefinementPass,
335    params: GraphParams,
336    seed: u64,
337    generation: u64,
338    checkpoint_path: &'a Path,
339    control: &'a QueryControl,
340    max_work_rows: u64,
341}
342
343impl<'a> CheckpointedRefinement<'a> {
344    /// Creates one request whose published generation is stamped into the checkpoint
345    /// for provenance; resume identity deliberately excludes it (see the decoder).
346    #[must_use]
347    pub const fn new(
348        pass: RefinementPass,
349        params: GraphParams,
350        seed: u64,
351        generation: u64,
352        checkpoint_path: &'a Path,
353        control: &'a QueryControl,
354    ) -> Self {
355        Self {
356            pass,
357            params,
358            seed,
359            generation,
360            checkpoint_path,
361            control,
362            max_work_rows: u64::MAX,
363        }
364    }
365
366    /// Caps deterministic row work performed by this invocation.
367    #[must_use]
368    pub const fn with_max_work_rows(mut self, max_work_rows: u64) -> Self {
369        self.max_work_rows = max_work_rows;
370        self
371    }
372}
373
374impl OwnedGraph {
375    fn read(graph: GraphNodeBlocks<'_>) -> Result<Self, RefinementError> {
376        let mut flags = Vec::with_capacity(graph.node_count() as usize);
377        let mut neighbors = Vec::with_capacity(graph.node_count() as usize);
378        for node_id in 0..graph.node_count() {
379            let block = graph.block(node_id)?;
380            flags.push(block.flags());
381            neighbors.push(
382                block
383                    .neighbors_padded()
384                    .take(usize::from(block.degree()))
385                    .collect(),
386            );
387        }
388        Ok(Self {
389            layout: graph.layout(),
390            passes: graph.refinement_passes(),
391            flags,
392            neighbors,
393        })
394    }
395}
396
397impl RefinementState {
398    fn new(
399        reader: &SegmentReader,
400        pass: RefinementPass,
401        params: GraphParams,
402        accounting: Option<&std::sync::Arc<crate::lifecycle::stats::Accounting>>,
403    ) -> Result<Self, RefinementError> {
404        let source = reader.graph_node_blocks()?;
405        let graph = OwnedGraph::read(source)?;
406        let node_count = graph.neighbors.len();
407        let (phase, permutation, queue, seen) = match pass {
408            RefinementPass::Renumber => {
409                let mut roots = graph
410                    .flags
411                    .iter()
412                    .enumerate()
413                    .filter_map(|(node, flags)| {
414                        (*flags & ENTRY_FLAG != 0)
415                            .then(|| u32::try_from(node).ok())
416                            .flatten()
417                    })
418                    .collect::<Vec<_>>();
419                if roots.is_empty() {
420                    return Err(RefinementError::Geometry(
421                        "renumbering requires at least one persisted entry seed".to_owned(),
422                    ));
423                }
424                sort_by_degree(&mut roots, &graph)?;
425                let mut queue = VecDeque::with_capacity(node_count);
426                let mut seen = vec![0_u8; node_count];
427                for root in roots {
428                    mark_and_enqueue_byte(root, &mut seen, &mut queue)?;
429                }
430                (
431                    RefinementPhase::Rows,
432                    Vec::with_capacity(node_count),
433                    queue,
434                    seen,
435                )
436            }
437            RefinementPass::ConnectivityRepair => {
438                if graph.layout.max_degree() != params.r_max() {
439                    return Err(RefinementError::Geometry(format!(
440                        "graph max degree {}, profile max degree {}",
441                        graph.layout.max_degree(),
442                        params.r_max()
443                    )));
444                }
445                let root = first_entry(&graph)?;
446                let mut queue = VecDeque::with_capacity(node_count);
447                let mut seen = vec![0_u8; node_count];
448                mark_and_enqueue_byte(root, &mut seen, &mut queue)?;
449                (
450                    RefinementPhase::ConnectivityTraverse,
451                    Vec::new(),
452                    queue,
453                    seen,
454                )
455            }
456            RefinementPass::SeedRefit => (
457                RefinementPhase::SeedMedoid,
458                Vec::new(),
459                VecDeque::new(),
460                Vec::new(),
461            ),
462            RefinementPass::AlphaReprune | RefinementPass::NeighborReorder => (
463                RefinementPhase::Rows,
464                Vec::new(),
465                VecDeque::new(),
466                Vec::new(),
467            ),
468        };
469        let mut state = Self {
470            graph,
471            phase,
472            next_row: 0,
473            work_rows: 0,
474            permutation,
475            queue,
476            seen,
477            entries: Vec::with_capacity(ENTRY_POINT_COUNT),
478            best_node: None,
479            best_score: 0.0,
480            scratch_bytes: refinement_scratch_bytes(pass, params)?,
481            memory: match accounting {
482                Some(accounting) => Some(AccountedCounter::new(
483                    accounting,
484                    AllocationComponent::Temporary,
485                )?),
486                None => None,
487            },
488        };
489        state.update_memory_charge()?;
490        Ok(state)
491    }
492
493    fn update_memory_charge(&mut self) -> Result<(), RefinementError> {
494        let graph_rows = self
495            .graph
496            .neighbors
497            .capacity()
498            .checked_mul(std::mem::size_of::<Vec<u32>>())
499            .ok_or_else(|| {
500                RefinementError::Geometry("refinement graph charge overflow".to_owned())
501            })?;
502        let neighbor_bytes = self
503            .graph
504            .neighbors
505            .iter()
506            .try_fold(0_usize, |total, row| {
507                row.capacity()
508                    .checked_mul(std::mem::size_of::<u32>())
509                    .and_then(|bytes| total.checked_add(bytes))
510                    .ok_or_else(|| {
511                        RefinementError::Geometry("refinement neighbor charge overflow".to_owned())
512                    })
513            })?;
514        let bytes = graph_rows
515            .checked_add(neighbor_bytes)
516            .and_then(|total| total.checked_add(self.graph.flags.capacity()))
517            .and_then(|total| {
518                self.permutation
519                    .capacity()
520                    .checked_mul(std::mem::size_of::<u32>())
521                    .and_then(|bytes| total.checked_add(bytes))
522            })
523            .and_then(|total| {
524                self.queue
525                    .capacity()
526                    .checked_mul(std::mem::size_of::<u32>())
527                    .and_then(|bytes| total.checked_add(bytes))
528            })
529            .and_then(|total| total.checked_add(self.seen.capacity()))
530            .and_then(|total| {
531                self.entries
532                    .capacity()
533                    .checked_mul(std::mem::size_of::<u32>())
534                    .and_then(|bytes| total.checked_add(bytes))
535            })
536            .and_then(|total| total.checked_add(self.scratch_bytes))
537            .ok_or_else(|| {
538                RefinementError::Geometry("refinement state charge overflow".to_owned())
539            })?;
540        if let Some(memory) = self.memory.as_mut() {
541            memory.set(bytes)?;
542        }
543        Ok(())
544    }
545
546    fn advance_batch(
547        &mut self,
548        reader: &SegmentReader,
549        pass: RefinementPass,
550        params: GraphParams,
551        seed: u64,
552        limit: u64,
553        cancellation: &QueryCancellation<'_>,
554    ) -> Result<bool, RefinementError> {
555        let started = self.work_rows;
556        while self.phase != RefinementPhase::Complete
557            && self.work_rows.saturating_sub(started) < limit
558        {
559            check_refinement_cancellation(cancellation)?;
560            match pass {
561                RefinementPass::Renumber => self.advance_renumber()?,
562                RefinementPass::AlphaReprune => {
563                    self.advance_alpha(reader, params)?;
564                }
565                RefinementPass::SeedRefit => self.advance_seed(reader)?,
566                RefinementPass::NeighborReorder => self.advance_reorder(reader, seed)?,
567                RefinementPass::ConnectivityRepair => self.advance_connectivity(params)?,
568            }
569        }
570        check_refinement_cancellation(cancellation)?;
571        Ok(self.phase == RefinementPhase::Complete)
572    }
573
574    fn advance_renumber(&mut self) -> Result<(), RefinementError> {
575        if self.queue.is_empty() {
576            if let Some(unseen) = self.seen.iter().position(|seen| *seen == 0) {
577                let root = u32::try_from(unseen).map_err(|_| {
578                    RefinementError::Geometry("renumber row exceeds u32".to_owned())
579                })?;
580                mark_and_enqueue_byte(root, &mut self.seen, &mut self.queue)?;
581            } else {
582                let node_count = u32::try_from(self.graph.neighbors.len())
583                    .map_err(|_| RefinementError::Geometry("graph rows exceed u32".to_owned()))?;
584                let old_to_new = invert_permutation(&self.permutation, node_count)?;
585                apply_renumber(&mut self.graph, &self.permutation, &old_to_new)?;
586                self.phase = RefinementPhase::Complete;
587                return Ok(());
588            }
589        }
590        let owner = self.queue.pop_front().ok_or_else(|| {
591            RefinementError::Geometry("renumber queue became empty before completion".to_owned())
592        })?;
593        self.permutation.push(owner);
594        let node_count = u32::try_from(self.graph.neighbors.len())
595            .map_err(|_| RefinementError::Geometry("graph rows exceed u32".to_owned()))?;
596        let owner_index = node_index(owner, node_count)?;
597        let mut candidates = self
598            .graph
599            .neighbors
600            .get(owner_index)
601            .ok_or_else(|| {
602                RefinementError::Geometry(format!("neighbors for node {owner} are unavailable"))
603            })?
604            .iter()
605            .copied()
606            .filter(|neighbor| {
607                self.seen
608                    .get(*neighbor as usize)
609                    .is_some_and(|seen| *seen == 0)
610            })
611            .collect::<Vec<_>>();
612        sort_by_degree(&mut candidates, &self.graph)?;
613        for candidate in candidates {
614            mark_and_enqueue_byte(candidate, &mut self.seen, &mut self.queue)?;
615        }
616        self.work_rows = self
617            .work_rows
618            .checked_add(1)
619            .ok_or_else(|| RefinementError::Geometry("renumber work overflow".to_owned()))?;
620        if self.permutation.len() == self.graph.neighbors.len() {
621            let old_to_new = invert_permutation(&self.permutation, node_count)?;
622            apply_renumber(&mut self.graph, &self.permutation, &old_to_new)?;
623            self.phase = RefinementPhase::Complete;
624        }
625        Ok(())
626    }
627
628    fn advance_connectivity(&mut self, params: GraphParams) -> Result<(), RefinementError> {
629        let node_count = u32::try_from(self.graph.neighbors.len())
630            .map_err(|_| RefinementError::Geometry("graph rows exceed u32".to_owned()))?;
631        match self.phase {
632            RefinementPhase::ConnectivityTraverse => {
633                if let Some(owner) = self.queue.pop_front() {
634                    let neighbors = self
635                        .graph
636                        .neighbors
637                        .get(node_index(owner, node_count)?)
638                        .ok_or_else(|| {
639                            RefinementError::Geometry(format!(
640                                "neighbors for repair node {owner} are unavailable"
641                            ))
642                        })?
643                        .clone();
644                    for neighbor in neighbors {
645                        mark_and_enqueue_byte(neighbor, &mut self.seen, &mut self.queue)?;
646                    }
647                    self.work_rows = self.work_rows.checked_add(1).ok_or_else(|| {
648                        RefinementError::Geometry("connectivity work overflow".to_owned())
649                    })?;
650                } else if let Some(unseen) = self.seen.iter().position(|marker| *marker == 0) {
651                    let target = u32::try_from(unseen).map_err(|_| {
652                        RefinementError::Geometry("unreachable row exceeds u32".to_owned())
653                    })?;
654                    stitch_component(&mut self.graph, &self.seen, target, params.r_max())?;
655                    mark_and_enqueue_byte(target, &mut self.seen, &mut self.queue)?;
656                } else {
657                    self.phase = RefinementPhase::ConnectivityDensify;
658                    self.next_row = 0;
659                }
660            }
661            RefinementPhase::ConnectivityDensify => {
662                let owner = self.next_row;
663                densify_hub(&mut self.graph, owner, params.r_max())?;
664                self.finish_row(node_count, "connectivity densification")?;
665            }
666            RefinementPhase::Rows
667            | RefinementPhase::SeedMedoid
668            | RefinementPhase::SeedSpread
669            | RefinementPhase::Complete => {
670                return Err(RefinementError::Geometry(
671                    "connectivity repair entered an invalid phase".to_owned(),
672                ));
673            }
674        }
675        Ok(())
676    }
677
678    fn advance_alpha(
679        &mut self,
680        reader: &SegmentReader,
681        params: GraphParams,
682    ) -> Result<(), RefinementError> {
683        let source = reader.graph_node_blocks()?;
684        if source.layout().max_degree() != params.r_max() {
685            return Err(RefinementError::Geometry(format!(
686                "graph max degree {}, profile max degree {}",
687                source.layout().max_degree(),
688                params.r_max()
689            )));
690        }
691        let owner = self.next_row;
692        let owner_block = source.block(owner)?;
693        let own = owner_block
694            .neighbors_padded()
695            .take(usize::from(owner_block.degree()))
696            .collect::<Vec<_>>();
697        let candidate_capacity = usize::from(params.r_max())
698            .checked_add(
699                usize::from(params.r_max())
700                    .checked_mul(usize::from(params.r_max()))
701                    .ok_or_else(|| {
702                        RefinementError::Geometry("re-prune candidates overflow".to_owned())
703                    })?,
704            )
705            .ok_or_else(|| RefinementError::Geometry("re-prune candidates overflow".to_owned()))?;
706        let mut candidates = Vec::with_capacity(candidate_capacity);
707        candidates.extend_from_slice(&own);
708        for neighbor in &own {
709            let block = source.block(*neighbor)?;
710            candidates.extend(block.neighbors_padded().take(usize::from(block.degree())));
711        }
712        let rescore = reader.rescore_f32()?;
713        let dimensions = reader.meta().dims as usize;
714        let node_count = reader.meta().row_count;
715        let mut pruned = robust_prune_rows(
716            rescore,
717            dimensions,
718            node_count,
719            owner,
720            &candidates,
721            params.alpha_refine(),
722            usize::from(params.r_target()),
723        )?;
724        supplement_long_range(
725            rescore,
726            dimensions,
727            owner,
728            &own,
729            usize::from(params.r_max()),
730            &mut pruned,
731        )?;
732        *self
733            .graph
734            .neighbors
735            .get_mut(node_index(owner, node_count)?)
736            .ok_or_else(|| {
737                RefinementError::Geometry(format!(
738                    "output neighbors for node {owner} are unavailable"
739                ))
740            })? = pruned;
741        self.finish_row(node_count, "re-prune")
742    }
743
744    fn advance_reorder(
745        &mut self,
746        reader: &SegmentReader,
747        seed: u64,
748    ) -> Result<(), RefinementError> {
749        let owner = self.next_row;
750        let node_count = reader.meta().row_count;
751        let neighbors = self
752            .graph
753            .neighbors
754            .get_mut(node_index(owner, node_count)?)
755            .ok_or_else(|| {
756                RefinementError::Geometry(format!("neighbors for node {owner} are unavailable"))
757            })?;
758        reorder_neighbors_bit4(reader, owner, neighbors, seed)?;
759        self.finish_row(node_count, "neighbor reorder")
760    }
761
762    fn advance_seed(&mut self, reader: &SegmentReader) -> Result<(), RefinementError> {
763        let node_count = reader.meta().row_count;
764        if node_count == 0 {
765            return Err(RefinementError::Geometry(
766                "seed refit requires at least one row".to_owned(),
767            ));
768        }
769        if self.next_row == node_count {
770            self.finish_seed_phase(node_count)?;
771            if self.phase == RefinementPhase::Complete {
772                return Ok(());
773            }
774        }
775        let candidate = self.next_row;
776        let rescore = reader.rescore_f32()?;
777        let dimensions = reader.meta().dims as usize;
778        let score = match self.phase {
779            RefinementPhase::SeedMedoid => {
780                let mut sum = 0.0_f64;
781                for other in 0..node_count {
782                    sum += exact_distance(rescore, dimensions, node_count, candidate, other)?;
783                }
784                sum
785            }
786            RefinementPhase::SeedSpread => {
787                if self.entries.contains(&candidate) {
788                    f64::NEG_INFINITY
789                } else {
790                    self.entries
791                        .iter()
792                        .try_fold(f64::INFINITY, |nearest, entry| {
793                            exact_distance(rescore, dimensions, node_count, candidate, *entry)
794                                .map(|distance| nearest.min(distance))
795                        })?
796                }
797            }
798            RefinementPhase::Rows
799            | RefinementPhase::ConnectivityTraverse
800            | RefinementPhase::ConnectivityDensify
801            | RefinementPhase::Complete => {
802                return Err(RefinementError::Geometry(
803                    "seed refit entered an invalid phase".to_owned(),
804                ));
805            }
806        };
807        let better = match (self.phase, self.best_node) {
808            (_, None) => true,
809            (RefinementPhase::SeedMedoid, Some(best)) => {
810                score < self.best_score || (score == self.best_score && candidate < best)
811            }
812            (RefinementPhase::SeedSpread, Some(best)) => {
813                score > self.best_score || (score == self.best_score && candidate < best)
814            }
815            (
816                RefinementPhase::Rows
817                | RefinementPhase::ConnectivityTraverse
818                | RefinementPhase::ConnectivityDensify
819                | RefinementPhase::Complete,
820                Some(_),
821            ) => false,
822        };
823        if better {
824            self.best_node = Some(candidate);
825            self.best_score = score;
826        }
827        self.next_row = self
828            .next_row
829            .checked_add(1)
830            .ok_or_else(|| RefinementError::Geometry("seed row overflow".to_owned()))?;
831        self.work_rows = self
832            .work_rows
833            .checked_add(1)
834            .ok_or_else(|| RefinementError::Geometry("seed work overflow".to_owned()))?;
835        if self.next_row == node_count {
836            self.finish_seed_phase(node_count)?;
837        }
838        Ok(())
839    }
840
841    fn finish_seed_phase(&mut self, node_count: u32) -> Result<(), RefinementError> {
842        let selected = self.best_node.take().ok_or_else(|| {
843            RefinementError::Geometry("seed phase produced no candidate".to_owned())
844        })?;
845        self.entries.push(selected);
846        self.best_score = 0.0;
847        self.next_row = 0;
848        if self.entries.len() >= ENTRY_POINT_COUNT || self.entries.len() >= node_count as usize {
849            for (node, flags) in self.graph.flags.iter_mut().enumerate() {
850                let node = u32::try_from(node)
851                    .map_err(|_| RefinementError::Geometry("entry row exceeds u32".to_owned()))?;
852                *flags = (*flags & HUB_FLAG) | u8::from(self.entries.contains(&node));
853            }
854            self.phase = RefinementPhase::Complete;
855        } else {
856            self.phase = RefinementPhase::SeedSpread;
857        }
858        Ok(())
859    }
860
861    fn finish_row(&mut self, node_count: u32, label: &str) -> Result<(), RefinementError> {
862        self.next_row = self
863            .next_row
864            .checked_add(1)
865            .ok_or_else(|| RefinementError::Geometry(format!("{label} row overflow")))?;
866        self.work_rows = self
867            .work_rows
868            .checked_add(1)
869            .ok_or_else(|| RefinementError::Geometry(format!("{label} work overflow")))?;
870        if self.next_row == node_count {
871            self.phase = RefinementPhase::Complete;
872        }
873        Ok(())
874    }
875}
876
877fn mark_and_enqueue_byte(
878    node: u32,
879    seen: &mut [u8],
880    queue: &mut VecDeque<u32>,
881) -> Result<(), RefinementError> {
882    let marker = seen.get_mut(node as usize).ok_or_else(|| {
883        RefinementError::Geometry(format!("renumber queue node {node} is out of range"))
884    })?;
885    if *marker == 0 {
886        *marker = 1;
887        queue.push_back(node);
888    }
889    Ok(())
890}
891
892fn first_entry(graph: &OwnedGraph) -> Result<u32, RefinementError> {
893    graph
894        .flags
895        .iter()
896        .position(|flags| *flags & ENTRY_FLAG != 0)
897        .ok_or_else(|| {
898            RefinementError::Geometry(
899                "connectivity repair requires at least one persisted entry seed".to_owned(),
900            )
901        })
902        .and_then(|node| {
903            u32::try_from(node)
904                .map_err(|_| RefinementError::Geometry("entry row exceeds u32".to_owned()))
905        })
906}
907
908fn stitch_component(
909    graph: &mut OwnedGraph,
910    seen: &[u8],
911    target: u32,
912    r_max: u8,
913) -> Result<(), RefinementError> {
914    let maximum = usize::from(r_max);
915    if maximum == 0 {
916        return Err(RefinementError::Geometry(
917            "connectivity repair requires positive maximum degree".to_owned(),
918        ));
919    }
920    let (source, replacement) = connectivity_source(graph, seen, maximum)?;
921    set_or_append_neighbor(graph, source, target, maximum, replacement)?;
922    let target_index = node_index(
923        target,
924        u32::try_from(graph.neighbors.len())
925            .map_err(|_| RefinementError::Geometry("graph rows exceed u32".to_owned()))?,
926    )?;
927    let target_neighbors = graph.neighbors.get_mut(target_index).ok_or_else(|| {
928        RefinementError::Geometry(format!(
929            "neighbors for stitch target {target} are unavailable"
930        ))
931    })?;
932    if !target_neighbors.contains(&source) {
933        if target_neighbors.len() < maximum {
934            target_neighbors.push(source);
935        } else {
936            let slot = target_neighbors.len().saturating_sub(1);
937            let destination = target_neighbors.get_mut(slot).ok_or_else(|| {
938                RefinementError::Geometry("full stitch target has no replaceable edge".to_owned())
939            })?;
940            *destination = source;
941        }
942    }
943    let source_flag = graph.flags.get_mut(source as usize).ok_or_else(|| {
944        RefinementError::Geometry(format!("flags for stitch source {source} are unavailable"))
945    })?;
946    *source_flag |= HUB_FLAG;
947    Ok(())
948}
949
950fn connectivity_source(
951    graph: &OwnedGraph,
952    seen: &[u8],
953    maximum: usize,
954) -> Result<(u32, Option<usize>), RefinementError> {
955    for required_flag in [HUB_FLAG, ENTRY_FLAG, 0] {
956        for (owner, marker) in seen.iter().enumerate() {
957            let flags = graph.flags.get(owner).copied().ok_or_else(|| {
958                RefinementError::Geometry("repair flags are shorter than seen state".to_owned())
959            })?;
960            let neighbors = graph.neighbors.get(owner).ok_or_else(|| {
961                RefinementError::Geometry("repair rows are shorter than seen state".to_owned())
962            })?;
963            if *marker != 0
964                && neighbors.len() < maximum
965                && (required_flag == 0 || flags & required_flag != 0)
966            {
967                return u32::try_from(owner)
968                    .map(|owner| (owner, None))
969                    .map_err(|_| {
970                        RefinementError::Geometry("repair source exceeds u32".to_owned())
971                    });
972            }
973        }
974    }
975    replaceable_non_tree_edge(graph, seen).map(|(owner, slot)| (owner, Some(slot)))
976}
977
978fn replaceable_non_tree_edge(
979    graph: &OwnedGraph,
980    seen: &[u8],
981) -> Result<(u32, usize), RefinementError> {
982    let node_count = u32::try_from(graph.neighbors.len())
983        .map_err(|_| RefinementError::Geometry("graph rows exceed u32".to_owned()))?;
984    let root = first_entry(graph)?;
985    let mut parents = vec![u32::MAX; graph.neighbors.len()];
986    let root_index = node_index(root, node_count)?;
987    *parents.get_mut(root_index).ok_or_else(|| {
988        RefinementError::Geometry("repair root parent is unavailable".to_owned())
989    })? = root;
990    let mut queue = VecDeque::from([root]);
991    while let Some(owner) = queue.pop_front() {
992        let neighbors = graph
993            .neighbors
994            .get(node_index(owner, node_count)?)
995            .ok_or_else(|| {
996                RefinementError::Geometry(format!("repair tree row {owner} is unavailable"))
997            })?;
998        for neighbor in neighbors {
999            let index = node_index(*neighbor, node_count)?;
1000            if seen.get(index).copied() == Some(1) && parents.get(index).copied() == Some(u32::MAX)
1001            {
1002                *parents.get_mut(index).ok_or_else(|| {
1003                    RefinementError::Geometry("repair parent row is unavailable".to_owned())
1004                })? = owner;
1005                queue.push_back(*neighbor);
1006            }
1007        }
1008    }
1009    for (owner, marker) in seen.iter().enumerate() {
1010        if *marker == 0 {
1011            continue;
1012        }
1013        let owner_u32 = u32::try_from(owner)
1014            .map_err(|_| RefinementError::Geometry("repair owner exceeds u32".to_owned()))?;
1015        let neighbors = graph.neighbors.get(owner).ok_or_else(|| {
1016            RefinementError::Geometry("repair source row is unavailable".to_owned())
1017        })?;
1018        for (slot, neighbor) in neighbors.iter().enumerate().rev() {
1019            let neighbor_index = node_index(*neighbor, node_count)?;
1020            if parents.get(neighbor_index).copied() != Some(owner_u32) {
1021                return Ok((owner_u32, slot));
1022            }
1023        }
1024    }
1025    Err(RefinementError::Geometry(
1026        "connected component has no spare or non-tree edge for stitching".to_owned(),
1027    ))
1028}
1029
1030fn set_or_append_neighbor(
1031    graph: &mut OwnedGraph,
1032    owner: u32,
1033    target: u32,
1034    maximum: usize,
1035    replacement: Option<usize>,
1036) -> Result<(), RefinementError> {
1037    let node_count = u32::try_from(graph.neighbors.len())
1038        .map_err(|_| RefinementError::Geometry("graph rows exceed u32".to_owned()))?;
1039    let neighbors = graph
1040        .neighbors
1041        .get_mut(node_index(owner, node_count)?)
1042        .ok_or_else(|| {
1043            RefinementError::Geometry(format!(
1044                "neighbors for stitch source {owner} are unavailable"
1045            ))
1046        })?;
1047    if neighbors.contains(&target) {
1048        return Ok(());
1049    }
1050    if neighbors.len() < maximum {
1051        neighbors.push(target);
1052        return Ok(());
1053    }
1054    let slot = replacement.ok_or_else(|| {
1055        RefinementError::Geometry(format!("full stitch source {owner} has no replacement"))
1056    })?;
1057    let destination = neighbors.get_mut(slot).ok_or_else(|| {
1058        RefinementError::Geometry(format!("stitch replacement {slot} is unavailable"))
1059    })?;
1060    *destination = target;
1061    Ok(())
1062}
1063
1064fn densify_hub(graph: &mut OwnedGraph, owner: u32, r_max: u8) -> Result<(), RefinementError> {
1065    let node_count = u32::try_from(graph.neighbors.len())
1066        .map_err(|_| RefinementError::Geometry("graph rows exceed u32".to_owned()))?;
1067    let owner_index = node_index(owner, node_count)?;
1068    if graph.flags.get(owner_index).copied().ok_or_else(|| {
1069        RefinementError::Geometry(format!("flags for repair row {owner} are unavailable"))
1070    })? & HUB_FLAG
1071        == 0
1072    {
1073        return Ok(());
1074    }
1075    let target_degree = usize::from(r_max).min(graph.neighbors.len().saturating_sub(1));
1076    let neighbors = graph.neighbors.get_mut(owner_index).ok_or_else(|| {
1077        RefinementError::Geometry(format!("neighbors for repair hub {owner} are unavailable"))
1078    })?;
1079    for candidate in 0..node_count {
1080        if neighbors.len() >= target_degree {
1081            break;
1082        }
1083        if candidate != owner && !neighbors.contains(&candidate) {
1084            neighbors.push(candidate);
1085        }
1086    }
1087    if neighbors.len() != target_degree {
1088        return Err(RefinementError::Geometry(format!(
1089            "repair hub {owner} densified to {}, expected {target_degree}",
1090            neighbors.len()
1091        )));
1092    }
1093    Ok(())
1094}
1095
1096/// Runs one refinement in deterministic, durably checkpointed row batches.
1097pub fn refine_graph_checkpointed(
1098    store: &Store,
1099    reader: &SegmentReader,
1100    request: CheckpointedRefinement<'_>,
1101    lease: &SnapshotLease,
1102) -> Result<RefinementArtifact, RefinementError> {
1103    let source = reader.graph_node_blocks()?;
1104    if source.refinement_passes().contains(request.pass) {
1105        return Err(RefinementError::AlreadyApplied(request.pass));
1106    }
1107    if source.node_count() != reader.meta().row_count {
1108        return Err(RefinementError::Geometry(format!(
1109            "graph rows {}, segment rows {}",
1110            source.node_count(),
1111            reader.meta().row_count
1112        )));
1113    }
1114    let mut state = match store.vfs.read(request.checkpoint_path) {
1115        Ok(bytes) => match decode_refinement_checkpoint(
1116            &bytes,
1117            reader,
1118            request.pass,
1119            request.params,
1120            request.seed,
1121            Some(&store.accounting),
1122        ) {
1123            Ok(state) => state,
1124            Err(error @ RefinementError::Graph(GraphBuildError::CheckpointCorrupt(_))) => {
1125                remove_refinement_checkpoint(store.vfs.as_ref(), request.checkpoint_path)?;
1126                return Err(error);
1127            }
1128            Err(error) => return Err(error),
1129        },
1130        Err(error) if error.kind() == std::io::ErrorKind::NotFound => RefinementState::new(
1131            reader,
1132            request.pass,
1133            request.params,
1134            Some(&store.accounting),
1135        )?,
1136        Err(source) => {
1137            return Err(GraphBuildError::CheckpointIo {
1138                path: request.checkpoint_path.to_path_buf(),
1139                source,
1140            }
1141            .into());
1142        }
1143    };
1144    let started = state.work_rows;
1145    let cancellation = QueryCancellation::new(request.control, lease);
1146    while state.phase != RefinementPhase::Complete {
1147        let completed = state.work_rows.saturating_sub(started);
1148        if completed >= request.max_work_rows {
1149            return Err(GraphBuildError::BudgetExhausted {
1150                rows_completed: completed,
1151            }
1152            .into());
1153        }
1154        let batch = u64::from(request.params.checkpoint_batch_rows())
1155            .min(request.max_work_rows.saturating_sub(completed));
1156        let complete = state.advance_batch(
1157            reader,
1158            request.pass,
1159            request.params,
1160            request.seed,
1161            batch,
1162            &cancellation,
1163        )?;
1164        write_refinement_checkpoint(
1165            store.vfs.as_ref(),
1166            store.durability_policy,
1167            request.checkpoint_path,
1168            reader,
1169            request.pass,
1170            request.params,
1171            request.seed,
1172            request.generation,
1173            &mut state,
1174        )?;
1175        if complete {
1176            break;
1177        }
1178    }
1179    let work_rows_completed = state.work_rows.saturating_sub(started);
1180    state.graph.passes = state.graph.passes.with(request.pass);
1181    let new_to_old =
1182        (request.pass == RefinementPass::Renumber).then(|| std::mem::take(&mut state.permutation));
1183    let old_to_new = match new_to_old.as_deref() {
1184        Some(order) => Some(invert_permutation(order, reader.meta().row_count)?),
1185        None => None,
1186    };
1187    let encoded_region = encode_owned_graph(reader, &state.graph, new_to_old.as_deref())?;
1188    remove_refinement_checkpoint(store.vfs.as_ref(), request.checkpoint_path)?;
1189    Ok(RefinementArtifact {
1190        pass: request.pass,
1191        encoded_region,
1192        old_to_new,
1193        new_to_old,
1194        work_rows_completed,
1195    })
1196}
1197
1198#[allow(clippy::too_many_arguments)]
1199fn write_refinement_checkpoint(
1200    vfs: &dyn Vfs,
1201    policy: DurabilityPolicy,
1202    path: &Path,
1203    reader: &SegmentReader,
1204    pass: RefinementPass,
1205    params: GraphParams,
1206    seed: u64,
1207    generation: u64,
1208    state: &mut RefinementState,
1209) -> Result<(), RefinementError> {
1210    let base_charge = state.memory.as_ref().map_or(0, AccountedCounter::bytes);
1211    let checkpoint_bytes = refinement_checkpoint_size(reader, params, state)?;
1212    if let Some(memory) = state.memory.as_mut() {
1213        let peak = usize::try_from(base_charge)
1214            .ok()
1215            .and_then(|base| base.checked_add(checkpoint_bytes))
1216            .ok_or_else(|| {
1217                RefinementError::Geometry("checkpoint peak charge overflow".to_owned())
1218            })?;
1219        memory.set(peak)?;
1220    }
1221    let result = encode_refinement_checkpoint(reader, pass, params, seed, generation, state)
1222        .and_then(|bytes| publish_refinement_checkpoint(vfs, policy, path, &bytes));
1223    if let Some(memory) = state.memory.as_mut() {
1224        memory.set(usize::try_from(base_charge).map_err(|_| {
1225            RefinementError::Geometry("checkpoint base charge exceeds usize".to_owned())
1226        })?)?;
1227    }
1228    result
1229}
1230
1231fn refinement_checkpoint_size(
1232    reader: &SegmentReader,
1233    params: GraphParams,
1234    state: &RefinementState,
1235) -> Result<usize, RefinementError> {
1236    let node_count = reader.meta().row_count as usize;
1237    let graph_bytes = node_count
1238        .checked_mul(usize::from(params.r_max()))
1239        .and_then(|slots| slots.checked_mul(std::mem::size_of::<u32>()))
1240        .and_then(|bytes| bytes.checked_add(node_count.checked_mul(2)?))
1241        .ok_or_else(|| RefinementError::Geometry("checkpoint graph size overflow".to_owned()))?;
1242    let aux_bytes = checkpoint_aux_bytes(state)?;
1243    CHECKPOINT_HEADER_BYTES
1244        .checked_add(graph_bytes)
1245        .and_then(|bytes| bytes.checked_add(aux_bytes))
1246        .and_then(|bytes| bytes.checked_add(CHECKPOINT_CHECKSUM_BYTES))
1247        .ok_or_else(|| RefinementError::Geometry("checkpoint length overflow".to_owned()))
1248}
1249
1250#[allow(clippy::too_many_arguments)]
1251fn encode_refinement_checkpoint(
1252    reader: &SegmentReader,
1253    pass: RefinementPass,
1254    params: GraphParams,
1255    seed: u64,
1256    generation: u64,
1257    state: &RefinementState,
1258) -> Result<Vec<u8>, RefinementError> {
1259    let node_count = reader.meta().row_count as usize;
1260    let r_max = usize::from(params.r_max());
1261    if state.graph.neighbors.len() != node_count || state.graph.flags.len() != node_count {
1262        return Err(RefinementError::Geometry(
1263            "checkpoint graph shape differs from source".to_owned(),
1264        ));
1265    }
1266    let graph_bytes = node_count
1267        .checked_mul(r_max)
1268        .and_then(|slots| slots.checked_mul(std::mem::size_of::<u32>()))
1269        .and_then(|bytes| bytes.checked_add(node_count.checked_mul(2)?))
1270        .ok_or_else(|| RefinementError::Geometry("checkpoint graph size overflow".to_owned()))?;
1271    let aux_bytes = checkpoint_aux_bytes(state)?;
1272    let total_bytes = CHECKPOINT_HEADER_BYTES
1273        .checked_add(graph_bytes)
1274        .and_then(|bytes| bytes.checked_add(aux_bytes))
1275        .and_then(|bytes| bytes.checked_add(CHECKPOINT_CHECKSUM_BYTES))
1276        .ok_or_else(|| RefinementError::Geometry("checkpoint length overflow".to_owned()))?;
1277    let mut bytes = Vec::with_capacity(total_bytes);
1278    bytes.extend_from_slice(&CHECKPOINT_MAGIC);
1279    bytes.extend_from_slice(&CHECKPOINT_VERSION.to_le_bytes());
1280    bytes.push(pass as u8);
1281    bytes.push(refinement_phase_code(state.phase));
1282    bytes.push(params.r_target());
1283    bytes.push(params.r_max());
1284    bytes.extend_from_slice(&0_u16.to_le_bytes());
1285    bytes.extend_from_slice(&params.checkpoint_batch_rows().to_le_bytes());
1286    bytes.extend_from_slice(&reader.meta().row_count.to_le_bytes());
1287    bytes.extend_from_slice(&reader.meta().dims.to_le_bytes());
1288    bytes.extend_from_slice(&state.next_row.to_le_bytes());
1289    bytes.extend_from_slice(&0_u32.to_le_bytes());
1290    bytes.extend_from_slice(&seed.to_le_bytes());
1291    bytes.extend_from_slice(&params.alpha_refine().to_bits().to_le_bytes());
1292    bytes.extend_from_slice(reader.meta().id.as_bytes());
1293    bytes.extend_from_slice(&state.graph.passes.bits().to_le_bytes());
1294    bytes.extend_from_slice(&0_u16.to_le_bytes());
1295    bytes.extend_from_slice(&generation.to_le_bytes());
1296    bytes.extend_from_slice(&state.work_rows.to_le_bytes());
1297    bytes.extend_from_slice(
1298        &u64::try_from(graph_bytes)
1299            .map_err(|_| RefinementError::Geometry("checkpoint graph bytes exceed u64".to_owned()))?
1300            .to_le_bytes(),
1301    );
1302    bytes.extend_from_slice(
1303        &u64::try_from(aux_bytes)
1304            .map_err(|_| RefinementError::Geometry("checkpoint aux bytes exceed u64".to_owned()))?
1305            .to_le_bytes(),
1306    );
1307    bytes.extend_from_slice(&0_u32.to_le_bytes());
1308    if bytes.len() != CHECKPOINT_HEADER_BYTES {
1309        return Err(RefinementError::Geometry(format!(
1310            "checkpoint header encoded {} bytes, expected {CHECKPOINT_HEADER_BYTES}",
1311            bytes.len()
1312        )));
1313    }
1314    bytes.extend_from_slice(&state.graph.flags);
1315    for neighbors in &state.graph.neighbors {
1316        let degree = u8::try_from(neighbors.len())
1317            .map_err(|_| RefinementError::Geometry("checkpoint degree exceeds u8".to_owned()))?;
1318        if usize::from(degree) > r_max {
1319            return Err(RefinementError::Geometry(format!(
1320                "checkpoint degree {degree} exceeds {}",
1321                params.r_max()
1322            )));
1323        }
1324        bytes.push(degree);
1325    }
1326    for (owner, neighbors) in state.graph.neighbors.iter().enumerate() {
1327        for neighbor in neighbors {
1328            if usize::try_from(*neighbor)
1329                .ok()
1330                .is_none_or(|node| node >= node_count)
1331                || usize::try_from(*neighbor).ok() == Some(owner)
1332            {
1333                return Err(RefinementError::Geometry(format!(
1334                    "checkpoint node {owner} has invalid neighbor {neighbor}"
1335                )));
1336            }
1337            bytes.extend_from_slice(&neighbor.to_le_bytes());
1338        }
1339        for _ in neighbors.len()..r_max {
1340            bytes.extend_from_slice(&u32::MAX.to_le_bytes());
1341        }
1342    }
1343    encode_u32_slice(&mut bytes, &state.permutation)?;
1344    encode_u32_slice(&mut bytes, state.queue.as_slices().0)?;
1345    encode_u32_slice(&mut bytes, state.queue.as_slices().1)?;
1346    bytes.extend_from_slice(
1347        &u32::try_from(state.seen.len())
1348            .map_err(|_| {
1349                RefinementError::Geometry("checkpoint seen length exceeds u32".to_owned())
1350            })?
1351            .to_le_bytes(),
1352    );
1353    bytes.extend_from_slice(&state.seen);
1354    encode_u32_slice(&mut bytes, &state.entries)?;
1355    bytes.extend_from_slice(&state.best_node.unwrap_or(u32::MAX).to_le_bytes());
1356    bytes.extend_from_slice(&state.best_score.to_bits().to_le_bytes());
1357    if bytes.len() != CHECKPOINT_HEADER_BYTES + graph_bytes + aux_bytes {
1358        return Err(RefinementError::Geometry(
1359            "checkpoint payload length changed while encoding".to_owned(),
1360        ));
1361    }
1362    bytes.extend_from_slice(&xxh3_64(&bytes).to_le_bytes());
1363    Ok(bytes)
1364}
1365
1366fn checkpoint_aux_bytes(state: &RefinementState) -> Result<usize, RefinementError> {
1367    let u32_items = state
1368        .permutation
1369        .len()
1370        .checked_add(state.queue.len())
1371        .and_then(|count| count.checked_add(state.entries.len()))
1372        .ok_or_else(|| RefinementError::Geometry("checkpoint aux count overflow".to_owned()))?;
1373    u32_items
1374        .checked_mul(std::mem::size_of::<u32>())
1375        .and_then(|bytes| bytes.checked_add(state.seen.len()))
1376        .and_then(|bytes| bytes.checked_add(6 * std::mem::size_of::<u32>()))
1377        .and_then(|bytes| bytes.checked_add(std::mem::size_of::<u64>()))
1378        .ok_or_else(|| RefinementError::Geometry("checkpoint aux size overflow".to_owned()))
1379}
1380
1381fn encode_u32_slice(bytes: &mut Vec<u8>, values: &[u32]) -> Result<(), RefinementError> {
1382    bytes.extend_from_slice(
1383        &u32::try_from(values.len())
1384            .map_err(|_| RefinementError::Geometry("checkpoint list exceeds u32".to_owned()))?
1385            .to_le_bytes(),
1386    );
1387    for value in values {
1388        bytes.extend_from_slice(&value.to_le_bytes());
1389    }
1390    Ok(())
1391}
1392
1393fn publish_refinement_checkpoint(
1394    vfs: &dyn Vfs,
1395    policy: DurabilityPolicy,
1396    path: &Path,
1397    bytes: &[u8],
1398) -> Result<(), RefinementError> {
1399    let directory = path.parent().ok_or_else(|| GraphBuildError::CheckpointIo {
1400        path: path.to_path_buf(),
1401        source: std::io::Error::new(
1402            std::io::ErrorKind::InvalidInput,
1403            "refinement checkpoint path has no parent directory",
1404        ),
1405    })?;
1406    let temporary = refinement_checkpoint_temp_path(path);
1407    if let Err(source) = vfs.write(&temporary, bytes) {
1408        let _ = vfs.delete(&temporary);
1409        return Err(GraphBuildError::CheckpointIo {
1410            path: temporary,
1411            source,
1412        }
1413        .into());
1414    }
1415    if let SyncRequirement::Sync(kind) = policy.data_file_sync()
1416        && let Err(source) = vfs.sync(&temporary, kind)
1417    {
1418        let _ = vfs.delete(&temporary);
1419        return Err(GraphBuildError::CheckpointIo {
1420            path: temporary,
1421            source,
1422        }
1423        .into());
1424    }
1425    if let Err(source) = vfs.rename(&temporary, path) {
1426        let _ = vfs.delete(&temporary);
1427        return Err(GraphBuildError::CheckpointIo {
1428            path: path.to_path_buf(),
1429            source,
1430        }
1431        .into());
1432    }
1433    match policy.directory_sync() {
1434        SyncRequirement::Skip => Ok(()),
1435        SyncRequirement::Sync(kind) => vfs.sync(directory, kind).map_err(|source| {
1436            RefinementError::Graph(GraphBuildError::CheckpointIo {
1437                path: directory.to_path_buf(),
1438                source,
1439            })
1440        }),
1441    }
1442}
1443
1444fn refinement_checkpoint_temp_path(path: &Path) -> PathBuf {
1445    let mut temporary = path.as_os_str().to_os_string();
1446    temporary.push(".tmp");
1447    PathBuf::from(temporary)
1448}
1449
1450/// Removes both durable and temporary forms of one refinement checkpoint.
1451pub fn remove_refinement_checkpoint(vfs: &dyn Vfs, path: &Path) -> Result<(), RefinementError> {
1452    for candidate in [refinement_checkpoint_temp_path(path), path.to_path_buf()] {
1453        match vfs.delete(&candidate) {
1454            Ok(()) => {}
1455            Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
1456            Err(source) => {
1457                return Err(GraphBuildError::CheckpointIo {
1458                    path: candidate,
1459                    source,
1460                }
1461                .into());
1462            }
1463        }
1464    }
1465    Ok(())
1466}
1467
1468fn decode_refinement_checkpoint(
1469    bytes: &[u8],
1470    reader: &SegmentReader,
1471    pass: RefinementPass,
1472    params: GraphParams,
1473    seed: u64,
1474    accounting: Option<&std::sync::Arc<crate::lifecycle::stats::Accounting>>,
1475) -> Result<RefinementState, RefinementError> {
1476    validate_refinement_checkpoint(bytes)?;
1477    let source_graph = reader.graph_node_blocks()?;
1478    // Store generation is provenance, not resume identity: live ingest advances per
1479    // acknowledged batch, replay advances per unabsorbed WAL record, and maintenance
1480    // can publish the counter without absorbing WAL, so it is not stable across reopen.
1481    if checkpoint_u8(bytes, 10, "pass")? != pass as u8
1482        || checkpoint_u8(bytes, 12, "target degree")? != params.r_target()
1483        || checkpoint_u8(bytes, 13, "maximum degree")? != params.r_max()
1484        || checkpoint_u32(bytes, 16, "checkpoint batch")? != params.checkpoint_batch_rows()
1485        || checkpoint_u32(bytes, 20, "node count")? != reader.meta().row_count
1486        || checkpoint_u32(bytes, 24, "dimensions")? != reader.meta().dims
1487        || checkpoint_u64(bytes, 36, "seed")? != seed
1488        || checkpoint_u32(bytes, 44, "refinement alpha")? != params.alpha_refine().to_bits()
1489        || bytes.get(48..64) != Some(reader.meta().id.as_bytes().as_slice())
1490        || checkpoint_u16(bytes, 64, "source passes")? != source_graph.refinement_passes().bits()
1491        || source_graph.layout().max_degree() != params.r_max()
1492    {
1493        return Err(checkpoint_corrupt(
1494            "checkpoint identity or refinement parameters do not match input",
1495        ));
1496    }
1497    let phase = decode_refinement_phase(checkpoint_u8(bytes, 11, "phase")?)?;
1498    let next_row = checkpoint_u32(bytes, 28, "next row")?;
1499    let work_rows = checkpoint_u64(bytes, 76, "work rows")?;
1500    let graph_bytes = usize::try_from(checkpoint_u64(bytes, 84, "graph bytes")?)
1501        .map_err(|_| checkpoint_corrupt("graph bytes exceed usize"))?;
1502    let graph_start = CHECKPOINT_HEADER_BYTES;
1503    let aux_start = graph_start
1504        .checked_add(graph_bytes)
1505        .ok_or_else(|| checkpoint_corrupt("aux offset overflow"))?;
1506    let node_count = reader.meta().row_count as usize;
1507    let degrees_start = graph_start
1508        .checked_add(node_count)
1509        .ok_or_else(|| checkpoint_corrupt("degree offset overflow"))?;
1510    let slots_start = degrees_start
1511        .checked_add(node_count)
1512        .ok_or_else(|| checkpoint_corrupt("slot offset overflow"))?;
1513    let mut flags = Vec::with_capacity(node_count);
1514    flags.extend_from_slice(
1515        bytes
1516            .get(graph_start..degrees_start)
1517            .ok_or_else(|| checkpoint_corrupt("flag bytes are truncated"))?,
1518    );
1519    let mut neighbors = Vec::with_capacity(node_count);
1520    let r_max = usize::from(params.r_max());
1521    for owner in 0..node_count {
1522        let degree = usize::from(checkpoint_u8(bytes, degrees_start + owner, "degree")?);
1523        let row_start = owner
1524            .checked_mul(r_max)
1525            .and_then(|slot| slot.checked_mul(std::mem::size_of::<u32>()))
1526            .and_then(|offset| slots_start.checked_add(offset))
1527            .ok_or_else(|| checkpoint_corrupt("neighbor row offset overflow"))?;
1528        let mut row = Vec::with_capacity(r_max);
1529        for slot in 0..degree {
1530            let offset = slot
1531                .checked_mul(std::mem::size_of::<u32>())
1532                .and_then(|offset| row_start.checked_add(offset))
1533                .ok_or_else(|| checkpoint_corrupt("neighbor offset overflow"))?;
1534            row.push(checkpoint_u32(bytes, offset, "neighbor")?);
1535        }
1536        neighbors.push(row);
1537    }
1538    let mut cursor = aux_start;
1539    let permutation = decode_u32_list(bytes, &mut cursor, node_count, "permutation")?;
1540    let queue_first = decode_u32_list(bytes, &mut cursor, node_count, "queue first")?;
1541    let queue_second = decode_u32_list(bytes, &mut cursor, node_count, "queue second")?;
1542    let seen_len = usize::try_from(checkpoint_u32(bytes, cursor, "seen length")?)
1543        .map_err(|_| checkpoint_corrupt("seen length exceeds usize"))?;
1544    cursor = cursor
1545        .checked_add(4)
1546        .ok_or_else(|| checkpoint_corrupt("seen offset overflow"))?;
1547    let seen_end = cursor
1548        .checked_add(seen_len)
1549        .ok_or_else(|| checkpoint_corrupt("seen end overflow"))?;
1550    let seen = bytes
1551        .get(cursor..seen_end)
1552        .ok_or_else(|| checkpoint_corrupt("seen bytes are truncated"))?
1553        .to_vec();
1554    cursor = seen_end;
1555    let entries = decode_u32_list(bytes, &mut cursor, ENTRY_POINT_COUNT, "entries")?;
1556    let best = checkpoint_u32(bytes, cursor, "best node")?;
1557    cursor = cursor
1558        .checked_add(4)
1559        .ok_or_else(|| checkpoint_corrupt("best score offset overflow"))?;
1560    let best_score = f64::from_bits(checkpoint_u64(bytes, cursor, "best score")?);
1561    let mut queue = VecDeque::with_capacity(queue_first.len() + queue_second.len());
1562    queue.extend(queue_first);
1563    queue.extend(queue_second);
1564    let layout = source_graph.layout();
1565    let source_passes = RefinementPasses::from_bits(checkpoint_u16(bytes, 64, "source passes")?)
1566        .ok_or_else(|| checkpoint_corrupt("source pass bits are unknown"))?;
1567    let mut state = RefinementState {
1568        graph: OwnedGraph {
1569            layout,
1570            passes: source_passes,
1571            flags,
1572            neighbors,
1573        },
1574        phase,
1575        next_row,
1576        work_rows,
1577        permutation,
1578        queue,
1579        seen,
1580        entries,
1581        best_node: (best != u32::MAX).then_some(best),
1582        best_score,
1583        scratch_bytes: refinement_scratch_bytes(pass, params)?,
1584        memory: match accounting {
1585            Some(accounting) => Some(AccountedCounter::new(
1586                accounting,
1587                AllocationComponent::Temporary,
1588            )?),
1589            None => None,
1590        },
1591    };
1592    state.update_memory_charge()?;
1593    Ok(state)
1594}
1595
1596fn refinement_scratch_bytes(
1597    pass: RefinementPass,
1598    params: GraphParams,
1599) -> Result<usize, RefinementError> {
1600    let r_max = usize::from(params.r_max());
1601    match pass {
1602        RefinementPass::Renumber => r_max
1603            .checked_mul(std::mem::size_of::<u32>())
1604            .ok_or_else(|| RefinementError::Geometry("renumber scratch overflow".to_owned())),
1605        RefinementPass::AlphaReprune => {
1606            let candidates = r_max
1607                .checked_mul(r_max)
1608                .and_then(|count| count.checked_add(r_max))
1609                .ok_or_else(|| RefinementError::Geometry("re-prune scratch overflow".to_owned()))?;
1610            candidates
1611                .checked_mul(24)
1612                .and_then(|bytes| {
1613                    usize::from(params.r_target())
1614                        .checked_mul(std::mem::size_of::<u32>())
1615                        .and_then(|selected| bytes.checked_add(selected))
1616                })
1617                .ok_or_else(|| RefinementError::Geometry("re-prune scratch overflow".to_owned()))
1618        }
1619        RefinementPass::SeedRefit => Ok(0),
1620        RefinementPass::NeighborReorder => r_max
1621            .checked_mul(std::mem::size_of::<(u32, f64)>())
1622            .ok_or_else(|| RefinementError::Geometry("reorder scratch overflow".to_owned())),
1623        RefinementPass::ConnectivityRepair => Ok(0),
1624    }
1625}
1626
1627fn decode_u32_list(
1628    bytes: &[u8],
1629    cursor: &mut usize,
1630    maximum: usize,
1631    field: &str,
1632) -> Result<Vec<u32>, RefinementError> {
1633    let count = usize::try_from(checkpoint_u32(bytes, *cursor, field)?)
1634        .map_err(|_| checkpoint_corrupt(format!("{field} count exceeds usize")))?;
1635    if count > maximum {
1636        return Err(checkpoint_corrupt(format!(
1637            "{field} count {count} exceeds {maximum}"
1638        )));
1639    }
1640    *cursor = cursor
1641        .checked_add(4)
1642        .ok_or_else(|| checkpoint_corrupt(format!("{field} offset overflow")))?;
1643    let mut values = Vec::with_capacity(count);
1644    for _ in 0..count {
1645        values.push(checkpoint_u32(bytes, *cursor, field)?);
1646        *cursor = cursor
1647            .checked_add(4)
1648            .ok_or_else(|| checkpoint_corrupt(format!("{field} offset overflow")))?;
1649    }
1650    Ok(values)
1651}
1652
1653/// Validates arbitrary refinement checkpoint bytes without allocating pass state.
1654pub fn validate_refinement_checkpoint(bytes: &[u8]) -> Result<(), RefinementError> {
1655    let minimum = CHECKPOINT_HEADER_BYTES
1656        .checked_add(CHECKPOINT_CHECKSUM_BYTES)
1657        .ok_or_else(|| checkpoint_corrupt("minimum checkpoint length overflow"))?;
1658    if bytes.len() < minimum {
1659        return Err(checkpoint_corrupt(format!(
1660            "truncated: need at least {minimum} bytes, got {}",
1661            bytes.len()
1662        )));
1663    }
1664    if bytes.get(..8) != Some(CHECKPOINT_MAGIC.as_slice()) {
1665        return Err(checkpoint_corrupt("magic does not match ZEREFCP1"));
1666    }
1667    let checksum_start = bytes
1668        .len()
1669        .checked_sub(CHECKPOINT_CHECKSUM_BYTES)
1670        .ok_or_else(|| checkpoint_corrupt("checksum offset underflow"))?;
1671    let stored_checksum = checkpoint_u64(bytes, checksum_start, "checksum")?;
1672    let prefix = bytes
1673        .get(..checksum_start)
1674        .ok_or_else(|| checkpoint_corrupt("checksummed prefix is unavailable"))?;
1675    let actual_checksum = xxh3_64(prefix);
1676    if stored_checksum != actual_checksum {
1677        return Err(checkpoint_corrupt(format!(
1678            "xxh3-64 expected {stored_checksum:#018x}, computed {actual_checksum:#018x}"
1679        )));
1680    }
1681    if checkpoint_u16(bytes, 8, "version")? != CHECKPOINT_VERSION {
1682        return Err(checkpoint_corrupt("checkpoint version is unsupported"));
1683    }
1684    let pass = decode_refinement_pass(checkpoint_u8(bytes, 10, "pass")?)?;
1685    let phase = decode_refinement_phase(checkpoint_u8(bytes, 11, "phase")?)?;
1686    let r_target = checkpoint_u8(bytes, 12, "target degree")?;
1687    let r_max = checkpoint_u8(bytes, 13, "maximum degree")?;
1688    let batch = checkpoint_u32(bytes, 16, "checkpoint batch")?;
1689    let node_count = checkpoint_u32(bytes, 20, "node count")?;
1690    let dimensions = checkpoint_u32(bytes, 24, "dimensions")?;
1691    let next_row = checkpoint_u32(bytes, 28, "next row")?;
1692    let alpha = f32::from_bits(checkpoint_u32(bytes, 44, "refinement alpha")?);
1693    let source_passes = checkpoint_u16(bytes, 64, "source passes")?;
1694    if checkpoint_u16(bytes, 14, "reserved degree bytes")? != 0
1695        || checkpoint_u32(bytes, 32, "reserved row bytes")? != 0
1696        || checkpoint_u16(bytes, 66, "reserved pass bytes")? != 0
1697        || checkpoint_u32(bytes, 100, "reserved tail bytes")? != 0
1698    {
1699        return Err(checkpoint_corrupt("reserved header bytes are nonzero"));
1700    }
1701    if r_target == 0
1702        || r_target > r_max
1703        || batch == 0
1704        || node_count == 0
1705        || dimensions == 0
1706        || next_row > node_count
1707        || !alpha.is_finite()
1708        || alpha < 1.0
1709        || RefinementPasses::from_bits(source_passes).is_none()
1710        || source_passes & pass.bit() != 0
1711    {
1712        return Err(checkpoint_corrupt(
1713            "stored shape, alpha, progress, or pass record is invalid",
1714        ));
1715    }
1716    let phase_matches_pass = match pass {
1717        RefinementPass::SeedRefit => matches!(
1718            phase,
1719            RefinementPhase::SeedMedoid | RefinementPhase::SeedSpread | RefinementPhase::Complete
1720        ),
1721        RefinementPass::ConnectivityRepair => matches!(
1722            phase,
1723            RefinementPhase::ConnectivityTraverse
1724                | RefinementPhase::ConnectivityDensify
1725                | RefinementPhase::Complete
1726        ),
1727        RefinementPass::Renumber
1728        | RefinementPass::AlphaReprune
1729        | RefinementPass::NeighborReorder => {
1730            matches!(phase, RefinementPhase::Rows | RefinementPhase::Complete)
1731        }
1732    };
1733    if !phase_matches_pass {
1734        return Err(checkpoint_corrupt(
1735            "refinement phase does not match its pass",
1736        ));
1737    }
1738    let node_count_usize = node_count as usize;
1739    let expected_graph = node_count_usize
1740        .checked_mul(usize::from(r_max))
1741        .and_then(|slots| slots.checked_mul(std::mem::size_of::<u32>()))
1742        .and_then(|bytes| bytes.checked_add(node_count_usize.checked_mul(2)?))
1743        .ok_or_else(|| checkpoint_corrupt("graph payload geometry overflow"))?;
1744    let graph_bytes = usize::try_from(checkpoint_u64(bytes, 84, "graph bytes")?)
1745        .map_err(|_| checkpoint_corrupt("graph payload exceeds usize"))?;
1746    let aux_bytes = usize::try_from(checkpoint_u64(bytes, 92, "aux bytes")?)
1747        .map_err(|_| checkpoint_corrupt("aux payload exceeds usize"))?;
1748    if graph_bytes != expected_graph {
1749        return Err(checkpoint_corrupt(format!(
1750            "graph payload {graph_bytes}, expected {expected_graph}"
1751        )));
1752    }
1753    let expected_len = CHECKPOINT_HEADER_BYTES
1754        .checked_add(graph_bytes)
1755        .and_then(|length| length.checked_add(aux_bytes))
1756        .and_then(|length| length.checked_add(CHECKPOINT_CHECKSUM_BYTES))
1757        .ok_or_else(|| checkpoint_corrupt("checkpoint length overflow"))?;
1758    if bytes.len() != expected_len {
1759        return Err(checkpoint_corrupt(format!(
1760            "checkpoint length {}, expected {expected_len}",
1761            bytes.len()
1762        )));
1763    }
1764    let flags_start = CHECKPOINT_HEADER_BYTES;
1765    let degrees_start = flags_start
1766        .checked_add(node_count_usize)
1767        .ok_or_else(|| checkpoint_corrupt("degree offset overflow"))?;
1768    let slots_start = degrees_start
1769        .checked_add(node_count_usize)
1770        .ok_or_else(|| checkpoint_corrupt("slot offset overflow"))?;
1771    for owner in 0..node_count_usize {
1772        let flags = checkpoint_u8(bytes, flags_start + owner, "node flags")?;
1773        if flags & !(ENTRY_FLAG | HUB_FLAG) != 0 {
1774            return Err(checkpoint_corrupt(format!(
1775                "node {owner} has unknown flags {flags:#04x}"
1776            )));
1777        }
1778        let degree = usize::from(checkpoint_u8(bytes, degrees_start + owner, "degree")?);
1779        if degree > usize::from(r_max) {
1780            return Err(checkpoint_corrupt(format!(
1781                "node {owner} degree {degree} exceeds {r_max}"
1782            )));
1783        }
1784        let row_start = owner
1785            .checked_mul(usize::from(r_max))
1786            .and_then(|slot| slot.checked_mul(std::mem::size_of::<u32>()))
1787            .and_then(|offset| slots_start.checked_add(offset))
1788            .ok_or_else(|| checkpoint_corrupt("neighbor row offset overflow"))?;
1789        for slot in 0..usize::from(r_max) {
1790            let offset = slot
1791                .checked_mul(std::mem::size_of::<u32>())
1792                .and_then(|offset| row_start.checked_add(offset))
1793                .ok_or_else(|| checkpoint_corrupt("neighbor offset overflow"))?;
1794            let neighbor = checkpoint_u32(bytes, offset, "neighbor")?;
1795            if slot < degree {
1796                if neighbor >= node_count || neighbor as usize == owner {
1797                    return Err(checkpoint_corrupt(format!(
1798                        "node {owner} has invalid active neighbor {neighbor}"
1799                    )));
1800                }
1801                for prior in 0..slot {
1802                    let prior_offset = prior
1803                        .checked_mul(std::mem::size_of::<u32>())
1804                        .and_then(|offset| row_start.checked_add(offset))
1805                        .ok_or_else(|| checkpoint_corrupt("prior neighbor offset overflow"))?;
1806                    if checkpoint_u32(bytes, prior_offset, "prior neighbor")? == neighbor {
1807                        return Err(checkpoint_corrupt(format!(
1808                            "node {owner} repeats neighbor {neighbor}"
1809                        )));
1810                    }
1811                }
1812            } else if neighbor != u32::MAX {
1813                return Err(checkpoint_corrupt(format!(
1814                    "node {owner} inactive slot {slot} is not the sentinel"
1815                )));
1816            }
1817        }
1818    }
1819    validate_refinement_aux(
1820        bytes,
1821        slots_start + node_count_usize * usize::from(r_max) * 4,
1822        checksum_start,
1823        node_count_usize,
1824        pass,
1825        phase,
1826    )
1827}
1828
1829fn validate_refinement_aux(
1830    bytes: &[u8],
1831    mut cursor: usize,
1832    end: usize,
1833    node_count: usize,
1834    pass: RefinementPass,
1835    phase: RefinementPhase,
1836) -> Result<(), RefinementError> {
1837    let permutation = validate_u32_list(bytes, &mut cursor, node_count, node_count, "permutation")?;
1838    let queue_first = validate_u32_list(bytes, &mut cursor, node_count, node_count, "queue first")?;
1839    let queue_second =
1840        validate_u32_list(bytes, &mut cursor, node_count, node_count, "queue second")?;
1841    let seen_len = usize::try_from(checkpoint_u32(bytes, cursor, "seen length")?)
1842        .map_err(|_| checkpoint_corrupt("seen length exceeds usize"))?;
1843    cursor = cursor
1844        .checked_add(4)
1845        .ok_or_else(|| checkpoint_corrupt("seen offset overflow"))?;
1846    if seen_len != 0 && seen_len != node_count {
1847        return Err(checkpoint_corrupt(
1848            "seen length is neither zero nor node count",
1849        ));
1850    }
1851    let seen_end = cursor
1852        .checked_add(seen_len)
1853        .ok_or_else(|| checkpoint_corrupt("seen end overflow"))?;
1854    if bytes
1855        .get(cursor..seen_end)
1856        .ok_or_else(|| checkpoint_corrupt("seen bytes are truncated"))?
1857        .iter()
1858        .any(|marker| *marker > 1)
1859    {
1860        return Err(checkpoint_corrupt(
1861            "seen bytes contain a non-boolean marker",
1862        ));
1863    }
1864    cursor = seen_end;
1865    let entries = validate_u32_list(bytes, &mut cursor, ENTRY_POINT_COUNT, node_count, "entries")?;
1866    let best_node = checkpoint_u32(bytes, cursor, "best node")?;
1867    cursor = cursor
1868        .checked_add(4)
1869        .ok_or_else(|| checkpoint_corrupt("best score offset overflow"))?;
1870    let best_score = f64::from_bits(checkpoint_u64(bytes, cursor, "best score")?);
1871    cursor = cursor
1872        .checked_add(8)
1873        .ok_or_else(|| checkpoint_corrupt("aux end overflow"))?;
1874    if cursor != end {
1875        return Err(checkpoint_corrupt(
1876            "aux payload has trailing or missing bytes",
1877        ));
1878    }
1879    if best_node != u32::MAX && best_node as usize >= node_count {
1880        return Err(checkpoint_corrupt("best seed node is out of range"));
1881    }
1882    if best_score.is_nan() {
1883        return Err(checkpoint_corrupt("best seed score is NaN"));
1884    }
1885    if pass == RefinementPass::Renumber {
1886        if permutation > node_count || queue_first + queue_second > node_count || entries != 0 {
1887            return Err(checkpoint_corrupt("renumber auxiliary state is invalid"));
1888        }
1889    } else if pass == RefinementPass::ConnectivityRepair {
1890        if permutation != 0
1891            || queue_first + queue_second > node_count
1892            || seen_len != node_count
1893            || entries != 0
1894        {
1895            return Err(checkpoint_corrupt(
1896                "connectivity-repair auxiliary state is invalid",
1897            ));
1898        }
1899    } else if permutation != 0 || queue_first != 0 || queue_second != 0 || seen_len != 0 {
1900        return Err(checkpoint_corrupt(
1901            "non-renumber pass carries renumber state",
1902        ));
1903    }
1904    if pass != RefinementPass::SeedRefit && (entries != 0 || best_node != u32::MAX) {
1905        return Err(checkpoint_corrupt("non-seed pass carries seed state"));
1906    }
1907    if phase == RefinementPhase::Complete && pass == RefinementPass::SeedRefit && entries == 0 {
1908        return Err(checkpoint_corrupt("completed seed pass has no entries"));
1909    }
1910    Ok(())
1911}
1912
1913fn validate_u32_list(
1914    bytes: &[u8],
1915    cursor: &mut usize,
1916    maximum_count: usize,
1917    node_count: usize,
1918    field: &str,
1919) -> Result<usize, RefinementError> {
1920    let count = usize::try_from(checkpoint_u32(bytes, *cursor, field)?)
1921        .map_err(|_| checkpoint_corrupt(format!("{field} count exceeds usize")))?;
1922    if count > maximum_count {
1923        return Err(checkpoint_corrupt(format!(
1924            "{field} count {count} exceeds {maximum_count}"
1925        )));
1926    }
1927    *cursor = cursor
1928        .checked_add(4)
1929        .ok_or_else(|| checkpoint_corrupt(format!("{field} offset overflow")))?;
1930    for index in 0..count {
1931        let offset = index
1932            .checked_mul(4)
1933            .and_then(|offset| cursor.checked_add(offset))
1934            .ok_or_else(|| checkpoint_corrupt(format!("{field} offset overflow")))?;
1935        let node = checkpoint_u32(bytes, offset, field)?;
1936        if node as usize >= node_count {
1937            return Err(checkpoint_corrupt(format!(
1938                "{field} node {node} is out of range"
1939            )));
1940        }
1941        for prior in 0..index {
1942            let prior_offset = prior
1943                .checked_mul(4)
1944                .and_then(|offset| cursor.checked_add(offset))
1945                .ok_or_else(|| checkpoint_corrupt(format!("{field} prior offset overflow")))?;
1946            if checkpoint_u32(bytes, prior_offset, field)? == node {
1947                return Err(checkpoint_corrupt(format!("{field} repeats node {node}")));
1948            }
1949        }
1950    }
1951    *cursor = cursor
1952        .checked_add(
1953            count
1954                .checked_mul(4)
1955                .ok_or_else(|| checkpoint_corrupt(format!("{field} byte count overflow")))?,
1956        )
1957        .ok_or_else(|| checkpoint_corrupt(format!("{field} end overflow")))?;
1958    Ok(count)
1959}
1960
1961fn check_refinement_cancellation(
1962    cancellation: &QueryCancellation<'_>,
1963) -> Result<(), RefinementError> {
1964    match cancellation.check_graph() {
1965        Ok(()) => Ok(()),
1966        Err(ScanError::Cancelled { .. }) => {
1967            Err(GraphBuildError::Cancelled { partial: false }.into())
1968        }
1969        Err(ScanError::Timeout { .. }) => Err(GraphBuildError::Timeout { partial: false }.into()),
1970        Err(ScanError::ReadCancelled { .. }) => {
1971            Err(GraphBuildError::ReadCancelled { partial: false }.into())
1972        }
1973        Err(error) => Err(RefinementError::Geometry(format!(
1974            "refinement cancellation seam returned non-cancellation error: {error}"
1975        ))),
1976    }
1977}
1978
1979fn refinement_phase_code(phase: RefinementPhase) -> u8 {
1980    match phase {
1981        RefinementPhase::Rows => 1,
1982        RefinementPhase::SeedMedoid => 2,
1983        RefinementPhase::SeedSpread => 3,
1984        RefinementPhase::Complete => 4,
1985        RefinementPhase::ConnectivityTraverse => 5,
1986        RefinementPhase::ConnectivityDensify => 6,
1987    }
1988}
1989
1990fn decode_refinement_phase(value: u8) -> Result<RefinementPhase, RefinementError> {
1991    match value {
1992        1 => Ok(RefinementPhase::Rows),
1993        2 => Ok(RefinementPhase::SeedMedoid),
1994        3 => Ok(RefinementPhase::SeedSpread),
1995        4 => Ok(RefinementPhase::Complete),
1996        5 => Ok(RefinementPhase::ConnectivityTraverse),
1997        6 => Ok(RefinementPhase::ConnectivityDensify),
1998        _ => Err(checkpoint_corrupt(format!(
1999            "unknown refinement phase {value}"
2000        ))),
2001    }
2002}
2003
2004fn decode_refinement_pass(value: u8) -> Result<RefinementPass, RefinementError> {
2005    match value {
2006        0 => Ok(RefinementPass::Renumber),
2007        1 => Ok(RefinementPass::AlphaReprune),
2008        2 => Ok(RefinementPass::SeedRefit),
2009        3 => Ok(RefinementPass::NeighborReorder),
2010        4 => Ok(RefinementPass::ConnectivityRepair),
2011        _ => Err(checkpoint_corrupt(format!(
2012            "unknown refinement pass {value}"
2013        ))),
2014    }
2015}
2016
2017fn checkpoint_corrupt(detail: impl Into<String>) -> RefinementError {
2018    RefinementError::Graph(GraphBuildError::CheckpointCorrupt(detail.into()))
2019}
2020
2021fn checkpoint_u8(bytes: &[u8], offset: usize, field: &str) -> Result<u8, RefinementError> {
2022    bytes
2023        .get(offset)
2024        .copied()
2025        .ok_or_else(|| checkpoint_corrupt(format!("{field} is truncated at byte {offset}")))
2026}
2027
2028fn checkpoint_u16(bytes: &[u8], offset: usize, field: &str) -> Result<u16, RefinementError> {
2029    let end = offset
2030        .checked_add(2)
2031        .ok_or_else(|| checkpoint_corrupt(format!("{field} offset overflow")))?;
2032    let array = bytes
2033        .get(offset..end)
2034        .ok_or_else(|| checkpoint_corrupt(format!("{field} is truncated")))?
2035        .try_into()
2036        .map_err(|_| checkpoint_corrupt(format!("{field} width is invalid")))?;
2037    Ok(u16::from_le_bytes(array))
2038}
2039
2040fn checkpoint_u32(bytes: &[u8], offset: usize, field: &str) -> Result<u32, RefinementError> {
2041    let end = offset
2042        .checked_add(4)
2043        .ok_or_else(|| checkpoint_corrupt(format!("{field} offset overflow")))?;
2044    let array = bytes
2045        .get(offset..end)
2046        .ok_or_else(|| checkpoint_corrupt(format!("{field} is truncated")))?
2047        .try_into()
2048        .map_err(|_| checkpoint_corrupt(format!("{field} width is invalid")))?;
2049    Ok(u32::from_le_bytes(array))
2050}
2051
2052fn checkpoint_u64(bytes: &[u8], offset: usize, field: &str) -> Result<u64, RefinementError> {
2053    let end = offset
2054        .checked_add(8)
2055        .ok_or_else(|| checkpoint_corrupt(format!("{field} offset overflow")))?;
2056    let array = bytes
2057        .get(offset..end)
2058        .ok_or_else(|| checkpoint_corrupt(format!("{field} is truncated")))?
2059        .try_into()
2060        .map_err(|_| checkpoint_corrupt(format!("{field} width is invalid")))?;
2061    Ok(u64::from_le_bytes(array))
2062}
2063
2064/// Computes one complete deterministic refinement artifact without publishing it.
2065pub fn refine_graph(
2066    reader: &SegmentReader,
2067    pass: RefinementPass,
2068    params: GraphParams,
2069    seed: u64,
2070) -> Result<RefinementArtifact, RefinementError> {
2071    let source = reader.graph_node_blocks()?;
2072    if source.refinement_passes().contains(pass) {
2073        return Err(RefinementError::AlreadyApplied(pass));
2074    }
2075    if source.node_count() != reader.meta().row_count {
2076        return Err(RefinementError::Geometry(format!(
2077            "graph rows {}, segment rows {}",
2078            source.node_count(),
2079            reader.meta().row_count
2080        )));
2081    }
2082    let mut graph = OwnedGraph::read(source)?;
2083    let rescore = reader.rescore_f32()?;
2084    let dimensions = reader.meta().dims as usize;
2085    let node_count = reader.meta().row_count;
2086    let (old_to_new, new_to_old) = match pass {
2087        RefinementPass::Renumber => {
2088            let new_to_old = gorder_lite(&graph)?;
2089            let old_to_new = invert_permutation(&new_to_old, node_count)?;
2090            apply_renumber(&mut graph, &new_to_old, &old_to_new)?;
2091            (Some(old_to_new), Some(new_to_old))
2092        }
2093        RefinementPass::AlphaReprune => {
2094            alpha_reprune(&mut graph, rescore, dimensions, node_count, params)?;
2095            (None, None)
2096        }
2097        RefinementPass::SeedRefit => {
2098            seed_refit(&mut graph, rescore, dimensions, node_count)?;
2099            (None, None)
2100        }
2101        RefinementPass::NeighborReorder => {
2102            neighbor_reorder(&mut graph, reader, node_count, seed)?;
2103            (None, None)
2104        }
2105        RefinementPass::ConnectivityRepair => {
2106            connectivity_repair(&mut graph, params)?;
2107            (None, None)
2108        }
2109    };
2110    graph.passes = graph.passes.with(pass);
2111    let encoded_region = encode_owned_graph(reader, &graph, new_to_old.as_deref())?;
2112    let work_rows_completed = match pass {
2113        RefinementPass::SeedRefit => u64::from(node_count)
2114            .checked_mul(u64::from(node_count.min(ENTRY_POINT_COUNT as u32)))
2115            .ok_or_else(|| RefinementError::Geometry("seed work rows overflow".to_owned()))?,
2116        RefinementPass::ConnectivityRepair => u64::from(node_count)
2117            .checked_mul(2)
2118            .ok_or_else(|| RefinementError::Geometry("repair work rows overflow".to_owned()))?,
2119        RefinementPass::Renumber
2120        | RefinementPass::AlphaReprune
2121        | RefinementPass::NeighborReorder => u64::from(node_count),
2122    };
2123    Ok(RefinementArtifact {
2124        pass,
2125        encoded_region,
2126        old_to_new,
2127        new_to_old,
2128        work_rows_completed,
2129    })
2130}
2131
2132fn gorder_lite(graph: &OwnedGraph) -> Result<Vec<u32>, RefinementError> {
2133    let node_count = u32::try_from(graph.neighbors.len())
2134        .map_err(|_| RefinementError::Geometry("graph rows exceed u32".to_owned()))?;
2135    let mut roots = graph
2136        .flags
2137        .iter()
2138        .enumerate()
2139        .filter_map(|(node, flags)| {
2140            (*flags & ENTRY_FLAG != 0)
2141                .then(|| u32::try_from(node).ok())
2142                .flatten()
2143        })
2144        .collect::<Vec<_>>();
2145    if roots.is_empty() {
2146        return Err(RefinementError::Geometry(
2147            "renumbering requires at least one persisted entry seed".to_owned(),
2148        ));
2149    }
2150    sort_by_degree(&mut roots, graph)?;
2151    let mut queued = vec![false; graph.neighbors.len()];
2152    let mut queue = VecDeque::new();
2153    for root in roots {
2154        mark_and_enqueue(root, &mut queued, &mut queue)?;
2155    }
2156    let mut order = Vec::with_capacity(graph.neighbors.len());
2157    loop {
2158        while let Some(owner) = queue.pop_front() {
2159            order.push(owner);
2160            let owner_index = node_index(owner, node_count)?;
2161            let mut candidates = graph
2162                .neighbors
2163                .get(owner_index)
2164                .ok_or_else(|| {
2165                    RefinementError::Geometry(format!("neighbors for node {owner} are unavailable"))
2166                })?
2167                .iter()
2168                .copied()
2169                .filter(|neighbor| {
2170                    usize::try_from(*neighbor)
2171                        .ok()
2172                        .and_then(|index| queued.get(index))
2173                        .is_some_and(|seen| !*seen)
2174                })
2175                .collect::<Vec<_>>();
2176            sort_by_degree(&mut candidates, graph)?;
2177            for candidate in candidates {
2178                mark_and_enqueue(candidate, &mut queued, &mut queue)?;
2179            }
2180        }
2181        let Some(unseen) = queued.iter().position(|seen| !*seen) else {
2182            break;
2183        };
2184        let root = u32::try_from(unseen)
2185            .map_err(|_| RefinementError::Geometry("unseen row exceeds u32".to_owned()))?;
2186        mark_and_enqueue(root, &mut queued, &mut queue)?;
2187    }
2188    if order.len() != graph.neighbors.len() {
2189        return Err(RefinementError::Geometry(format!(
2190            "renumbering emitted {} rows, expected {}",
2191            order.len(),
2192            graph.neighbors.len()
2193        )));
2194    }
2195    Ok(order)
2196}
2197
2198fn sort_by_degree(nodes: &mut [u32], graph: &OwnedGraph) -> Result<(), RefinementError> {
2199    for node in nodes.iter().copied() {
2200        let _ = graph.neighbors.get(node as usize).ok_or_else(|| {
2201            RefinementError::Geometry(format!("degree for node {node} is unavailable"))
2202        })?;
2203    }
2204    nodes.sort_unstable_by(|left, right| {
2205        let left_degree = graph.neighbors.get(*left as usize).map_or(0, Vec::len);
2206        let right_degree = graph.neighbors.get(*right as usize).map_or(0, Vec::len);
2207        right_degree.cmp(&left_degree).then_with(|| left.cmp(right))
2208    });
2209    Ok(())
2210}
2211
2212fn mark_and_enqueue(
2213    node: u32,
2214    queued: &mut [bool],
2215    queue: &mut VecDeque<u32>,
2216) -> Result<(), RefinementError> {
2217    let marker = queued.get_mut(node as usize).ok_or_else(|| {
2218        RefinementError::Geometry(format!("renumber queue node {node} is out of range"))
2219    })?;
2220    if !*marker {
2221        *marker = true;
2222        queue.push_back(node);
2223    }
2224    Ok(())
2225}
2226
2227fn invert_permutation(new_to_old: &[u32], node_count: u32) -> Result<Vec<u32>, RefinementError> {
2228    if new_to_old.len() != node_count as usize {
2229        return Err(RefinementError::Geometry(
2230            "renumber permutation length differs from row count".to_owned(),
2231        ));
2232    }
2233    let mut old_to_new = vec![u32::MAX; new_to_old.len()];
2234    for (new, old) in new_to_old.iter().copied().enumerate() {
2235        let slot = old_to_new
2236            .get_mut(node_index(old, node_count)?)
2237            .ok_or_else(|| {
2238                RefinementError::Geometry(format!("permutation row {old} is unavailable"))
2239            })?;
2240        if *slot != u32::MAX {
2241            return Err(RefinementError::Geometry(format!(
2242                "renumber permutation repeats row {old}"
2243            )));
2244        }
2245        *slot = u32::try_from(new)
2246            .map_err(|_| RefinementError::Geometry("new row id exceeds u32".to_owned()))?;
2247    }
2248    if old_to_new.contains(&u32::MAX) {
2249        return Err(RefinementError::Geometry(
2250            "renumber permutation omits a row".to_owned(),
2251        ));
2252    }
2253    Ok(old_to_new)
2254}
2255
2256fn apply_renumber(
2257    graph: &mut OwnedGraph,
2258    new_to_old: &[u32],
2259    old_to_new: &[u32],
2260) -> Result<(), RefinementError> {
2261    let node_count = u32::try_from(graph.neighbors.len())
2262        .map_err(|_| RefinementError::Geometry("graph rows exceed u32".to_owned()))?;
2263    let mut flags = Vec::with_capacity(new_to_old.len());
2264    let mut neighbors = Vec::with_capacity(new_to_old.len());
2265    for old in new_to_old {
2266        let old_index = node_index(*old, node_count)?;
2267        flags.push(*graph.flags.get(old_index).ok_or_else(|| {
2268            RefinementError::Geometry(format!("flags for old row {old} are unavailable"))
2269        })?);
2270        let remapped = graph
2271            .neighbors
2272            .get(old_index)
2273            .ok_or_else(|| {
2274                RefinementError::Geometry(format!("neighbors for old row {old} are unavailable"))
2275            })?
2276            .iter()
2277            .map(|neighbor| {
2278                old_to_new
2279                    .get(node_index(*neighbor, node_count)?)
2280                    .copied()
2281                    .ok_or_else(|| {
2282                        RefinementError::Geometry(format!(
2283                            "new id for old neighbor {neighbor} is unavailable"
2284                        ))
2285                    })
2286            })
2287            .collect::<Result<Vec<_>, _>>()?;
2288        neighbors.push(remapped);
2289    }
2290    graph.flags = flags;
2291    graph.neighbors = neighbors;
2292    Ok(())
2293}
2294
2295fn alpha_reprune(
2296    graph: &mut OwnedGraph,
2297    rescore: &[f32],
2298    dimensions: usize,
2299    node_count: u32,
2300    params: GraphParams,
2301) -> Result<(), RefinementError> {
2302    if graph.layout.max_degree() != params.r_max() {
2303        return Err(RefinementError::Geometry(format!(
2304            "graph max degree {}, profile max degree {}",
2305            graph.layout.max_degree(),
2306            params.r_max()
2307        )));
2308    }
2309    let source = graph.neighbors.clone();
2310    for owner in 0..node_count {
2311        let owner_index = node_index(owner, node_count)?;
2312        let own = source.get(owner_index).ok_or_else(|| {
2313            RefinementError::Geometry(format!("neighbors for node {owner} are unavailable"))
2314        })?;
2315        let mut candidates = own.clone();
2316        for neighbor in own {
2317            let neighbor_index = node_index(*neighbor, node_count)?;
2318            candidates.extend_from_slice(source.get(neighbor_index).ok_or_else(|| {
2319                RefinementError::Geometry(format!(
2320                    "neighbors for candidate {neighbor} are unavailable"
2321                ))
2322            })?);
2323        }
2324        let mut pruned = robust_prune_rows(
2325            rescore,
2326            dimensions,
2327            node_count,
2328            owner,
2329            &candidates,
2330            params.alpha_refine(),
2331            usize::from(params.r_target()),
2332        )
2333        .map_err(RefinementError::Graph)?;
2334        supplement_long_range(
2335            rescore,
2336            dimensions,
2337            owner,
2338            own,
2339            usize::from(params.r_max()),
2340            &mut pruned,
2341        )?;
2342        let destination = graph.neighbors.get_mut(owner_index).ok_or_else(|| {
2343            RefinementError::Geometry(format!("output neighbors for node {owner} are unavailable"))
2344        })?;
2345        *destination = pruned;
2346    }
2347    Ok(())
2348}
2349
2350fn seed_refit(
2351    graph: &mut OwnedGraph,
2352    rescore: &[f32],
2353    dimensions: usize,
2354    node_count: u32,
2355) -> Result<(), RefinementError> {
2356    let mut medoid = None::<(u32, f64)>;
2357    for candidate in 0..node_count {
2358        let mut sum = 0.0_f64;
2359        for other in 0..node_count {
2360            sum += exact_distance(rescore, dimensions, node_count, candidate, other)?;
2361        }
2362        if medoid
2363            .is_none_or(|current| sum < current.1 || (sum == current.1 && candidate < current.0))
2364        {
2365            medoid = Some((candidate, sum));
2366        }
2367    }
2368    let Some((medoid, _)) = medoid else {
2369        return Err(RefinementError::Geometry(
2370            "seed refit requires at least one row".to_owned(),
2371        ));
2372    };
2373    let mut entries = vec![medoid];
2374    while entries.len() < ENTRY_POINT_COUNT && entries.len() < node_count as usize {
2375        let mut farthest = None::<(u32, f64)>;
2376        for candidate in 0..node_count {
2377            if entries.contains(&candidate) {
2378                continue;
2379            }
2380            let nearest = entries.iter().try_fold(f64::INFINITY, |nearest, entry| {
2381                exact_distance(rescore, dimensions, node_count, candidate, *entry)
2382                    .map(|distance| nearest.min(distance))
2383            })?;
2384            if farthest.is_none_or(|current| {
2385                nearest > current.1 || (nearest == current.1 && candidate < current.0)
2386            }) {
2387                farthest = Some((candidate, nearest));
2388            }
2389        }
2390        let Some((next, _)) = farthest else {
2391            break;
2392        };
2393        entries.push(next);
2394    }
2395    for (node, flags) in graph.flags.iter_mut().enumerate() {
2396        let node = u32::try_from(node)
2397            .map_err(|_| RefinementError::Geometry("entry row exceeds u32".to_owned()))?;
2398        *flags = (*flags & HUB_FLAG) | u8::from(entries.contains(&node));
2399    }
2400    Ok(())
2401}
2402
2403fn neighbor_reorder(
2404    graph: &mut OwnedGraph,
2405    reader: &SegmentReader,
2406    node_count: u32,
2407    seed: u64,
2408) -> Result<(), RefinementError> {
2409    for owner in 0..node_count {
2410        let owner_index = node_index(owner, node_count)?;
2411        let neighbors = graph.neighbors.get_mut(owner_index).ok_or_else(|| {
2412            RefinementError::Geometry(format!("neighbors for node {owner} are unavailable"))
2413        })?;
2414        reorder_neighbors_bit4(reader, owner, neighbors, seed)?;
2415    }
2416    Ok(())
2417}
2418
2419fn connectivity_repair(graph: &mut OwnedGraph, params: GraphParams) -> Result<(), RefinementError> {
2420    if graph.layout.max_degree() != params.r_max() {
2421        return Err(RefinementError::Geometry(format!(
2422            "graph max degree {}, profile max degree {}",
2423            graph.layout.max_degree(),
2424            params.r_max()
2425        )));
2426    }
2427    let node_count = u32::try_from(graph.neighbors.len())
2428        .map_err(|_| RefinementError::Geometry("graph rows exceed u32".to_owned()))?;
2429    let root = first_entry(graph)?;
2430    let mut seen = vec![0_u8; graph.neighbors.len()];
2431    let mut queue = VecDeque::with_capacity(graph.neighbors.len());
2432    mark_and_enqueue_byte(root, &mut seen, &mut queue)?;
2433    loop {
2434        while let Some(owner) = queue.pop_front() {
2435            let neighbors = graph
2436                .neighbors
2437                .get(node_index(owner, node_count)?)
2438                .ok_or_else(|| {
2439                    RefinementError::Geometry(format!(
2440                        "neighbors for repair node {owner} are unavailable"
2441                    ))
2442                })?
2443                .clone();
2444            for neighbor in neighbors {
2445                mark_and_enqueue_byte(neighbor, &mut seen, &mut queue)?;
2446            }
2447        }
2448        let Some(unseen) = seen.iter().position(|marker| *marker == 0) else {
2449            break;
2450        };
2451        let target = u32::try_from(unseen)
2452            .map_err(|_| RefinementError::Geometry("unreachable row exceeds u32".to_owned()))?;
2453        stitch_component(graph, &seen, target, params.r_max())?;
2454        mark_and_enqueue_byte(target, &mut seen, &mut queue)?;
2455    }
2456    for owner in 0..node_count {
2457        densify_hub(graph, owner, params.r_max())?;
2458    }
2459    Ok(())
2460}
2461
2462fn reorder_neighbors_bit4(
2463    reader: &SegmentReader,
2464    owner: u32,
2465    neighbors: &mut Vec<u32>,
2466    seed: u64,
2467) -> Result<(), RefinementError> {
2468    let dimensions = reader.meta().dims as usize;
2469    let node_count = reader.meta().row_count;
2470    let query = exact_row(reader.rescore_f32()?, dimensions, node_count, owner)?;
2471    let query_norm_squared = query
2472        .iter()
2473        .map(|value| f64::from(*value) * f64::from(*value))
2474        .sum::<f64>();
2475    let prepared =
2476        prepare_bit4_query(query, seed ^ u64::from(owner)).map_err(GraphBuildError::Quant)?;
2477    let code_stride = dimensions.div_ceil(2);
2478    let codes = reader.bit4_codes()?;
2479    let factors = reader.bit4_factors()?;
2480    let mut scored = Vec::with_capacity(neighbors.len());
2481    for neighbor in neighbors.iter().copied() {
2482        let index = node_index(neighbor, node_count)?;
2483        let start = index
2484            .checked_mul(code_stride)
2485            .ok_or_else(|| RefinementError::Geometry("Bit4 row offset overflow".to_owned()))?;
2486        let end = start
2487            .checked_add(code_stride)
2488            .ok_or_else(|| RefinementError::Geometry("Bit4 row end overflow".to_owned()))?;
2489        let code = codes.get(start..end).ok_or_else(|| {
2490            RefinementError::Geometry(format!("Bit4 row {neighbor} is unavailable"))
2491        })?;
2492        let factor = factors.get(index).copied().ok_or_else(|| {
2493            RefinementError::Geometry(format!("Bit4 factors {neighbor} are unavailable"))
2494        })?;
2495        let dot = est_dot_bit4(&prepared, code, factor).map_err(GraphBuildError::Quant)?;
2496        let norm = factor.norm();
2497        let distance = (query_norm_squared + norm * norm - 2.0 * f64::from(dot)).max(0.0);
2498        scored.push((neighbor, distance));
2499    }
2500    scored.sort_unstable_by(|left, right| {
2501        left.1
2502            .total_cmp(&right.1)
2503            .then_with(|| left.0.cmp(&right.0))
2504    });
2505    neighbors.clear();
2506    neighbors.extend(scored.into_iter().map(|(neighbor, _)| neighbor));
2507    Ok(())
2508}
2509
2510fn supplement_long_range(
2511    rescore: &[f32],
2512    dimensions: usize,
2513    owner: u32,
2514    neighbors: &[u32],
2515    maximum_degree: usize,
2516    selected: &mut Vec<u32>,
2517) -> Result<(), RefinementError> {
2518    let node_count = u32::try_from(rescore.len() / dimensions.max(1))
2519        .map_err(|_| RefinementError::Geometry("exact row count exceeds u32".to_owned()))?;
2520    let mut long_range = Vec::with_capacity(neighbors.len());
2521    for neighbor in neighbors {
2522        let distance = exact_distance(rescore, dimensions, node_count, owner, *neighbor)?;
2523        long_range.push((*neighbor, distance));
2524    }
2525    long_range.sort_unstable_by(|left, right| {
2526        right
2527            .1
2528            .total_cmp(&left.1)
2529            .then_with(|| left.0.cmp(&right.0))
2530    });
2531    for (neighbor, _) in long_range {
2532        if selected.len() >= maximum_degree {
2533            break;
2534        }
2535        if !selected.contains(&neighbor) {
2536            selected.push(neighbor);
2537        }
2538    }
2539    Ok(())
2540}
2541
2542fn exact_distance(
2543    rescore: &[f32],
2544    dimensions: usize,
2545    node_count: u32,
2546    left: u32,
2547    right: u32,
2548) -> Result<f64, RefinementError> {
2549    let left = exact_row(rescore, dimensions, node_count, left)?;
2550    let right = exact_row(rescore, dimensions, node_count, right)?;
2551    Ok(left
2552        .iter()
2553        .zip(right)
2554        .map(|(left, right)| {
2555            let difference = f64::from(*left) - f64::from(*right);
2556            difference * difference
2557        })
2558        .sum())
2559}
2560
2561fn exact_row(
2562    rescore: &[f32],
2563    dimensions: usize,
2564    node_count: u32,
2565    node: u32,
2566) -> Result<&[f32], RefinementError> {
2567    let index = node_index(node, node_count)?;
2568    let expected = (node_count as usize)
2569        .checked_mul(dimensions)
2570        .ok_or_else(|| RefinementError::Geometry("exact row length overflow".to_owned()))?;
2571    if dimensions == 0 || rescore.len() != expected {
2572        return Err(RefinementError::Geometry(format!(
2573            "exact rows contain {} values, expected {expected}",
2574            rescore.len()
2575        )));
2576    }
2577    let start = index
2578        .checked_mul(dimensions)
2579        .ok_or_else(|| RefinementError::Geometry("exact row offset overflow".to_owned()))?;
2580    let end = start
2581        .checked_add(dimensions)
2582        .ok_or_else(|| RefinementError::Geometry("exact row end overflow".to_owned()))?;
2583    rescore
2584        .get(start..end)
2585        .ok_or_else(|| RefinementError::Geometry(format!("exact row {node} is unavailable")))
2586}
2587
2588fn node_index(node: u32, node_count: u32) -> Result<usize, RefinementError> {
2589    if node >= node_count {
2590        return Err(RefinementError::Geometry(format!(
2591            "node {node} is outside row count {node_count}"
2592        )));
2593    }
2594    usize::try_from(node).map_err(|_| RefinementError::Geometry("node exceeds usize".to_owned()))
2595}
2596
2597fn encode_owned_graph(
2598    reader: &SegmentReader,
2599    graph: &OwnedGraph,
2600    new_to_old: Option<&[u32]>,
2601) -> Result<Vec<u8>, RefinementError> {
2602    let source = reader.graph_node_blocks()?;
2603    let node_count = source.node_count();
2604    let mut blocks = Vec::with_capacity(node_count as usize);
2605    for new_node in 0..node_count {
2606        let source_node = match new_to_old {
2607            Some(order) => *order
2608                .get(node_index(new_node, node_count)?)
2609                .ok_or_else(|| {
2610                    RefinementError::Geometry(format!(
2611                        "source row for new row {new_node} is missing"
2612                    ))
2613                })?,
2614            None => new_node,
2615        };
2616        blocks.push(source.block(source_node)?);
2617    }
2618    let mut nodes = Vec::with_capacity(node_count as usize);
2619    for new_node in 0..node_count {
2620        let index = node_index(new_node, node_count)?;
2621        let block = blocks.get(index).ok_or_else(|| {
2622            RefinementError::Geometry(format!("source block for new row {new_node} is missing"))
2623        })?;
2624        nodes.push(GraphNodeBlockInput {
2625            codes: block.codes(),
2626            factors: block.factors(),
2627            flags: *graph.flags.get(index).ok_or_else(|| {
2628                RefinementError::Geometry(format!("flags for new row {new_node} are missing"))
2629            })?,
2630            neighbors: graph.neighbors.get(index).ok_or_else(|| {
2631                RefinementError::Geometry(format!("neighbors for new row {new_node} are missing"))
2632            })?,
2633        });
2634    }
2635    encode_node_blocks_with_refinement_passes(
2636        GraphNodeBlockBuild {
2637            layout: graph.layout,
2638            nodes: &nodes,
2639        },
2640        graph.passes,
2641    )
2642    .map(|encoded| encoded.into_bytes())
2643    .map_err(RefinementError::NodeBlock)
2644}
2645
2646fn refinement_intermediate_id(output_id: SegmentId) -> SegmentId {
2647    let first = xxh3_64_with_seed(output_id.as_bytes(), INTERMEDIATE_ID_SEED).to_be_bytes();
2648    let second = xxh3_64_with_seed(output_id.as_bytes(), !INTERMEDIATE_ID_SEED).to_be_bytes();
2649    let mut bytes = [0_u8; 16];
2650    if let Some(prefix) = bytes.get_mut(..8) {
2651        prefix.copy_from_slice(&first);
2652    }
2653    if let Some(suffix) = bytes.get_mut(8..) {
2654        suffix.copy_from_slice(&second);
2655    }
2656    SegmentId::from_bytes(bytes)
2657}
2658
2659fn write_permuted_segment(
2660    vfs: &dyn Vfs,
2661    directory: &Path,
2662    input: &SegmentReader,
2663    new_to_old: &[u32],
2664    output_id: SegmentId,
2665    analyzer: &Analyzer,
2666    policy: DurabilityPolicy,
2667) -> Result<SegmentMeta, RefinementError> {
2668    for entry in input.directory() {
2669        if !crate::graph::consolidate::consolidation_supports_region(entry.kind) {
2670            return Err(RefinementError::Geometry(format!(
2671                "renumber cannot carry source region kind {}",
2672                entry.kind
2673            )));
2674        }
2675    }
2676    let rows = new_to_old
2677        .iter()
2678        .map(|row| usize::try_from(*row))
2679        .collect::<Result<Vec<_>, _>>()
2680        .map_err(|_| RefinementError::Geometry("permuted row exceeds usize".to_owned()))?;
2681    let dimensions = input.meta().dims as usize;
2682    let (codes, factors) =
2683        crate::ingest::purge_support::gather_survivor_codes(input, &rows, dimensions)?;
2684    let crate::ingest::purge_support::OwnedFactors::Bit4(factors) = factors else {
2685        return Err(RefinementError::Geometry(
2686            "renumbering requires Bit4 factors".to_owned(),
2687        ));
2688    };
2689    let rescore = crate::ingest::purge_support::gather_survivor_rescore(input, &rows, dimensions)?;
2690    let source_columns = input.columns()?;
2691    let mut columns = ColumnStoreBuilder::new(source_columns.schema().clone());
2692    crate::ingest::purge_support::append_survivor_columns(&mut columns, &source_columns, &rows)?;
2693    let columns = columns
2694        .finish()
2695        .map_err(|error| RefinementError::Segment(SegmentError::Columns(error.to_string())))?;
2696    let source_alive = input.alive()?;
2697    let mut alive = AliveSet::new(input.meta().row_count);
2698    for (new_row, old_row) in new_to_old.iter().copied().enumerate() {
2699        if !source_alive.is_alive(old_row) {
2700            alive
2701                .tombstone(u32::try_from(new_row).map_err(|_| {
2702                    RefinementError::Geometry("new alive row exceeds u32".to_owned())
2703                })?)
2704                .map_err(|_| RefinementError::Geometry("alive permutation failed".to_owned()))?;
2705        }
2706    }
2707    let documents_present = input
2708        .directory()
2709        .iter()
2710        .any(|entry| entry.kind == RegionKind::DocumentVersions.id());
2711    let (doc_ids, revisions) = if documents_present {
2712        crate::ingest::purge_support::gather_survivor_documents(input, &rows)?
2713    } else {
2714        (Vec::new(), Vec::new())
2715    };
2716    let metadata = gather_metadata(input, &rows)?;
2717    let text = gather_text(input, &rows)?;
2718    let postings = rebuild_postings(input, &rows, analyzer)?;
2719    if !documents_present && (metadata.is_some() || text.is_some() || postings.is_some()) {
2720        return Err(RefinementError::Geometry(
2721            "row payloads require document identities".to_owned(),
2722        ));
2723    }
2724    let build = SegmentBuild {
2725        id: output_id,
2726        scheme: input.meta().scheme,
2727        dims: input.meta().dims,
2728        codes: &codes,
2729        factors: SegmentFactors::Bit4(&factors),
2730        rescore: &rescore,
2731        columns: &columns,
2732        alive: &alive,
2733    };
2734    if documents_present {
2735        let documents = SegmentDocumentVersions {
2736            doc_ids: &doc_ids,
2737            revisions: &revisions,
2738        };
2739        let metadata = metadata
2740            .as_ref()
2741            .map(|(end_offsets, bytes)| SegmentStoredMetadata { end_offsets, bytes });
2742        let text = text
2743            .as_ref()
2744            .map(|(present, end_offsets, bytes)| SegmentStoredText {
2745                present,
2746                end_offsets,
2747                bytes,
2748            });
2749        let postings = postings.as_deref().map(|bytes| SegmentPostings { bytes });
2750        write_segment_with_documents_payloads(
2751            vfs,
2752            directory,
2753            build,
2754            SegmentPayloads {
2755                documents,
2756                metadata,
2757                text,
2758                postings,
2759            },
2760            policy,
2761        )
2762        .map_err(RefinementError::Segment)
2763    } else {
2764        write_segment(vfs, directory, build, policy).map_err(RefinementError::Segment)
2765    }
2766}
2767
2768type StoredMetadata = (Vec<u64>, Vec<u8>);
2769
2770fn gather_metadata(
2771    input: &SegmentReader,
2772    rows: &[usize],
2773) -> Result<Option<StoredMetadata>, RefinementError> {
2774    let Some(source) = input.stored_metadata()? else {
2775        return Ok(None);
2776    };
2777    let mut end_offsets = Vec::with_capacity(rows.len());
2778    let mut bytes = Vec::new();
2779    for row in rows {
2780        bytes.extend_from_slice(source.row(*row).ok_or_else(|| {
2781            RefinementError::Geometry(format!("stored metadata row {row} is unavailable"))
2782        })?);
2783        end_offsets.push(u64::try_from(bytes.len()).map_err(|_| {
2784            RefinementError::Geometry("stored metadata bytes exceed u64".to_owned())
2785        })?);
2786    }
2787    Ok(Some((end_offsets, bytes)))
2788}
2789
2790type StoredText = (Vec<u8>, Vec<u64>, Vec<u8>);
2791
2792fn gather_text(
2793    input: &SegmentReader,
2794    rows: &[usize],
2795) -> Result<Option<StoredText>, RefinementError> {
2796    let Some(source) = input.stored_text()? else {
2797        return Ok(None);
2798    };
2799    let mut present = Vec::with_capacity(rows.len());
2800    let mut end_offsets = Vec::with_capacity(rows.len());
2801    let mut bytes = Vec::new();
2802    for row in rows {
2803        match source.row(*row).ok_or_else(|| {
2804            RefinementError::Geometry(format!("stored text row {row} is unavailable"))
2805        })? {
2806            Some(text) => {
2807                present.push(1);
2808                bytes.extend_from_slice(text.as_bytes());
2809            }
2810            None => present.push(0),
2811        }
2812        end_offsets.push(
2813            u64::try_from(bytes.len()).map_err(|_| {
2814                RefinementError::Geometry("stored text bytes exceed u64".to_owned())
2815            })?,
2816        );
2817    }
2818    Ok(Some((present, end_offsets, bytes)))
2819}
2820
2821fn rebuild_postings(
2822    input: &SegmentReader,
2823    rows: &[usize],
2824    analyzer: &Analyzer,
2825) -> Result<Option<Vec<u8>>, RefinementError> {
2826    if input.postings()?.is_none() {
2827        return Ok(None);
2828    }
2829    let source = input.stored_text()?.ok_or_else(|| {
2830        RefinementError::Geometry("renumber cannot rebuild postings without stored text".to_owned())
2831    })?;
2832    let mut index = SegmentIndex::new();
2833    for row in rows {
2834        let text = source.row(*row).ok_or_else(|| {
2835            RefinementError::Geometry(format!("stored text row {row} is unavailable"))
2836        })?;
2837        let document = text.map_or_else(LexicalDocument::new, LexicalDocument::with_text);
2838        index
2839            .push_document(analyzer, &document)
2840            .map_err(|error| RefinementError::Geometry(error.to_string()))?;
2841    }
2842    SealedSegment::seal(&index)
2843        .map_err(|error| RefinementError::Geometry(error.to_string()))?
2844        .encode_region()
2845        .map(Some)
2846        .map_err(|error| RefinementError::Geometry(error.to_string()))
2847}