1mod native;
4
5pub(crate) use native::{NativeGraphBuildError, build_native_graph};
6
7use std::path::{Path, PathBuf};
8
9use xxhash_rust::xxh3::xxh3_64;
10
11use crate::graph::GraphParams;
12use crate::graph::block::{
13 GraphNodeBlockBuild, GraphNodeBlockInput, GraphNodeError, GraphNodeLayout,
14 NODE_BLOCK_TRAILER_LEN, encode_node_blocks_controlled,
15};
16use crate::lifecycle::QueryCancellation;
17use crate::lifecycle::durability::{DurabilityPolicy, SyncRequirement};
18use crate::lifecycle::stats::{AccountedCounter, Accounting, AllocationComponent};
19use crate::lifecycle::{QueryControl, SnapshotLease, Store, StoreError};
20use crate::quant::{
21 Bit4ControlError, Bit4Factors, QuantError, est_dot_bit4_controlled,
22 prepare_bit4_query_controlled,
23};
24use crate::scan::ScanError;
25use crate::segment::layout::RegionKind;
26use crate::segment::reader::SegmentReader;
27use crate::segment::writer::{
28 SegmentBuild, SegmentCopiedRegion, SegmentDocumentVersions, SegmentFactors,
29 write_segment_with_encoded_graph_and_copied_regions,
30};
31use crate::segment::{SegmentError, SegmentId, SegmentMeta};
32use crate::vfs::Vfs;
33
34const ENTRY_POINT_COUNT: usize = 4;
35const MEDOID_SAMPLE_ROWS: usize = 1_000;
36const CHECKPOINT_MAGIC: [u8; 8] = *b"ZEVAMCP1";
37const CHECKPOINT_VERSION: u16 = 1;
38const CHECKPOINT_HEADER_BYTES: usize = 80;
39const CHECKPOINT_CHECKSUM_BYTES: usize = 8;
40const CANCELLATION_CHECK_ROWS: usize = 64;
41
42#[derive(Clone, Copy, Debug, Eq, PartialEq)]
44pub(crate) enum GraphRewriteRegion {
45 Reencoded,
47 Replaced,
49 CopyForward,
51 Unsupported,
53}
54
55pub(crate) const fn graph_rewrite_source_region(kind: u16) -> GraphRewriteRegion {
57 match RegionKind::from_id(kind) {
58 Some(
59 RegionKind::Columns
60 | RegionKind::Alive
61 | RegionKind::VectorCodes
62 | RegionKind::VectorFactors
63 | RegionKind::VectorRescore
64 | RegionKind::DocumentVersions,
65 ) => GraphRewriteRegion::Reencoded,
66 Some(RegionKind::GraphNodeBlocks | RegionKind::ChecksumTable) => {
67 GraphRewriteRegion::Replaced
68 }
69 Some(RegionKind::Postings | RegionKind::StoredMetadata | RegionKind::StoredText) | None => {
70 GraphRewriteRegion::CopyForward
71 }
72 Some(
73 RegionKind::GraphColocatedCodes
74 | RegionKind::SignPlane
75 | RegionKind::PdxClusteredBlocks
76 | RegionKind::VectorSpaceN,
77 ) => GraphRewriteRegion::Unsupported,
78 }
79}
80
81#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
83pub enum GraphBuildPasses {
84 #[default]
87 One,
88 Two,
90}
91
92pub struct GraphBuildArtifact {
94 encoded_region: Vec<u8>,
95 entry_points: Vec<u32>,
96 work_rows_completed: u64,
97 _memory: Option<AccountedCounter>,
98}
99
100impl GraphBuildArtifact {
101 #[must_use]
103 pub fn encoded_region(&self) -> &[u8] {
104 &self.encoded_region
105 }
106
107 #[must_use]
109 pub fn entry_points(&self) -> &[u32] {
110 &self.entry_points
111 }
112
113 #[must_use]
115 pub const fn work_rows_completed(&self) -> u64 {
116 self.work_rows_completed
117 }
118
119 pub(crate) fn resident_bytes(&self) -> Result<usize, GraphBuildError> {
120 self.encoded_region
121 .capacity()
122 .checked_add(
123 self.entry_points
124 .capacity()
125 .checked_mul(std::mem::size_of::<u32>())
126 .ok_or_else(|| {
127 GraphBuildError::Geometry("entry capacity overflow".to_owned())
128 })?,
129 )
130 .ok_or_else(|| GraphBuildError::Geometry("graph capacity overflow".to_owned()))
131 }
132
133 pub fn write_segment_with_graph(
135 &self,
136 vfs: &dyn Vfs,
137 directory: &Path,
138 input: &SegmentReader,
139 output_id: SegmentId,
140 policy: DurabilityPolicy,
141 ) -> Result<SegmentMeta, GraphBuildError> {
142 write_segment_with_encoded_graph(
143 &self.encoded_region,
144 vfs,
145 directory,
146 input,
147 output_id,
148 policy,
149 )
150 }
151}
152
153pub(crate) fn write_segment_with_encoded_graph(
154 encoded_region: &[u8],
155 vfs: &dyn Vfs,
156 directory: &Path,
157 input: &SegmentReader,
158 output_id: SegmentId,
159 policy: DurabilityPolicy,
160) -> Result<SegmentMeta, GraphBuildError> {
161 let columns = input.columns()?;
162 let alive = input.alive()?;
163 let factors = input.bit4_factors()?;
164 let mut doc_ids = Vec::with_capacity(input.meta().row_count as usize);
165 let mut revisions = Vec::with_capacity(input.meta().row_count as usize);
166 let mut has_documents = None;
167 for row in 0..input.meta().row_count as usize {
168 let document = input.document_version(row)?;
169 let present = document.is_some();
170 if has_documents.is_some_and(|expected| expected != present) {
171 return Err(GraphBuildError::Segment(SegmentError::Geometry(
172 "document-version region is present for only part of the segment".to_owned(),
173 )));
174 }
175 has_documents = Some(present);
176 if let Some(document) = document {
177 doc_ids.push(document.doc_id());
178 revisions.push(document.revision());
179 }
180 }
181 let build = SegmentBuild {
182 id: output_id,
183 scheme: input.meta().scheme,
184 dims: input.meta().dims,
185 codes: input.bit4_codes()?,
186 factors: SegmentFactors::Bit4(factors),
187 rescore: input.rescore_f32()?,
188 columns: &columns,
189 alive: &alive,
190 };
191 let documents = (has_documents == Some(true)).then_some(SegmentDocumentVersions {
192 doc_ids: &doc_ids,
193 revisions: &revisions,
194 });
195 let mut copied_regions = Vec::new();
196 for entry in input.directory() {
197 match graph_rewrite_source_region(entry.kind) {
198 GraphRewriteRegion::CopyForward => {
199 copied_regions.push(SegmentCopiedRegion {
200 kind: entry.kind,
201 version: entry.version,
202 bytes: input.region_by_id(entry.kind)?,
203 });
204 }
205 GraphRewriteRegion::Unsupported => {
206 return Err(GraphBuildError::Segment(SegmentError::Geometry(format!(
207 "graph rewrite cannot carry source region kind {}",
208 entry.kind
209 ))));
210 }
211 GraphRewriteRegion::Reencoded | GraphRewriteRegion::Replaced => {}
212 }
213 }
214 write_segment_with_encoded_graph_and_copied_regions(
215 vfs,
216 directory,
217 build,
218 encoded_region,
219 documents,
220 &copied_regions,
221 policy,
222 )
223 .map_err(GraphBuildError::Segment)
224}
225
226#[derive(Debug)]
228pub enum GraphBuildError {
229 Profile(crate::graph::search::GraphProfileError),
231 CheckpointIo {
233 path: PathBuf,
235 source: std::io::Error,
237 },
238 CheckpointCorrupt(String),
240 BudgetExhausted {
242 rows_completed: u64,
244 },
245 UnitNorm {
247 row_id: u32,
249 squared_norm: f64,
251 tolerance: f64,
253 },
254 Cancelled {
256 partial: bool,
258 },
259 Timeout {
261 partial: bool,
263 },
264 ReadCancelled {
266 partial: bool,
268 },
269 Segment(SegmentError),
271 Quant(QuantError),
273 NodeBlock(GraphNodeError),
275 Geometry(String),
277 NodeIdOutOfRange {
279 node_id: u32,
281 node_count: u32,
283 },
284 Store(StoreError),
286}
287
288impl std::fmt::Display for GraphBuildError {
289 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
290 match self {
291 Self::Profile(error) => write!(formatter, "graph build {error}"),
292 Self::CheckpointIo { path, source } => {
293 write!(
294 formatter,
295 "graph build checkpoint I/O {}: {source}",
296 path.display()
297 )
298 }
299 Self::CheckpointCorrupt(detail) => {
300 write!(formatter, "graph build checkpoint is corrupt: {detail}")
301 }
302 Self::BudgetExhausted { rows_completed } => write!(
303 formatter,
304 "graph build work budget exhausted after {rows_completed} checkpointed rows"
305 ),
306 Self::UnitNorm {
307 row_id,
308 squared_norm,
309 tolerance,
310 } => write!(
311 formatter,
312 "angular graph row {row_id} has squared norm {squared_norm}; expected 1 within {tolerance}"
313 ),
314 Self::Cancelled { partial } => {
315 write!(formatter, "graph build was cancelled (partial={partial})")
316 }
317 Self::Timeout { partial } => {
318 write!(
319 formatter,
320 "graph build deadline expired (partial={partial})"
321 )
322 }
323 Self::ReadCancelled { partial } => write!(
324 formatter,
325 "store close cancelled graph build (partial={partial})"
326 ),
327 Self::Segment(error) => write!(formatter, "graph build segment input: {error}"),
328 Self::Quant(error) => write!(formatter, "graph build Bit4 score: {error}"),
329 Self::NodeBlock(error) => write!(formatter, "graph build node blocks: {error}"),
330 Self::Geometry(detail) => write!(formatter, "graph build geometry: {detail}"),
331 Self::NodeIdOutOfRange {
332 node_id,
333 node_count,
334 } => write!(
335 formatter,
336 "graph build node id {node_id} is outside dense row count {node_count}"
337 ),
338 Self::Store(error) => write!(formatter, "graph build memory accounting: {error}"),
339 }
340 }
341}
342
343impl std::error::Error for GraphBuildError {
344 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
345 match self {
346 Self::Profile(error) => Some(error),
347 Self::CheckpointIo { source, .. } => Some(source),
348 Self::Segment(error) => Some(error),
349 Self::Quant(error) => Some(error),
350 Self::NodeBlock(error) => Some(error),
351 Self::Store(error) => Some(error),
352 Self::CheckpointCorrupt(_)
353 | Self::BudgetExhausted { .. }
354 | Self::UnitNorm { .. }
355 | Self::Cancelled { .. }
356 | Self::Timeout { .. }
357 | Self::ReadCancelled { .. }
358 | Self::Geometry(_)
359 | Self::NodeIdOutOfRange { .. } => None,
360 }
361 }
362}
363
364struct GraphBuildSession<'a> {
365 vfs: &'a dyn Vfs,
366 policy: DurabilityPolicy,
367 vectors: SegmentVectors<'a>,
368 params: GraphParams,
369 seed: u64,
370 passes: GraphBuildPasses,
371 checkpoint_path: PathBuf,
372 state: BuildState,
373}
374
375impl<'a> GraphBuildSession<'a> {
376 #[allow(clippy::too_many_arguments)]
377 fn new(
378 vfs: &'a dyn Vfs,
379 reader: &'a SegmentReader,
380 params: GraphParams,
381 seed: u64,
382 passes: GraphBuildPasses,
383 checkpoint_path: &Path,
384 policy: DurabilityPolicy,
385 accounting: Option<&std::sync::Arc<Accounting>>,
386 ) -> Result<Self, GraphBuildError> {
387 let vectors = SegmentVectors::new(reader)?;
388 let state = match vfs.read(checkpoint_path) {
389 Ok(bytes) => {
390 match decode_checkpoint(&bytes, &vectors, params, seed, passes, accounting) {
391 Ok(state) => state,
392 Err(error @ GraphBuildError::CheckpointCorrupt(_)) => {
393 remove_checkpoint(vfs, checkpoint_path)?;
394 return Err(error);
395 }
396 Err(other) => return Err(other),
397 }
398 }
399 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
400 BuildState::new(&vectors, params, seed, accounting)?
401 }
402 Err(source) => {
403 return Err(GraphBuildError::CheckpointIo {
404 path: checkpoint_path.to_path_buf(),
405 source,
406 });
407 }
408 };
409 Ok(Self {
410 vfs,
411 policy,
412 vectors,
413 params,
414 seed,
415 passes,
416 checkpoint_path: checkpoint_path.to_path_buf(),
417 state,
418 })
419 }
420
421 fn advance_batch(
422 &mut self,
423 cancellation: &QueryCancellation<'_>,
424 ) -> Result<bool, GraphBuildError> {
425 check_cancellation(cancellation)?;
426 let complete = self.state.advance_batch(
427 &self.vectors,
428 self.params,
429 self.seed,
430 self.passes,
431 Some(cancellation),
432 )?;
433 check_cancellation(cancellation)?;
434 {
435 let checkpoint_bytes = checkpoint_encoded_bytes(&self.state)?;
436 let base_charge = match self.state.memory.as_mut() {
437 Some(charge) => {
438 let base = charge.bytes();
439 let peak = usize::try_from(base)
440 .ok()
441 .and_then(|base| base.checked_add(checkpoint_bytes))
442 .ok_or_else(|| {
443 GraphBuildError::Geometry("checkpoint accounting overflow".to_owned())
444 })?;
445 charge.set(peak)?;
446 Some(base)
447 }
448 None => None,
449 };
450 let written = write_checkpoint(
451 self.vfs,
452 self.policy,
453 &self.checkpoint_path,
454 &self.vectors,
455 self.params,
456 self.seed,
457 self.passes,
458 &self.state,
459 );
460 if let (Some(charge), Some(base)) = (self.state.memory.as_mut(), base_charge) {
461 let base = usize::try_from(base).map_err(|_| {
462 GraphBuildError::Geometry("base accounting exceeds usize".to_owned())
463 })?;
464 charge.set(base)?;
465 }
466 written?;
467 }
468 Ok(complete)
469 }
470
471 fn completed_work_rows(&self) -> Result<u64, GraphBuildError> {
472 let rows = u64::try_from(self.state.order.len())
473 .map_err(|_| GraphBuildError::Geometry("graph row count exceeds u64".to_owned()))?;
474 let next = u64::try_from(self.state.next_index)
475 .map_err(|_| GraphBuildError::Geometry("graph progress exceeds u64".to_owned()))?;
476 match self.state.phase {
477 BuildPhase::Build => Ok(next),
478 BuildPhase::Refine => rows
479 .checked_add(next)
480 .ok_or_else(|| GraphBuildError::Geometry("graph progress overflow".to_owned())),
481 BuildPhase::Complete => match self.passes {
482 GraphBuildPasses::One => Ok(rows),
483 GraphBuildPasses::Two => rows
484 .checked_mul(2)
485 .ok_or_else(|| GraphBuildError::Geometry("graph progress overflow".to_owned())),
486 },
487 }
488 }
489
490 fn artifact(self) -> Result<GraphBuildArtifact, GraphBuildError> {
491 if !self.state.is_complete(self.passes) {
492 return Err(GraphBuildError::Geometry(
493 "graph artifact requested before every pass completed".to_owned(),
494 ));
495 }
496 let BuildState {
497 entries,
498 adjacency,
499 memory,
500 ..
501 } = self.state;
502 let artifact = encode_artifact(&self.vectors, self.params, &entries, &adjacency, memory)?;
503 remove_checkpoint(self.vfs, &self.checkpoint_path)?;
504 Ok(artifact)
505 }
506}
507
508impl From<SegmentError> for GraphBuildError {
509 fn from(error: SegmentError) -> Self {
510 Self::Segment(error)
511 }
512}
513
514impl From<QuantError> for GraphBuildError {
515 fn from(error: QuantError) -> Self {
516 Self::Quant(error)
517 }
518}
519
520impl From<GraphNodeError> for GraphBuildError {
521 fn from(error: GraphNodeError) -> Self {
522 Self::NodeBlock(error)
523 }
524}
525
526impl From<StoreError> for GraphBuildError {
527 fn from(error: StoreError) -> Self {
528 Self::Store(error)
529 }
530}
531
532#[cfg(test)]
533fn build_graph(
534 reader: &SegmentReader,
535 params: GraphParams,
536 seed: u64,
537 passes: GraphBuildPasses,
538) -> Result<GraphBuildArtifact, GraphBuildError> {
539 let vectors = SegmentVectors::new(reader)?;
540 let mut state = BuildState::new(&vectors, params, seed, None)?;
541 while !state.advance_batch(&vectors, params, seed, passes, None)? {}
542 let BuildState {
543 entries,
544 adjacency,
545 memory,
546 ..
547 } = state;
548 encode_artifact(&vectors, params, &entries, &adjacency, memory)
549}
550
551#[cfg(test)]
552pub(crate) fn build_graph_for_test(
553 reader: &SegmentReader,
554 params: GraphParams,
555 seed: u64,
556 passes: GraphBuildPasses,
557) -> Result<GraphBuildArtifact, GraphBuildError> {
558 build_graph(reader, params, seed, passes)
559}
560
561#[derive(Clone, Copy, Debug)]
563pub struct CheckpointedGraphBuild<'a> {
564 params: GraphParams,
565 seed: u64,
566 passes: GraphBuildPasses,
567 checkpoint_path: &'a Path,
568 control: &'a QueryControl,
569 max_work_rows: Option<u64>,
570}
571
572impl<'a> CheckpointedGraphBuild<'a> {
573 #[must_use]
575 pub fn new(
576 params: GraphParams,
577 seed: u64,
578 passes: GraphBuildPasses,
579 checkpoint_path: &'a Path,
580 control: &'a QueryControl,
581 ) -> Self {
582 Self {
583 params,
584 seed,
585 passes,
586 checkpoint_path,
587 control,
588 max_work_rows: None,
589 }
590 }
591
592 #[must_use]
596 pub const fn with_max_work_rows(mut self, max_work_rows: u64) -> Self {
597 self.max_work_rows = Some(max_work_rows);
598 self
599 }
600}
601
602pub fn build_graph_checkpointed(
609 store: &Store,
610 reader: &SegmentReader,
611 request: CheckpointedGraphBuild<'_>,
612 lease: &SnapshotLease,
613) -> Result<GraphBuildArtifact, GraphBuildError> {
614 let cancellation = QueryCancellation::new(request.control, lease);
615 check_cancellation(&cancellation)?;
616 if lease.graph_profile() == Ok(crate::graph::search::EpochGraphProfile::AngularClass) {
617 validate_angular_rows(reader)?;
618 }
619 let mut session = GraphBuildSession::new(
620 store.vfs.as_ref(),
621 reader,
622 request.params,
623 request.seed,
624 request.passes,
625 request.checkpoint_path,
626 store.durability_policy,
627 Some(&store.accounting),
628 )?;
629 let started_rows = session.completed_work_rows()?;
630 loop {
631 let complete = session.advance_batch(&cancellation)?;
632 let completed_rows = session
633 .completed_work_rows()?
634 .checked_sub(started_rows)
635 .ok_or_else(|| GraphBuildError::Geometry("graph work progress regressed".to_owned()))?;
636 if complete {
637 let mut artifact = session.artifact()?;
638 artifact.work_rows_completed = completed_rows;
639 return Ok(artifact);
640 }
641 if request
642 .max_work_rows
643 .is_some_and(|maximum| completed_rows >= maximum)
644 {
645 return Err(GraphBuildError::BudgetExhausted {
646 rows_completed: completed_rows,
647 });
648 }
649 }
650}
651
652fn validate_angular_rows(reader: &SegmentReader) -> Result<(), GraphBuildError> {
653 let dimensions = reader.meta().dims as usize;
654 for (row, vector) in reader.rescore_f32()?.chunks_exact(dimensions).enumerate() {
655 if let Some(squared_norm) = crate::graph::search::non_unit_squared_norm(vector) {
656 return Err(GraphBuildError::UnitNorm {
657 row_id: u32::try_from(row).map_err(|_| {
658 GraphBuildError::Geometry("angular row id exceeds u32".to_owned())
659 })?,
660 squared_norm,
661 tolerance: crate::graph::search::UNIT_NORM_SQUARED_TOLERANCE,
662 });
663 }
664 }
665 Ok(())
666}
667
668struct SegmentVectors<'a> {
669 segment_id: [u8; 16],
670 node_count: u32,
671 dimensions: usize,
672 code_stride: usize,
673 codes: &'a [u8],
674 factors: &'a [Bit4Factors],
675 rescore: &'a [f32],
676}
677
678impl<'a> SegmentVectors<'a> {
679 fn new(reader: &'a SegmentReader) -> Result<Self, GraphBuildError> {
680 if reader.meta().scheme != 4 {
681 return Err(GraphBuildError::Geometry(format!(
682 "flat Vamana requires Bit4 scheme 4, got {}",
683 reader.meta().scheme
684 )));
685 }
686 let dimensions = reader.meta().dims as usize;
687 let node_count = reader.meta().row_count;
688 if node_count == 0 {
689 return Err(GraphBuildError::Geometry(
690 "flat Vamana requires at least one row".to_owned(),
691 ));
692 }
693 Ok(Self {
694 segment_id: *reader.meta().id.as_bytes(),
695 node_count,
696 dimensions,
697 code_stride: dimensions.div_ceil(2),
698 codes: reader.bit4_codes()?,
699 factors: reader.bit4_factors()?,
700 rescore: reader.rescore_f32()?,
701 })
702 }
703
704 fn from_native(
705 dimensions: usize,
706 codes: &'a [u8],
707 factors: &'a [Bit4Factors],
708 rescore: &'a [f32],
709 ) -> Result<Self, GraphBuildError> {
710 if dimensions == 0 || factors.is_empty() {
711 return Err(GraphBuildError::Geometry(
712 "native flat Vamana requires dimensions and rows".to_owned(),
713 ));
714 }
715 let code_stride = dimensions.div_ceil(2);
716 let expected_codes = factors
717 .len()
718 .checked_mul(code_stride)
719 .ok_or_else(|| GraphBuildError::Geometry("native code length overflow".to_owned()))?;
720 let expected_rescore = factors.len().checked_mul(dimensions).ok_or_else(|| {
721 GraphBuildError::Geometry("native rescore length overflow".to_owned())
722 })?;
723 if codes.len() != expected_codes || rescore.len() != expected_rescore {
724 return Err(GraphBuildError::Geometry(
725 "native flat Vamana vector geometry".to_owned(),
726 ));
727 }
728 Ok(Self {
729 segment_id: [0; 16],
730 node_count: u32::try_from(factors.len()).map_err(|_| {
731 GraphBuildError::Geometry("native row count exceeds u32".to_owned())
732 })?,
733 dimensions,
734 code_stride,
735 codes,
736 factors,
737 rescore,
738 })
739 }
740
741 fn validate_node(&self, node_id: u32) -> Result<usize, GraphBuildError> {
742 if node_id >= self.node_count {
743 return Err(GraphBuildError::NodeIdOutOfRange {
744 node_id,
745 node_count: self.node_count,
746 });
747 }
748 usize::try_from(node_id)
749 .map_err(|_| GraphBuildError::Geometry("node id exceeds usize".to_owned()))
750 }
751
752 fn f32_row(&self, node_id: u32) -> Result<&'a [f32], GraphBuildError> {
753 let node = self.validate_node(node_id)?;
754 let start = node
755 .checked_mul(self.dimensions)
756 .ok_or_else(|| GraphBuildError::Geometry("f32 row offset overflow".to_owned()))?;
757 let end = start
758 .checked_add(self.dimensions)
759 .ok_or_else(|| GraphBuildError::Geometry("f32 row end overflow".to_owned()))?;
760 self.rescore.get(start..end).ok_or_else(|| {
761 GraphBuildError::Geometry(format!("f32 row {node_id} is outside validated rescore"))
762 })
763 }
764
765 fn code_row(&self, node_id: u32) -> Result<&'a [u8], GraphBuildError> {
766 let node = self.validate_node(node_id)?;
767 let start = node
768 .checked_mul(self.code_stride)
769 .ok_or_else(|| GraphBuildError::Geometry("Bit4 row offset overflow".to_owned()))?;
770 let end = start
771 .checked_add(self.code_stride)
772 .ok_or_else(|| GraphBuildError::Geometry("Bit4 row end overflow".to_owned()))?;
773 self.codes.get(start..end).ok_or_else(|| {
774 GraphBuildError::Geometry(format!("Bit4 row {node_id} is outside validated codes"))
775 })
776 }
777
778 fn factor(&self, node_id: u32) -> Result<Bit4Factors, GraphBuildError> {
779 let node = self.validate_node(node_id)?;
780 self.factors.get(node).copied().ok_or_else(|| {
781 GraphBuildError::Geometry(format!("Bit4 factor {node_id} is unavailable"))
782 })
783 }
784}
785
786fn build_native_graph_inner(
787 vectors: &SegmentVectors<'_>,
788 params: GraphParams,
789 seed: u64,
790 control: &mut BuildControl<'_>,
791) -> Result<GraphBuildArtifact, GraphBuildError> {
792 let mut state = BuildState::new_controlled(vectors, params, seed, None, control)?;
793 while !state.advance_batch_controlled(
794 vectors,
795 params,
796 seed,
797 GraphBuildPasses::One,
798 None,
799 control,
800 )? {}
801 let transient_bytes = state
802 .order
803 .capacity()
804 .checked_mul(std::mem::size_of::<u32>())
805 .and_then(|bytes| bytes.checked_add(state.inserted.capacity()))
806 .and_then(|bytes| {
807 state
808 .visited
809 .capacity()
810 .checked_mul(std::mem::size_of::<u32>())
811 .and_then(|visited| bytes.checked_add(visited))
812 })
813 .ok_or_else(|| GraphBuildError::Geometry("build state capacity overflow".to_owned()))?;
814 let BuildState {
815 entries,
816 adjacency,
817 order,
818 inserted,
819 visited,
820 memory,
821 ..
822 } = state;
823 drop(order);
824 drop(inserted);
825 drop(visited);
826 control.release(transient_bytes)?;
827 let source_bytes = entries
828 .capacity()
829 .checked_mul(std::mem::size_of::<u32>())
830 .and_then(|bytes| {
831 adjacency
832 .slots
833 .capacity()
834 .checked_mul(std::mem::size_of::<u32>())
835 .and_then(|slots| bytes.checked_add(slots))
836 })
837 .and_then(|bytes| bytes.checked_add(adjacency.degrees.capacity()))
838 .ok_or_else(|| GraphBuildError::Geometry("graph source capacity overflow".to_owned()))?;
839 let artifact =
840 encode_artifact_controlled(vectors, params, &entries, &adjacency, memory, control)?;
841 drop(entries);
842 drop(adjacency);
843 control.release(source_bytes)?;
844 Ok(artifact)
845}
846
847struct BuildControl<'a> {
848 poll: Option<&'a mut dyn FnMut(u64) -> bool>,
849 memory: Option<&'a mut dyn FnMut(BuildMemoryEvent) -> bool>,
850}
851
852fn reserve_vec<T>(
853 values: &mut Vec<T>,
854 capacity: usize,
855 label: &'static str,
856 control: &mut BuildControl<'_>,
857) -> Result<usize, GraphBuildError> {
858 let requested = capacity
859 .checked_mul(std::mem::size_of::<T>())
860 .ok_or_else(|| GraphBuildError::Geometry(format!("{label} allocation overflow")))?;
861 control.acquire(requested)?;
862 #[cfg(feature = "allocation-audit")]
863 let allocation = crate::allocation_audit::attributed(|| values.try_reserve_exact(capacity));
864 #[cfg(not(feature = "allocation-audit"))]
865 let allocation = values.try_reserve_exact(capacity);
866 allocation.map_err(|_| {
867 GraphBuildError::Store(StoreError::AllocationFailed {
868 needed: requested as u64,
869 component: label,
870 })
871 })?;
872 let actual = values
873 .capacity()
874 .checked_mul(std::mem::size_of::<T>())
875 .ok_or_else(|| GraphBuildError::Geometry(format!("{label} capacity overflow")))?;
876 control.reconcile(requested, actual)?;
877 Ok(actual)
878}
879
880fn resize_vec_controlled<T: Clone>(
881 values: &mut Vec<T>,
882 length: usize,
883 value: T,
884 label: &str,
885 control: &mut BuildControl<'_>,
886) -> Result<(), GraphBuildError> {
887 if length > values.capacity() {
888 return Err(GraphBuildError::Geometry(format!(
889 "{label} fixed capacity exhausted"
890 )));
891 }
892 while values.len() < length {
893 let count = (length - values.len()).min(256);
894 control.step(count as u64)?;
895 values.resize(values.len() + count, value.clone());
896 }
897 Ok(())
898}
899
900fn fill_controlled<T: Clone>(
901 values: &mut [T],
902 value: T,
903 control: &mut BuildControl<'_>,
904) -> Result<(), GraphBuildError> {
905 for chunk in values.chunks_mut(256) {
906 control.step(chunk.len() as u64)?;
907 chunk.fill(value.clone());
908 }
909 Ok(())
910}
911
912#[derive(Clone, Copy)]
913pub(crate) enum BuildMemoryEvent {
914 Acquire(usize),
915 Reconcile { requested: usize, actual: usize },
916 Release(usize),
917}
918
919impl BuildControl<'_> {
920 fn disabled() -> Self {
921 Self {
922 poll: None,
923 memory: None,
924 }
925 }
926
927 fn step(&mut self, units: u64) -> Result<(), GraphBuildError> {
928 if self.poll.as_mut().is_some_and(|poll| !poll(units)) {
929 return Err(GraphBuildError::Geometry(
930 "native graph control interrupted construction".to_owned(),
931 ));
932 }
933 Ok(())
934 }
935
936 fn acquire(&mut self, bytes: usize) -> Result<(), GraphBuildError> {
937 if self
938 .memory
939 .as_mut()
940 .is_some_and(|memory| !memory(BuildMemoryEvent::Acquire(bytes)))
941 {
942 return Err(GraphBuildError::Geometry(
943 "native graph memory admission failed".to_owned(),
944 ));
945 }
946 Ok(())
947 }
948
949 fn reconcile(&mut self, requested: usize, actual: usize) -> Result<(), GraphBuildError> {
950 if self
951 .memory
952 .as_mut()
953 .is_some_and(|memory| !memory(BuildMemoryEvent::Reconcile { requested, actual }))
954 {
955 return Err(GraphBuildError::Geometry(
956 "native graph memory reconciliation failed".to_owned(),
957 ));
958 }
959 Ok(())
960 }
961
962 fn release(&mut self, bytes: usize) -> Result<(), GraphBuildError> {
963 if self
964 .memory
965 .as_mut()
966 .is_some_and(|memory| !memory(BuildMemoryEvent::Release(bytes)))
967 {
968 return Err(GraphBuildError::Geometry(
969 "native graph memory release failed".to_owned(),
970 ));
971 }
972 Ok(())
973 }
974}
975
976struct Adjacency {
977 slots: Vec<u32>,
978 degrees: Vec<u8>,
979 node_count: u32,
980 r_max: u8,
981}
982
983impl Adjacency {
984 #[cfg(test)]
985 fn new(node_count: u32, r_max: u8) -> Result<Self, GraphBuildError> {
986 Self::new_controlled(node_count, r_max, &mut BuildControl::disabled())
987 }
988
989 fn new_controlled(
990 node_count: u32,
991 r_max: u8,
992 control: &mut BuildControl<'_>,
993 ) -> Result<Self, GraphBuildError> {
994 let slot_count = (node_count as usize)
995 .checked_mul(usize::from(r_max))
996 .ok_or_else(|| GraphBuildError::Geometry("adjacency size overflow".to_owned()))?;
997 let mut slots = Vec::new();
998 reserve_vec(&mut slots, slot_count, "adjacency", control)?;
999 resize_vec_controlled(&mut slots, slot_count, u32::MAX, "adjacency", control)?;
1000 let mut degrees = Vec::new();
1001 reserve_vec(&mut degrees, node_count as usize, "degree", control)?;
1002 resize_vec_controlled(&mut degrees, node_count as usize, 0, "degree", control)?;
1003 Ok(Self {
1004 slots,
1005 degrees,
1006 node_count,
1007 r_max,
1008 })
1009 }
1010
1011 fn row_range(&self, node_id: u32) -> Result<std::ops::Range<usize>, GraphBuildError> {
1012 if node_id >= self.node_count {
1013 return Err(GraphBuildError::NodeIdOutOfRange {
1014 node_id,
1015 node_count: self.node_count,
1016 });
1017 }
1018 let start = (node_id as usize)
1019 .checked_mul(usize::from(self.r_max))
1020 .ok_or_else(|| GraphBuildError::Geometry("adjacency offset overflow".to_owned()))?;
1021 let end = start
1022 .checked_add(usize::from(self.r_max))
1023 .ok_or_else(|| GraphBuildError::Geometry("adjacency end overflow".to_owned()))?;
1024 Ok(start..end)
1025 }
1026
1027 fn neighbors(&self, node_id: u32) -> Result<&[u32], GraphBuildError> {
1028 let node = usize::try_from(node_id)
1029 .map_err(|_| GraphBuildError::Geometry("node id exceeds usize".to_owned()))?;
1030 let degree = usize::from(*self.degrees.get(node).ok_or(
1031 GraphBuildError::NodeIdOutOfRange {
1032 node_id,
1033 node_count: self.node_count,
1034 },
1035 )?);
1036 let range = self.row_range(node_id)?;
1037 let start = range.start;
1038 let active_end = start
1039 .checked_add(degree)
1040 .ok_or_else(|| GraphBuildError::Geometry("degree end overflow".to_owned()))?;
1041 self.slots.get(start..active_end).ok_or_else(|| {
1042 GraphBuildError::Geometry(format!("adjacency row {node_id} is unavailable"))
1043 })
1044 }
1045
1046 fn set_neighbors(&mut self, node_id: u32, neighbors: &[u32]) -> Result<(), GraphBuildError> {
1047 if neighbors.len() > usize::from(self.r_max) {
1048 return Err(GraphBuildError::Geometry(format!(
1049 "node {node_id} degree {} exceeds hard cap {}",
1050 neighbors.len(),
1051 self.r_max
1052 )));
1053 }
1054 for neighbor in neighbors {
1055 if *neighbor >= self.node_count {
1056 return Err(GraphBuildError::NodeIdOutOfRange {
1057 node_id: *neighbor,
1058 node_count: self.node_count,
1059 });
1060 }
1061 }
1062 let range = self.row_range(node_id)?;
1063 let row = self.slots.get_mut(range).ok_or_else(|| {
1064 GraphBuildError::Geometry(format!("adjacency row {node_id} is unavailable"))
1065 })?;
1066 row.fill(u32::MAX);
1067 let destination = row.get_mut(..neighbors.len()).ok_or_else(|| {
1068 GraphBuildError::Geometry(format!("adjacency row {node_id} active slice failed"))
1069 })?;
1070 destination.copy_from_slice(neighbors);
1071 let node = usize::try_from(node_id)
1072 .map_err(|_| GraphBuildError::Geometry("node id exceeds usize".to_owned()))?;
1073 let degree = self
1074 .degrees
1075 .get_mut(node)
1076 .ok_or(GraphBuildError::NodeIdOutOfRange {
1077 node_id,
1078 node_count: self.node_count,
1079 })?;
1080 *degree = u8::try_from(neighbors.len())
1081 .map_err(|_| GraphBuildError::Geometry("degree exceeds u8".to_owned()))?;
1082 Ok(())
1083 }
1084}
1085
1086#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1087enum BuildPhase {
1088 Build,
1089 Refine,
1090 Complete,
1091}
1092
1093fn build_arena_bytes(node_count: u32, params: GraphParams) -> Result<usize, GraphBuildError> {
1094 let rows = usize::try_from(node_count)
1095 .map_err(|_| GraphBuildError::Geometry("node count exceeds usize".to_owned()))?;
1096 let slots = rows
1097 .checked_mul(usize::from(params.r_max()))
1098 .and_then(|count| count.checked_mul(std::mem::size_of::<u32>()))
1099 .ok_or_else(|| GraphBuildError::Geometry("accounted adjacency overflow".to_owned()))?;
1100 let row_state = rows
1101 .checked_mul(
1102 std::mem::size_of::<u32>()
1103 .checked_mul(2)
1104 .and_then(|bytes| bytes.checked_add(2))
1105 .ok_or_else(|| GraphBuildError::Geometry("row state width overflow".to_owned()))?,
1106 )
1107 .ok_or_else(|| GraphBuildError::Geometry("accounted row state overflow".to_owned()))?;
1108 let construction = usize::from(params.l_build());
1109 let scratch_per_candidate = std::mem::size_of::<ScoredNode>()
1110 .checked_mul(4)
1111 .and_then(|bytes| bytes.checked_add(std::mem::size_of::<u32>() * 6))
1112 .ok_or_else(|| GraphBuildError::Geometry("scratch width overflow".to_owned()))?;
1113 let candidate_scratch = construction
1114 .checked_add(usize::from(params.r_max()))
1115 .and_then(|candidates| candidates.checked_mul(scratch_per_candidate))
1116 .ok_or_else(|| GraphBuildError::Geometry("scratch arena overflow".to_owned()))?;
1117 let medoid_scratch = MEDOID_SAMPLE_ROWS
1118 .checked_mul(std::mem::size_of::<ScoredNode>())
1119 .ok_or_else(|| GraphBuildError::Geometry("medoid scratch overflow".to_owned()))?;
1120 let scratch = candidate_scratch.max(medoid_scratch);
1121 slots
1122 .checked_add(row_state)
1123 .and_then(|bytes| bytes.checked_add(ENTRY_POINT_COUNT * std::mem::size_of::<u32>()))
1124 .and_then(|bytes| bytes.checked_add(scratch))
1125 .ok_or_else(|| GraphBuildError::Geometry("graph build arena overflow".to_owned()))
1126}
1127
1128struct BuildState {
1129 phase: BuildPhase,
1130 next_index: usize,
1131 order: Vec<u32>,
1132 entries: Vec<u32>,
1133 adjacency: Adjacency,
1134 inserted: Vec<u8>,
1135 visited: Vec<u32>,
1136 visit_epoch: u32,
1137 memory: Option<AccountedCounter>,
1138}
1139
1140impl BuildState {
1141 fn new(
1142 vectors: &SegmentVectors<'_>,
1143 params: GraphParams,
1144 seed: u64,
1145 accounting: Option<&std::sync::Arc<Accounting>>,
1146 ) -> Result<Self, GraphBuildError> {
1147 Self::new_controlled(
1148 vectors,
1149 params,
1150 seed,
1151 accounting,
1152 &mut BuildControl::disabled(),
1153 )
1154 }
1155
1156 fn new_controlled(
1157 vectors: &SegmentVectors<'_>,
1158 params: GraphParams,
1159 seed: u64,
1160 accounting: Option<&std::sync::Arc<Accounting>>,
1161 control: &mut BuildControl<'_>,
1162 ) -> Result<Self, GraphBuildError> {
1163 let mut memory = match accounting {
1164 Some(accounting) => Some(AccountedCounter::new(
1165 accounting,
1166 AllocationComponent::Temporary,
1167 )?),
1168 None => None,
1169 };
1170 if let Some(charge) = memory.as_mut() {
1171 charge.set(build_arena_bytes(vectors.node_count, params)?)?;
1172 }
1173 let mut random = SplitMix64::new(seed);
1174 let mut order = Vec::new();
1175 reserve_vec(
1176 &mut order,
1177 vectors.node_count as usize,
1178 "build order",
1179 control,
1180 )?;
1181 for node in 0..vectors.node_count {
1182 control.step(1)?;
1183 order.push(node);
1184 }
1185 shuffle_controlled(&mut order, &mut random, control)?;
1186 let entries = refined_entry_points(vectors, &order, control)?;
1187 move_entries_to_front(&mut order, &entries, control)?;
1188 let mut inserted = Vec::new();
1189 reserve_vec(
1190 &mut inserted,
1191 vectors.node_count as usize,
1192 "inserted",
1193 control,
1194 )?;
1195 resize_vec_controlled(
1196 &mut inserted,
1197 vectors.node_count as usize,
1198 0,
1199 "inserted",
1200 control,
1201 )?;
1202 let mut visited = Vec::new();
1203 reserve_vec(
1204 &mut visited,
1205 vectors.node_count as usize,
1206 "visited",
1207 control,
1208 )?;
1209 resize_vec_controlled(
1210 &mut visited,
1211 vectors.node_count as usize,
1212 0,
1213 "visited",
1214 control,
1215 )?;
1216 Ok(Self {
1217 phase: BuildPhase::Build,
1218 next_index: 0,
1219 order,
1220 entries,
1221 adjacency: Adjacency::new_controlled(vectors.node_count, params.r_max(), control)?,
1222 inserted,
1223 visited,
1224 visit_epoch: 0,
1225 memory,
1226 })
1227 }
1228
1229 fn is_complete(&self, _passes: GraphBuildPasses) -> bool {
1230 self.phase == BuildPhase::Complete
1231 }
1232
1233 fn advance_batch(
1234 &mut self,
1235 vectors: &SegmentVectors<'_>,
1236 params: GraphParams,
1237 seed: u64,
1238 passes: GraphBuildPasses,
1239 cancellation: Option<&QueryCancellation<'_>>,
1240 ) -> Result<bool, GraphBuildError> {
1241 self.advance_batch_controlled(
1242 vectors,
1243 params,
1244 seed,
1245 passes,
1246 cancellation,
1247 &mut BuildControl::disabled(),
1248 )
1249 }
1250
1251 fn advance_batch_controlled(
1252 &mut self,
1253 vectors: &SegmentVectors<'_>,
1254 params: GraphParams,
1255 seed: u64,
1256 passes: GraphBuildPasses,
1257 cancellation: Option<&QueryCancellation<'_>>,
1258 control: &mut BuildControl<'_>,
1259 ) -> Result<bool, GraphBuildError> {
1260 if self.phase == BuildPhase::Complete {
1261 return Ok(true);
1262 }
1263 let batch_rows = usize::try_from(params.checkpoint_batch_rows()).map_err(|_| {
1264 GraphBuildError::Geometry("checkpoint batch rows exceed usize".to_owned())
1265 })?;
1266 let end = self
1267 .next_index
1268 .checked_add(batch_rows)
1269 .map(|end| end.min(self.order.len()))
1270 .ok_or_else(|| GraphBuildError::Geometry("batch end overflow".to_owned()))?;
1271 let alpha = match self.phase {
1272 BuildPhase::Build => params.alpha_build(),
1273 BuildPhase::Refine => params.alpha_refine(),
1274 BuildPhase::Complete => return Ok(true),
1275 };
1276 let pass_seed = match self.phase {
1277 BuildPhase::Build => seed,
1278 BuildPhase::Refine => seed ^ 0xa1fa_1200_0000_0002,
1279 BuildPhase::Complete => seed,
1280 };
1281 for position in self.next_index..end {
1282 control.step(1)?;
1283 if position.is_multiple_of(CANCELLATION_CHECK_ROWS)
1284 && let Some(cancellation) = cancellation
1285 {
1286 check_cancellation(cancellation)?;
1287 }
1288 let order_position = match self.phase {
1289 BuildPhase::Build => position,
1290 BuildPhase::Refine => {
1291 self.order.len().checked_sub(position + 1).ok_or_else(|| {
1292 GraphBuildError::Geometry("refinement order underflow".to_owned())
1293 })?
1294 }
1295 BuildPhase::Complete => position,
1296 };
1297 let node_id = *self.order.get(order_position).ok_or_else(|| {
1298 GraphBuildError::Geometry(format!(
1299 "build order position {order_position} is unavailable"
1300 ))
1301 })?;
1302 process_node(
1303 vectors,
1304 params,
1305 pass_seed,
1306 alpha,
1307 node_id,
1308 &self.entries,
1309 &mut self.adjacency,
1310 &mut self.inserted,
1311 &mut self.visited,
1312 &mut self.visit_epoch,
1313 control,
1314 )?;
1315 }
1316 self.next_index = end;
1317 if self.next_index == self.order.len() {
1318 match (self.phase, passes) {
1319 (BuildPhase::Build, GraphBuildPasses::Two) => {
1320 self.phase = BuildPhase::Refine;
1321 self.next_index = 0;
1322 fill_controlled(&mut self.inserted, 1, control)?;
1323 }
1324 (BuildPhase::Build, GraphBuildPasses::One) | (BuildPhase::Refine, _) => {
1325 self.phase = BuildPhase::Complete;
1326 }
1327 (BuildPhase::Complete, _) => {}
1328 }
1329 }
1330 Ok(self.phase == BuildPhase::Complete)
1331 }
1332}
1333
1334fn check_cancellation(cancellation: &QueryCancellation<'_>) -> Result<(), GraphBuildError> {
1335 match cancellation.check_graph() {
1336 Ok(()) => Ok(()),
1337 Err(ScanError::Cancelled { .. }) => Err(GraphBuildError::Cancelled { partial: false }),
1338 Err(ScanError::Timeout { .. }) => Err(GraphBuildError::Timeout { partial: false }),
1339 Err(ScanError::ReadCancelled { .. }) => {
1340 Err(GraphBuildError::ReadCancelled { partial: false })
1341 }
1342 Err(error) => Err(GraphBuildError::Geometry(format!(
1343 "query cancellation seam returned non-cancellation error: {error}"
1344 ))),
1345 }
1346}
1347
1348fn checkpoint_temp_path(path: &Path) -> PathBuf {
1349 let mut temporary = path.as_os_str().to_os_string();
1350 temporary.push(".tmp");
1351 PathBuf::from(temporary)
1352}
1353
1354#[allow(clippy::too_many_arguments)]
1355fn write_checkpoint(
1356 vfs: &dyn Vfs,
1357 policy: DurabilityPolicy,
1358 path: &Path,
1359 vectors: &SegmentVectors<'_>,
1360 params: GraphParams,
1361 seed: u64,
1362 passes: GraphBuildPasses,
1363 state: &BuildState,
1364) -> Result<(), GraphBuildError> {
1365 let directory = path.parent().ok_or_else(|| GraphBuildError::CheckpointIo {
1366 path: path.to_path_buf(),
1367 source: std::io::Error::new(
1368 std::io::ErrorKind::InvalidInput,
1369 "checkpoint path has no parent directory",
1370 ),
1371 })?;
1372 let slot_bytes = state
1373 .adjacency
1374 .slots
1375 .len()
1376 .checked_mul(std::mem::size_of::<u32>())
1377 .ok_or_else(|| GraphBuildError::Geometry("checkpoint slot bytes overflow".to_owned()))?;
1378 let payload_bytes = state
1379 .adjacency
1380 .degrees
1381 .len()
1382 .checked_add(slot_bytes)
1383 .ok_or_else(|| GraphBuildError::Geometry("checkpoint payload overflow".to_owned()))?;
1384 let total_bytes = CHECKPOINT_HEADER_BYTES
1385 .checked_add(payload_bytes)
1386 .and_then(|bytes| bytes.checked_add(CHECKPOINT_CHECKSUM_BYTES))
1387 .ok_or_else(|| GraphBuildError::Geometry("checkpoint length overflow".to_owned()))?;
1388 let mut bytes = Vec::with_capacity(total_bytes);
1389 bytes.extend_from_slice(&CHECKPOINT_MAGIC);
1390 bytes.extend_from_slice(&CHECKPOINT_VERSION.to_le_bytes());
1391 bytes.push(phase_code(state.phase));
1392 bytes.push(pass_code(passes));
1393 bytes.push(params.r_target());
1394 bytes.push(params.r_max());
1395 bytes.extend_from_slice(&0_u16.to_le_bytes());
1396 bytes.extend_from_slice(¶ms.l_build().to_le_bytes());
1397 bytes.extend_from_slice(&0_u16.to_le_bytes());
1398 bytes.extend_from_slice(¶ms.checkpoint_batch_rows().to_le_bytes());
1399 bytes.extend_from_slice(&vectors.node_count.to_le_bytes());
1400 let dimensions = u32::try_from(vectors.dimensions)
1401 .map_err(|_| GraphBuildError::Geometry("checkpoint dimensions exceed u32".to_owned()))?;
1402 bytes.extend_from_slice(&dimensions.to_le_bytes());
1403 let next_index = u32::try_from(state.next_index)
1404 .map_err(|_| GraphBuildError::Geometry("checkpoint position exceeds u32".to_owned()))?;
1405 bytes.extend_from_slice(&next_index.to_le_bytes());
1406 bytes.extend_from_slice(&0_u32.to_le_bytes());
1407 bytes.extend_from_slice(&seed.to_le_bytes());
1408 bytes.extend_from_slice(¶ms.alpha_build().to_bits().to_le_bytes());
1409 bytes.extend_from_slice(¶ms.alpha_refine().to_bits().to_le_bytes());
1410 bytes.extend_from_slice(&vectors.segment_id);
1411 bytes.extend_from_slice(
1412 &u64::try_from(payload_bytes)
1413 .map_err(|_| GraphBuildError::Geometry("checkpoint payload exceeds u64".to_owned()))?
1414 .to_le_bytes(),
1415 );
1416 if bytes.len() != CHECKPOINT_HEADER_BYTES {
1417 return Err(GraphBuildError::Geometry(format!(
1418 "checkpoint header encoded {} bytes, expected {CHECKPOINT_HEADER_BYTES}",
1419 bytes.len()
1420 )));
1421 }
1422 bytes.extend_from_slice(&state.adjacency.degrees);
1423 for neighbor in &state.adjacency.slots {
1424 bytes.extend_from_slice(&neighbor.to_le_bytes());
1425 }
1426 bytes.extend_from_slice(&xxh3_64(&bytes).to_le_bytes());
1427 let temporary = checkpoint_temp_path(path);
1428 if let Err(source) = vfs.write(&temporary, &bytes) {
1429 let _ = vfs.delete(&temporary);
1430 return Err(GraphBuildError::CheckpointIo {
1431 path: temporary,
1432 source,
1433 });
1434 }
1435 if let SyncRequirement::Sync(kind) = policy.data_file_sync()
1436 && let Err(source) = vfs.sync(&temporary, kind)
1437 {
1438 let _ = vfs.delete(&temporary);
1439 return Err(GraphBuildError::CheckpointIo {
1440 path: temporary,
1441 source,
1442 });
1443 }
1444 if let Err(source) = vfs.rename(&temporary, path) {
1445 let _ = vfs.delete(&temporary);
1446 return Err(GraphBuildError::CheckpointIo {
1447 path: path.to_path_buf(),
1448 source,
1449 });
1450 }
1451 match policy.directory_sync() {
1452 SyncRequirement::Skip => Ok(()),
1453 SyncRequirement::Sync(kind) => {
1454 vfs.sync(directory, kind)
1455 .map_err(|source| GraphBuildError::CheckpointIo {
1456 path: directory.to_path_buf(),
1457 source,
1458 })
1459 }
1460 }
1461}
1462
1463fn checkpoint_encoded_bytes(state: &BuildState) -> Result<usize, GraphBuildError> {
1464 let slot_bytes = state
1465 .adjacency
1466 .slots
1467 .len()
1468 .checked_mul(std::mem::size_of::<u32>())
1469 .ok_or_else(|| GraphBuildError::Geometry("checkpoint slot bytes overflow".to_owned()))?;
1470 CHECKPOINT_HEADER_BYTES
1471 .checked_add(state.adjacency.degrees.len())
1472 .and_then(|bytes| bytes.checked_add(slot_bytes))
1473 .and_then(|bytes| bytes.checked_add(CHECKPOINT_CHECKSUM_BYTES))
1474 .ok_or_else(|| GraphBuildError::Geometry("checkpoint encoded bytes overflow".to_owned()))
1475}
1476
1477fn remove_checkpoint(vfs: &dyn Vfs, path: &Path) -> Result<(), GraphBuildError> {
1478 let temporary = checkpoint_temp_path(path);
1479 match vfs.delete(&temporary) {
1480 Ok(()) => {}
1481 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
1482 Err(source) => {
1483 return Err(GraphBuildError::CheckpointIo {
1484 path: temporary,
1485 source,
1486 });
1487 }
1488 }
1489 match vfs.delete(path) {
1490 Ok(()) => Ok(()),
1491 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
1492 Err(source) => Err(GraphBuildError::CheckpointIo {
1493 path: path.to_path_buf(),
1494 source,
1495 }),
1496 }
1497}
1498
1499fn decode_checkpoint(
1500 bytes: &[u8],
1501 vectors: &SegmentVectors<'_>,
1502 params: GraphParams,
1503 seed: u64,
1504 passes: GraphBuildPasses,
1505 accounting: Option<&std::sync::Arc<Accounting>>,
1506) -> Result<BuildState, GraphBuildError> {
1507 validate_graph_build_checkpoint(bytes)?;
1508 let minimum = CHECKPOINT_HEADER_BYTES
1509 .checked_add(CHECKPOINT_CHECKSUM_BYTES)
1510 .ok_or_else(|| GraphBuildError::CheckpointCorrupt("minimum length overflow".to_owned()))?;
1511 if bytes.len() < minimum {
1512 return Err(GraphBuildError::CheckpointCorrupt(format!(
1513 "truncated: need at least {minimum} bytes, got {}",
1514 bytes.len()
1515 )));
1516 }
1517 if bytes.get(..8) != Some(CHECKPOINT_MAGIC.as_slice()) {
1518 return Err(GraphBuildError::CheckpointCorrupt(
1519 "magic does not match ZEVAMCP1".to_owned(),
1520 ));
1521 }
1522 let checksum_start = bytes
1523 .len()
1524 .checked_sub(CHECKPOINT_CHECKSUM_BYTES)
1525 .ok_or_else(|| GraphBuildError::CheckpointCorrupt("checksum underflow".to_owned()))?;
1526 let stored_checksum = read_checkpoint_u64(bytes, checksum_start, "checksum")?;
1527 let checksummed = bytes.get(..checksum_start).ok_or_else(|| {
1528 GraphBuildError::CheckpointCorrupt("checksummed prefix is unavailable".to_owned())
1529 })?;
1530 let actual_checksum = xxh3_64(checksummed);
1531 if stored_checksum != actual_checksum {
1532 return Err(GraphBuildError::CheckpointCorrupt(format!(
1533 "xxh3-64 expected {stored_checksum:#018x}, computed {actual_checksum:#018x}"
1534 )));
1535 }
1536 let version = read_checkpoint_u16(bytes, 8, "version")?;
1537 if version != CHECKPOINT_VERSION {
1538 return Err(GraphBuildError::CheckpointCorrupt(format!(
1539 "version {version}, expected {CHECKPOINT_VERSION}"
1540 )));
1541 }
1542 let phase = decode_phase(read_checkpoint_u8(bytes, 10, "phase")?)?;
1543 let stored_passes = decode_passes(read_checkpoint_u8(bytes, 11, "passes")?)?;
1544 if stored_passes != passes {
1545 return Err(GraphBuildError::CheckpointCorrupt(
1546 "pass count does not match requested build".to_owned(),
1547 ));
1548 }
1549 let stored_r_target = read_checkpoint_u8(bytes, 12, "target degree")?;
1550 let stored_r_max = read_checkpoint_u8(bytes, 13, "maximum degree")?;
1551 let reserved_14 = read_checkpoint_u16(bytes, 14, "reserved degree bytes")?;
1552 let stored_l_build = read_checkpoint_u16(bytes, 16, "construction width")?;
1553 let reserved_18 = read_checkpoint_u16(bytes, 18, "reserved width bytes")?;
1554 let stored_batch = read_checkpoint_u32(bytes, 20, "checkpoint batch")?;
1555 let node_count = read_checkpoint_u32(bytes, 24, "node count")?;
1556 let dimensions = read_checkpoint_u32(bytes, 28, "dimensions")?;
1557 let next_index = read_checkpoint_u32(bytes, 32, "next index")?;
1558 let reserved_36 = read_checkpoint_u32(bytes, 36, "reserved index bytes")?;
1559 let stored_seed = read_checkpoint_u64(bytes, 40, "seed")?;
1560 let stored_alpha_build = f32::from_bits(read_checkpoint_u32(bytes, 48, "build alpha")?);
1561 let stored_alpha_refine = f32::from_bits(read_checkpoint_u32(bytes, 52, "refine alpha")?);
1562 let segment_id = bytes
1563 .get(56..72)
1564 .ok_or_else(|| GraphBuildError::CheckpointCorrupt("segment id is truncated".to_owned()))?;
1565 let payload_bytes = read_checkpoint_u64(bytes, 72, "payload length")?;
1566 if reserved_14 != 0 || reserved_18 != 0 || reserved_36 != 0 {
1567 return Err(GraphBuildError::CheckpointCorrupt(
1568 "reserved header bytes are non-zero".to_owned(),
1569 ));
1570 }
1571 if stored_r_target != params.r_target()
1572 || stored_r_max != params.r_max()
1573 || stored_l_build != params.l_build()
1574 || stored_batch != params.checkpoint_batch_rows()
1575 || stored_alpha_build.to_bits() != params.alpha_build().to_bits()
1576 || stored_alpha_refine.to_bits() != params.alpha_refine().to_bits()
1577 || node_count != vectors.node_count
1578 || dimensions != vectors.dimensions as u32
1579 || stored_seed != seed
1580 || segment_id != vectors.segment_id
1581 {
1582 return Err(GraphBuildError::CheckpointCorrupt(
1583 "checkpoint identity or graph parameters do not match input".to_owned(),
1584 ));
1585 }
1586 if passes == GraphBuildPasses::One && phase == BuildPhase::Refine {
1587 return Err(GraphBuildError::CheckpointCorrupt(
1588 "one-pass build cannot resume refinement".to_owned(),
1589 ));
1590 }
1591 let actual_payload = checksum_start
1592 .checked_sub(CHECKPOINT_HEADER_BYTES)
1593 .ok_or_else(|| GraphBuildError::CheckpointCorrupt("payload underflow".to_owned()))?;
1594 if usize::try_from(payload_bytes).ok() != Some(actual_payload) {
1595 return Err(GraphBuildError::CheckpointCorrupt(format!(
1596 "declared payload {payload_bytes} bytes, actual {actual_payload}"
1597 )));
1598 }
1599 let degree_bytes = node_count as usize;
1600 let degrees_end = CHECKPOINT_HEADER_BYTES
1601 .checked_add(degree_bytes)
1602 .ok_or_else(|| GraphBuildError::CheckpointCorrupt("degree end overflow".to_owned()))?;
1603 let degrees = bytes
1604 .get(CHECKPOINT_HEADER_BYTES..degrees_end)
1605 .ok_or_else(|| GraphBuildError::CheckpointCorrupt("degree bytes are truncated".to_owned()))?
1606 .to_vec();
1607 let expected_slots = (node_count as usize)
1608 .checked_mul(usize::from(params.r_max()))
1609 .ok_or_else(|| GraphBuildError::CheckpointCorrupt("slot count overflow".to_owned()))?;
1610 let slot_payload = bytes.get(degrees_end..checksum_start).ok_or_else(|| {
1611 GraphBuildError::CheckpointCorrupt("neighbor payload is unavailable".to_owned())
1612 })?;
1613 let expected_slot_bytes = expected_slots
1614 .checked_mul(std::mem::size_of::<u32>())
1615 .ok_or_else(|| GraphBuildError::CheckpointCorrupt("slot bytes overflow".to_owned()))?;
1616 if slot_payload.len() != expected_slot_bytes {
1617 return Err(GraphBuildError::CheckpointCorrupt(format!(
1618 "neighbor payload {} bytes, expected {expected_slot_bytes}",
1619 slot_payload.len()
1620 )));
1621 }
1622 let slots = slot_payload
1623 .chunks_exact(4)
1624 .map(|chunk| {
1625 let raw: [u8; 4] = chunk.try_into().map_err(|_| {
1626 GraphBuildError::CheckpointCorrupt("neighbor width is not four bytes".to_owned())
1627 })?;
1628 Ok(u32::from_le_bytes(raw))
1629 })
1630 .collect::<Result<Vec<_>, GraphBuildError>>()?;
1631 let adjacency = Adjacency {
1632 slots,
1633 degrees,
1634 node_count,
1635 r_max: params.r_max(),
1636 };
1637 validate_checkpoint_adjacency(&adjacency)?;
1638 let mut state = BuildState::new(vectors, params, seed, accounting)?;
1639 state.phase = phase;
1640 state.next_index = usize::try_from(next_index)
1641 .map_err(|_| GraphBuildError::CheckpointCorrupt("next index exceeds usize".to_owned()))?;
1642 if state.next_index > state.order.len() {
1643 return Err(GraphBuildError::CheckpointCorrupt(format!(
1644 "next index {} exceeds row count {}",
1645 state.next_index,
1646 state.order.len()
1647 )));
1648 }
1649 state.adjacency = adjacency;
1650 state.inserted.fill(u8::from(phase != BuildPhase::Build));
1651 if phase == BuildPhase::Build {
1652 for node_id in state.order.iter().take(state.next_index) {
1653 let node = vectors.validate_node(*node_id)?;
1654 let marker = state.inserted.get_mut(node).ok_or_else(|| {
1655 GraphBuildError::CheckpointCorrupt("inserted marker is unavailable".to_owned())
1656 })?;
1657 *marker = 1;
1658 }
1659 }
1660 Ok(state)
1661}
1662
1663pub fn validate_graph_build_checkpoint(bytes: &[u8]) -> Result<(), GraphBuildError> {
1667 let minimum = CHECKPOINT_HEADER_BYTES
1668 .checked_add(CHECKPOINT_CHECKSUM_BYTES)
1669 .ok_or_else(|| GraphBuildError::CheckpointCorrupt("minimum length overflow".to_owned()))?;
1670 if bytes.len() < minimum {
1671 return Err(GraphBuildError::CheckpointCorrupt(format!(
1672 "truncated: need at least {minimum} bytes, got {}",
1673 bytes.len()
1674 )));
1675 }
1676 if bytes.get(..8) != Some(CHECKPOINT_MAGIC.as_slice()) {
1677 return Err(GraphBuildError::CheckpointCorrupt(
1678 "magic does not match ZEVAMCP1".to_owned(),
1679 ));
1680 }
1681 let checksum_start = bytes
1682 .len()
1683 .checked_sub(CHECKPOINT_CHECKSUM_BYTES)
1684 .ok_or_else(|| GraphBuildError::CheckpointCorrupt("checksum underflow".to_owned()))?;
1685 let stored_checksum = read_checkpoint_u64(bytes, checksum_start, "checksum")?;
1686 let checksummed = bytes.get(..checksum_start).ok_or_else(|| {
1687 GraphBuildError::CheckpointCorrupt("checksummed prefix is unavailable".to_owned())
1688 })?;
1689 let actual_checksum = xxh3_64(checksummed);
1690 if stored_checksum != actual_checksum {
1691 return Err(GraphBuildError::CheckpointCorrupt(format!(
1692 "xxh3-64 expected {stored_checksum:#018x}, computed {actual_checksum:#018x}"
1693 )));
1694 }
1695 let version = read_checkpoint_u16(bytes, 8, "version")?;
1696 if version != CHECKPOINT_VERSION {
1697 return Err(GraphBuildError::CheckpointCorrupt(format!(
1698 "version {version}, expected {CHECKPOINT_VERSION}"
1699 )));
1700 }
1701 let _ = decode_phase(read_checkpoint_u8(bytes, 10, "phase")?)?;
1702 let _ = decode_passes(read_checkpoint_u8(bytes, 11, "passes")?)?;
1703 let r_target = read_checkpoint_u8(bytes, 12, "target degree")?;
1704 let r_max = read_checkpoint_u8(bytes, 13, "maximum degree")?;
1705 let reserved_14 = read_checkpoint_u16(bytes, 14, "reserved degree bytes")?;
1706 let l_build = read_checkpoint_u16(bytes, 16, "construction width")?;
1707 let reserved_18 = read_checkpoint_u16(bytes, 18, "reserved width bytes")?;
1708 let checkpoint_batch = read_checkpoint_u32(bytes, 20, "checkpoint batch")?;
1709 let node_count = read_checkpoint_u32(bytes, 24, "node count")?;
1710 let dimensions = read_checkpoint_u32(bytes, 28, "dimensions")?;
1711 let next_index = read_checkpoint_u32(bytes, 32, "next position")?;
1712 let reserved_36 = read_checkpoint_u32(bytes, 36, "reserved position bytes")?;
1713 if reserved_14 != 0 || reserved_18 != 0 || reserved_36 != 0 {
1714 return Err(GraphBuildError::CheckpointCorrupt(
1715 "reserved header bytes are nonzero".to_owned(),
1716 ));
1717 }
1718 if r_target == 0 || r_target > r_max || l_build < u16::from(r_target) {
1719 return Err(GraphBuildError::CheckpointCorrupt(
1720 "stored degree and construction width are invalid".to_owned(),
1721 ));
1722 }
1723 if checkpoint_batch == 0 || node_count == 0 || dimensions == 0 || next_index > node_count {
1724 return Err(GraphBuildError::CheckpointCorrupt(
1725 "stored batch, shape, or next position is invalid".to_owned(),
1726 ));
1727 }
1728 let alpha_build = f32::from_bits(read_checkpoint_u32(bytes, 48, "build alpha")?);
1729 let alpha_refine = f32::from_bits(read_checkpoint_u32(bytes, 52, "refine alpha")?);
1730 if !alpha_build.is_finite()
1731 || !alpha_refine.is_finite()
1732 || alpha_build < 1.0
1733 || alpha_refine < alpha_build
1734 {
1735 return Err(GraphBuildError::CheckpointCorrupt(
1736 "stored alpha values are invalid".to_owned(),
1737 ));
1738 }
1739 let payload_bytes = usize::try_from(read_checkpoint_u64(bytes, 72, "payload bytes")?)
1740 .map_err(|_| GraphBuildError::CheckpointCorrupt("payload exceeds usize".to_owned()))?;
1741 let node_count_usize = usize::try_from(node_count)
1742 .map_err(|_| GraphBuildError::CheckpointCorrupt("node count exceeds usize".to_owned()))?;
1743 let slot_count = node_count_usize
1744 .checked_mul(usize::from(r_max))
1745 .ok_or_else(|| GraphBuildError::CheckpointCorrupt("slot count overflow".to_owned()))?;
1746 let expected_payload = slot_count
1747 .checked_mul(std::mem::size_of::<u32>())
1748 .and_then(|slot_bytes| slot_bytes.checked_add(node_count_usize))
1749 .ok_or_else(|| {
1750 GraphBuildError::CheckpointCorrupt("payload geometry overflow".to_owned())
1751 })?;
1752 if payload_bytes != expected_payload {
1753 return Err(GraphBuildError::CheckpointCorrupt(format!(
1754 "payload bytes {payload_bytes}, expected {expected_payload}"
1755 )));
1756 }
1757 let expected_length = CHECKPOINT_HEADER_BYTES
1758 .checked_add(payload_bytes)
1759 .and_then(|length| length.checked_add(CHECKPOINT_CHECKSUM_BYTES))
1760 .ok_or_else(|| GraphBuildError::CheckpointCorrupt("file length overflow".to_owned()))?;
1761 if bytes.len() != expected_length {
1762 return Err(GraphBuildError::CheckpointCorrupt(format!(
1763 "file length {}, expected {expected_length}",
1764 bytes.len()
1765 )));
1766 }
1767 let degrees_start = CHECKPOINT_HEADER_BYTES;
1768 let slots_start = degrees_start
1769 .checked_add(node_count_usize)
1770 .ok_or_else(|| GraphBuildError::CheckpointCorrupt("slot offset overflow".to_owned()))?;
1771 for owner in 0..node_count_usize {
1772 let degree = usize::from(read_checkpoint_u8(
1773 bytes,
1774 degrees_start + owner,
1775 "node degree",
1776 )?);
1777 if degree > usize::from(r_max) {
1778 return Err(GraphBuildError::CheckpointCorrupt(format!(
1779 "node {owner} degree {degree} exceeds {r_max}"
1780 )));
1781 }
1782 let row_start = owner
1783 .checked_mul(usize::from(r_max))
1784 .and_then(|slot| slot.checked_mul(std::mem::size_of::<u32>()))
1785 .and_then(|offset| slots_start.checked_add(offset))
1786 .ok_or_else(|| GraphBuildError::CheckpointCorrupt("row offset overflow".to_owned()))?;
1787 let mut active = Vec::with_capacity(degree);
1788 for slot in 0..usize::from(r_max) {
1789 let offset = slot
1790 .checked_mul(std::mem::size_of::<u32>())
1791 .and_then(|offset| row_start.checked_add(offset))
1792 .ok_or_else(|| {
1793 GraphBuildError::CheckpointCorrupt("neighbor offset overflow".to_owned())
1794 })?;
1795 let neighbor = read_checkpoint_u32(bytes, offset, "neighbor")?;
1796 if slot < degree {
1797 if neighbor >= node_count
1798 || neighbor as usize == owner
1799 || active.contains(&neighbor)
1800 {
1801 return Err(GraphBuildError::CheckpointCorrupt(format!(
1802 "node {owner} has invalid active neighbor {neighbor}"
1803 )));
1804 }
1805 active.push(neighbor);
1806 } else if neighbor != u32::MAX {
1807 return Err(GraphBuildError::CheckpointCorrupt(format!(
1808 "node {owner} inactive slot {slot} is not the sentinel"
1809 )));
1810 }
1811 }
1812 }
1813 Ok(())
1814}
1815
1816fn validate_checkpoint_adjacency(adjacency: &Adjacency) -> Result<(), GraphBuildError> {
1817 for node_id in 0..adjacency.node_count {
1818 let neighbors = adjacency.neighbors(node_id)?;
1819 if neighbors.len() > usize::from(adjacency.r_max) {
1820 return Err(GraphBuildError::CheckpointCorrupt(format!(
1821 "node {node_id} exceeds maximum degree"
1822 )));
1823 }
1824 let mut seen = Vec::with_capacity(neighbors.len());
1825 for neighbor in neighbors {
1826 if *neighbor >= adjacency.node_count || *neighbor == node_id || seen.contains(neighbor)
1827 {
1828 return Err(GraphBuildError::CheckpointCorrupt(format!(
1829 "node {node_id} has an out-of-range, self, or duplicate neighbor {neighbor}"
1830 )));
1831 }
1832 seen.push(*neighbor);
1833 }
1834 let range = adjacency.row_range(node_id)?;
1835 let inactive_start = range.start.checked_add(neighbors.len()).ok_or_else(|| {
1836 GraphBuildError::CheckpointCorrupt("inactive offset overflow".to_owned())
1837 })?;
1838 let inactive = adjacency
1839 .slots
1840 .get(inactive_start..range.end)
1841 .ok_or_else(|| {
1842 GraphBuildError::CheckpointCorrupt("inactive slots are unavailable".to_owned())
1843 })?;
1844 if inactive.iter().any(|neighbor| *neighbor != u32::MAX) {
1845 return Err(GraphBuildError::CheckpointCorrupt(format!(
1846 "node {node_id} has non-sentinel inactive slots"
1847 )));
1848 }
1849 }
1850 Ok(())
1851}
1852
1853fn phase_code(phase: BuildPhase) -> u8 {
1854 match phase {
1855 BuildPhase::Build => 1,
1856 BuildPhase::Refine => 2,
1857 BuildPhase::Complete => 3,
1858 }
1859}
1860
1861fn decode_phase(value: u8) -> Result<BuildPhase, GraphBuildError> {
1862 match value {
1863 1 => Ok(BuildPhase::Build),
1864 2 => Ok(BuildPhase::Refine),
1865 3 => Ok(BuildPhase::Complete),
1866 _ => Err(GraphBuildError::CheckpointCorrupt(format!(
1867 "unknown phase {value}"
1868 ))),
1869 }
1870}
1871
1872fn pass_code(passes: GraphBuildPasses) -> u8 {
1873 match passes {
1874 GraphBuildPasses::One => 1,
1875 GraphBuildPasses::Two => 2,
1876 }
1877}
1878
1879fn decode_passes(value: u8) -> Result<GraphBuildPasses, GraphBuildError> {
1880 match value {
1881 1 => Ok(GraphBuildPasses::One),
1882 2 => Ok(GraphBuildPasses::Two),
1883 _ => Err(GraphBuildError::CheckpointCorrupt(format!(
1884 "unknown pass count {value}"
1885 ))),
1886 }
1887}
1888
1889fn read_checkpoint_u8(bytes: &[u8], offset: usize, field: &str) -> Result<u8, GraphBuildError> {
1890 bytes.get(offset).copied().ok_or_else(|| {
1891 GraphBuildError::CheckpointCorrupt(format!("{field} is truncated at byte {offset}"))
1892 })
1893}
1894
1895fn read_checkpoint_u16(bytes: &[u8], offset: usize, field: &str) -> Result<u16, GraphBuildError> {
1896 let end = offset
1897 .checked_add(2)
1898 .ok_or_else(|| GraphBuildError::CheckpointCorrupt(format!("{field} offset overflow")))?;
1899 let raw: [u8; 2] = bytes
1900 .get(offset..end)
1901 .ok_or_else(|| GraphBuildError::CheckpointCorrupt(format!("{field} is truncated")))?
1902 .try_into()
1903 .map_err(|_| GraphBuildError::CheckpointCorrupt(format!("{field} width is invalid")))?;
1904 Ok(u16::from_le_bytes(raw))
1905}
1906
1907fn read_checkpoint_u32(bytes: &[u8], offset: usize, field: &str) -> Result<u32, GraphBuildError> {
1908 let end = offset
1909 .checked_add(4)
1910 .ok_or_else(|| GraphBuildError::CheckpointCorrupt(format!("{field} offset overflow")))?;
1911 let raw: [u8; 4] = bytes
1912 .get(offset..end)
1913 .ok_or_else(|| GraphBuildError::CheckpointCorrupt(format!("{field} is truncated")))?
1914 .try_into()
1915 .map_err(|_| GraphBuildError::CheckpointCorrupt(format!("{field} width is invalid")))?;
1916 Ok(u32::from_le_bytes(raw))
1917}
1918
1919fn read_checkpoint_u64(bytes: &[u8], offset: usize, field: &str) -> Result<u64, GraphBuildError> {
1920 let end = offset
1921 .checked_add(8)
1922 .ok_or_else(|| GraphBuildError::CheckpointCorrupt(format!("{field} offset overflow")))?;
1923 let raw: [u8; 8] = bytes
1924 .get(offset..end)
1925 .ok_or_else(|| GraphBuildError::CheckpointCorrupt(format!("{field} is truncated")))?
1926 .try_into()
1927 .map_err(|_| GraphBuildError::CheckpointCorrupt(format!("{field} width is invalid")))?;
1928 Ok(u64::from_le_bytes(raw))
1929}
1930
1931#[derive(Clone, Copy, Debug)]
1932struct ScoredNode {
1933 node_id: u32,
1934 distance: f64,
1935}
1936
1937#[allow(clippy::too_many_arguments)]
1938fn process_node(
1939 vectors: &SegmentVectors<'_>,
1940 params: GraphParams,
1941 seed: u64,
1942 alpha: f32,
1943 node_id: u32,
1944 entries: &[u32],
1945 adjacency: &mut Adjacency,
1946 inserted: &mut [u8],
1947 visited: &mut [u32],
1948 visit_epoch: &mut u32,
1949 control: &mut BuildControl<'_>,
1950) -> Result<(), GraphBuildError> {
1951 vectors.validate_node(node_id)?;
1952 let candidates = search_candidates(
1953 vectors,
1954 adjacency,
1955 inserted,
1956 node_id,
1957 entries,
1958 usize::from(params.l_build()),
1959 seed,
1960 visited,
1961 visit_epoch,
1962 control,
1963 )?;
1964 let candidate_bytes = candidates
1965 .capacity()
1966 .checked_mul(std::mem::size_of::<ScoredNode>())
1967 .ok_or_else(|| GraphBuildError::Geometry("candidate capacity overflow".to_owned()))?;
1968 let mut candidate_ids = Vec::new();
1969 let candidate_id_bytes = reserve_vec(
1970 &mut candidate_ids,
1971 candidates
1972 .len()
1973 .saturating_add(adjacency.neighbors(node_id)?.len()),
1974 "candidate id",
1975 control,
1976 )?;
1977 for candidate in candidates {
1978 control.step(1)?;
1979 candidate_ids.push(candidate.node_id);
1980 }
1981 control.release(candidate_bytes)?;
1982 candidate_ids.extend_from_slice(adjacency.neighbors(node_id)?);
1983 let pruned = robust_prune(
1984 vectors,
1985 node_id,
1986 &candidate_ids,
1987 alpha,
1988 usize::from(params.r_target()),
1989 control,
1990 )?;
1991 drop(candidate_ids);
1992 control.release(candidate_id_bytes)?;
1993 adjacency.set_neighbors(node_id, &pruned)?;
1994 let node = vectors.validate_node(node_id)?;
1995 let inserted_node = inserted.get_mut(node).ok_or_else(|| {
1996 GraphBuildError::Geometry(format!("inserted marker {node_id} is unavailable"))
1997 })?;
1998 *inserted_node = 1;
1999 add_reciprocal_edges(vectors, adjacency, node_id, &pruned, params, alpha, control)?;
2000 let pruned_bytes = pruned
2001 .capacity()
2002 .checked_mul(std::mem::size_of::<u32>())
2003 .ok_or_else(|| GraphBuildError::Geometry("pruned capacity overflow".to_owned()))?;
2004 drop(pruned);
2005 control.release(pruned_bytes)?;
2006 Ok(())
2007}
2008
2009#[allow(clippy::too_many_arguments)]
2010fn search_candidates(
2011 vectors: &SegmentVectors<'_>,
2012 adjacency: &Adjacency,
2013 inserted: &[u8],
2014 query_id: u32,
2015 entries: &[u32],
2016 width: usize,
2017 seed: u64,
2018 visited: &mut [u32],
2019 visit_epoch: &mut u32,
2020 control: &mut BuildControl<'_>,
2021) -> Result<Vec<ScoredNode>, GraphBuildError> {
2022 let query = vectors.f32_row(query_id)?;
2023 let mut query_norm_squared = 0.0_f64;
2024 for chunk in query.chunks(256) {
2025 control.step(1)?;
2026 for value in chunk {
2027 query_norm_squared += f64::from(*value) * f64::from(*value);
2028 }
2029 }
2030 let query_requested = query
2031 .len()
2032 .checked_mul(2)
2033 .ok_or_else(|| GraphBuildError::Geometry("query allocation overflow".to_owned()))?;
2034 control.acquire(query_requested)?;
2035 let (prepared, query_peak) =
2036 prepare_bit4_query_controlled(query, seed ^ u64::from(query_id), &mut |units| {
2037 control.step(units)
2038 })
2039 .map_err(|error| match error {
2040 Bit4ControlError::Quant(error) => GraphBuildError::Quant(error),
2041 Bit4ControlError::Memory => GraphBuildError::Store(StoreError::AllocationFailed {
2042 needed: query.len() as u64,
2043 component: "graph query codes",
2044 }),
2045 Bit4ControlError::Control(error) => error,
2046 })?;
2047 control.reconcile(query_requested, query_peak)?;
2048 let query_bytes = prepared.resident_bytes();
2049 control.release(query_peak.checked_sub(query_bytes).ok_or_else(|| {
2050 GraphBuildError::Geometry("query peak accounting underflow".to_owned())
2051 })?)?;
2052 *visit_epoch = visit_epoch.wrapping_add(1);
2053 if *visit_epoch == 0 {
2054 fill_controlled(visited, 0, control)?;
2055 *visit_epoch = 1;
2056 }
2057 let epoch = *visit_epoch;
2058 let mut result = Vec::new();
2059 let _result_bytes = reserve_vec(&mut result, width.saturating_add(1), "result", control)?;
2060 let frontier_capacity = vectors.node_count as usize;
2061 let mut frontier = Vec::new();
2062 let frontier_bytes = reserve_vec(&mut frontier, frontier_capacity, "frontier", control)?;
2063 for &entry in entries {
2064 control.step(1)?;
2065 let entry_index = vectors.validate_node(entry)?;
2066 if entry == query_id
2067 || inserted.get(entry_index).copied() != Some(1)
2068 || visited.get(entry_index).copied() == Some(epoch)
2069 {
2070 continue;
2071 }
2072 if let Some(marker) = visited.get_mut(entry_index) {
2073 *marker = epoch;
2074 }
2075 let scored = score_node(vectors, &prepared, query_norm_squared, entry, control)?;
2076 insert_result(&mut result, scored, width, control)?;
2077 push_fixed(&mut frontier, scored, "frontier")?;
2078 }
2079 if frontier.is_empty()
2080 && let Some(start) = inserted.iter().position(|present| *present == 1)
2081 {
2082 let start = u32::try_from(start)
2083 .map_err(|_| GraphBuildError::Geometry("inserted start exceeds u32".to_owned()))?;
2084 let scored = score_node(vectors, &prepared, query_norm_squared, start, control)?;
2085 if let Some(marker) = visited.get_mut(start as usize) {
2086 *marker = epoch;
2087 }
2088 insert_result(&mut result, scored, width, control)?;
2089 push_fixed(&mut frontier, scored, "frontier")?;
2090 }
2091 while !frontier.is_empty() {
2092 control.step(1)?;
2093 sort_scored(&mut frontier, scored_worst_first, control)?;
2094 let Some(candidate) = frontier.pop() else {
2095 break;
2096 };
2097 let worst = result.last().map(|node| node.distance);
2098 if result.len() >= width && worst.is_some_and(|distance| candidate.distance > distance) {
2099 break;
2100 }
2101 for &neighbor in adjacency.neighbors(candidate.node_id)? {
2102 control.step(1)?;
2103 let neighbor_index = vectors.validate_node(neighbor)?;
2104 if neighbor == query_id
2105 || inserted.get(neighbor_index).copied() != Some(1)
2106 || visited.get(neighbor_index).copied() == Some(epoch)
2107 {
2108 continue;
2109 }
2110 if let Some(marker) = visited.get_mut(neighbor_index) {
2111 *marker = epoch;
2112 }
2113 let scored = score_node(vectors, &prepared, query_norm_squared, neighbor, control)?;
2114 let qualifies = result.len() < width
2115 || result
2116 .last()
2117 .is_some_and(|current_worst| scored.distance < current_worst.distance);
2118 if qualifies {
2119 insert_result(&mut result, scored, width, control)?;
2120 push_fixed(&mut frontier, scored, "frontier")?;
2121 }
2122 }
2123 }
2124 drop(frontier);
2125 control.release(frontier_bytes)?;
2126 drop(prepared);
2127 control.release(query_bytes)?;
2128 Ok(result)
2129}
2130
2131fn score_node(
2132 vectors: &SegmentVectors<'_>,
2133 query: &crate::quant::Bit4Query,
2134 query_norm_squared: f64,
2135 node_id: u32,
2136 control: &mut BuildControl<'_>,
2137) -> Result<ScoredNode, GraphBuildError> {
2138 vectors.validate_node(node_id)?;
2139 let factors = vectors.factor(node_id)?;
2140 let dot = f64::from(
2141 est_dot_bit4_controlled(query, vectors.code_row(node_id)?, factors, &mut |units| {
2142 control.step(units)
2143 })
2144 .map_err(|error| match error {
2145 Bit4ControlError::Quant(error) => GraphBuildError::Quant(error),
2146 Bit4ControlError::Memory => {
2147 GraphBuildError::Geometry("controlled Bit4 score overflow".to_owned())
2148 }
2149 Bit4ControlError::Control(error) => error,
2150 })?,
2151 );
2152 let norm = factors.norm();
2153 Ok(ScoredNode {
2154 node_id,
2155 distance: (query_norm_squared + norm * norm - 2.0 * dot).max(0.0),
2156 })
2157}
2158
2159fn push_fixed<T>(values: &mut Vec<T>, value: T, label: &str) -> Result<(), GraphBuildError> {
2160 if values.len() == values.capacity() {
2161 return Err(GraphBuildError::Geometry(format!(
2162 "{label} fixed capacity exhausted"
2163 )));
2164 }
2165 values.push(value);
2166 Ok(())
2167}
2168
2169fn insert_result(
2170 result: &mut Vec<ScoredNode>,
2171 node: ScoredNode,
2172 width: usize,
2173 control: &mut BuildControl<'_>,
2174) -> Result<(), GraphBuildError> {
2175 push_fixed(result, node, "result")?;
2176 sort_scored(result, scored_best_first, control)?;
2177 if result.len() > width {
2178 let _ = result.pop();
2179 }
2180 Ok(())
2181}
2182
2183fn sort_scored(
2184 values: &mut [ScoredNode],
2185 compare: fn(&ScoredNode, &ScoredNode) -> std::cmp::Ordering,
2186 control: &mut BuildControl<'_>,
2187) -> Result<(), GraphBuildError> {
2188 fn sift(
2189 values: &mut [ScoredNode],
2190 mut root: usize,
2191 end: usize,
2192 compare: fn(&ScoredNode, &ScoredNode) -> std::cmp::Ordering,
2193 control: &mut BuildControl<'_>,
2194 ) -> Result<(), GraphBuildError> {
2195 loop {
2196 let child = root
2197 .checked_mul(2)
2198 .and_then(|value| value.checked_add(1))
2199 .ok_or_else(|| GraphBuildError::Geometry("sort offset overflow".to_owned()))?;
2200 if child >= end {
2201 return Ok(());
2202 }
2203 let mut selected = child;
2204 if child + 1 < end {
2205 control.step(1)?;
2206 if compare(
2207 values.get(child).ok_or_else(|| {
2208 GraphBuildError::Geometry("sort child unavailable".to_owned())
2209 })?,
2210 values.get(child + 1).ok_or_else(|| {
2211 GraphBuildError::Geometry("sort sibling unavailable".to_owned())
2212 })?,
2213 )
2214 .is_lt()
2215 {
2216 selected = child + 1;
2217 }
2218 }
2219 control.step(1)?;
2220 if !compare(
2221 values
2222 .get(root)
2223 .ok_or_else(|| GraphBuildError::Geometry("sort root unavailable".to_owned()))?,
2224 values.get(selected).ok_or_else(|| {
2225 GraphBuildError::Geometry("sort selection unavailable".to_owned())
2226 })?,
2227 )
2228 .is_lt()
2229 {
2230 return Ok(());
2231 }
2232 values.swap(root, selected);
2233 root = selected;
2234 }
2235 }
2236 for root in (0..values.len() / 2).rev() {
2237 sift(values, root, values.len(), compare, control)?;
2238 }
2239 for end in (1..values.len()).rev() {
2240 control.step(1)?;
2241 values.swap(0, end);
2242 sift(values, 0, end, compare, control)?;
2243 }
2244 Ok(())
2245}
2246
2247fn scored_best_first(left: &ScoredNode, right: &ScoredNode) -> std::cmp::Ordering {
2248 left.distance
2249 .total_cmp(&right.distance)
2250 .then_with(|| left.node_id.cmp(&right.node_id))
2251}
2252
2253fn scored_worst_first(left: &ScoredNode, right: &ScoredNode) -> std::cmp::Ordering {
2254 scored_best_first(right, left)
2255}
2256
2257fn robust_prune(
2258 vectors: &SegmentVectors<'_>,
2259 owner: u32,
2260 candidates: &[u32],
2261 alpha: f32,
2262 target: usize,
2263 control: &mut BuildControl<'_>,
2264) -> Result<Vec<u32>, GraphBuildError> {
2265 robust_prune_rows_controlled(
2266 vectors.rescore,
2267 vectors.dimensions,
2268 vectors.node_count,
2269 owner,
2270 candidates,
2271 alpha,
2272 target,
2273 control,
2274 )
2275}
2276
2277pub(crate) fn robust_prune_rows(
2279 rescore: &[f32],
2280 dimensions: usize,
2281 node_count: u32,
2282 owner: u32,
2283 candidates: &[u32],
2284 alpha: f32,
2285 target: usize,
2286) -> Result<Vec<u32>, GraphBuildError> {
2287 robust_prune_rows_controlled(
2288 rescore,
2289 dimensions,
2290 node_count,
2291 owner,
2292 candidates,
2293 alpha,
2294 target,
2295 &mut BuildControl::disabled(),
2296 )
2297}
2298
2299#[allow(clippy::too_many_arguments)]
2300fn robust_prune_rows_controlled(
2301 rescore: &[f32],
2302 dimensions: usize,
2303 node_count: u32,
2304 owner: u32,
2305 candidates: &[u32],
2306 alpha: f32,
2307 target: usize,
2308 control: &mut BuildControl<'_>,
2309) -> Result<Vec<u32>, GraphBuildError> {
2310 validate_exact_rows(rescore, dimensions, node_count)?;
2311 validate_exact_node(owner, node_count)?;
2312 let mut remaining = Vec::new();
2313 let remaining_bytes = reserve_vec(&mut remaining, candidates.len(), "prune", control)?;
2314 for &candidate in candidates {
2315 control.step(1)?;
2316 validate_exact_node(candidate, node_count)?;
2317 if candidate != owner && !remaining.contains(&candidate) {
2318 remaining.push(candidate);
2319 }
2320 }
2321 let mut owner_distances = Vec::new();
2322 let distance_bytes = reserve_vec(&mut owner_distances, remaining.len(), "distance", control)?;
2323 for candidate in remaining {
2324 owner_distances.push(ScoredNode {
2325 node_id: candidate,
2326 distance: exact_row_distance_controlled(
2327 rescore, dimensions, owner, candidate, control,
2328 )?,
2329 });
2330 }
2331 control.release(remaining_bytes)?;
2332 sort_scored(&mut owner_distances, scored_best_first, control)?;
2333 let mut selected = Vec::new();
2334 reserve_vec(&mut selected, target, "selected", control)?;
2335 let alpha_squared = f64::from(alpha) * f64::from(alpha);
2336 while selected.len() < target {
2337 control.step(1)?;
2338 let Some(best) = owner_distances.first().copied() else {
2339 break;
2340 };
2341 selected.push(best.node_id);
2342 let mut index = 0;
2343 while index < owner_distances.len() {
2344 let candidate = *owner_distances.get(index).ok_or_else(|| {
2345 GraphBuildError::Geometry("prune candidate unavailable".to_owned())
2346 })?;
2347 let keep = if candidate.node_id == best.node_id {
2348 false
2349 } else {
2350 let between = exact_row_distance_controlled(
2351 rescore,
2352 dimensions,
2353 best.node_id,
2354 candidate.node_id,
2355 control,
2356 )?;
2357 alpha_squared * between > candidate.distance
2358 };
2359 if keep {
2360 index += 1;
2361 } else {
2362 owner_distances.remove(index);
2363 }
2364 }
2365 }
2366 drop(owner_distances);
2367 control.release(distance_bytes)?;
2368 Ok(selected)
2369}
2370
2371fn validate_exact_rows(
2372 rescore: &[f32],
2373 dimensions: usize,
2374 node_count: u32,
2375) -> Result<(), GraphBuildError> {
2376 if dimensions == 0 {
2377 return Err(GraphBuildError::Geometry(
2378 "alpha prune requires non-zero dimensions".to_owned(),
2379 ));
2380 }
2381 let expected = (node_count as usize)
2382 .checked_mul(dimensions)
2383 .ok_or_else(|| GraphBuildError::Geometry("exact row length overflow".to_owned()))?;
2384 if rescore.len() != expected {
2385 return Err(GraphBuildError::Geometry(format!(
2386 "exact rows contain {} values, expected {expected}",
2387 rescore.len()
2388 )));
2389 }
2390 Ok(())
2391}
2392
2393fn validate_exact_node(node_id: u32, node_count: u32) -> Result<(), GraphBuildError> {
2394 if node_id >= node_count {
2395 return Err(GraphBuildError::NodeIdOutOfRange {
2396 node_id,
2397 node_count,
2398 });
2399 }
2400 Ok(())
2401}
2402
2403fn exact_row_distance_controlled(
2404 rescore: &[f32],
2405 dimensions: usize,
2406 left: u32,
2407 right: u32,
2408 control: &mut BuildControl<'_>,
2409) -> Result<f64, GraphBuildError> {
2410 let row = |node_id: u32| {
2411 let start = (node_id as usize)
2412 .checked_mul(dimensions)
2413 .ok_or_else(|| GraphBuildError::Geometry("exact row offset overflow".to_owned()))?;
2414 let end = start
2415 .checked_add(dimensions)
2416 .ok_or_else(|| GraphBuildError::Geometry("exact row end overflow".to_owned()))?;
2417 rescore
2418 .get(start..end)
2419 .ok_or_else(|| GraphBuildError::Geometry(format!("exact row {node_id} is unavailable")))
2420 };
2421 let mut sum = 0.0;
2422 for (left_chunk, right_chunk) in row(left)?.chunks(256).zip(row(right)?.chunks(256)) {
2423 control.step(1)?;
2424 for (left, right) in left_chunk.iter().zip(right_chunk) {
2425 let difference = f64::from(*left) - f64::from(*right);
2426 sum += difference * difference;
2427 }
2428 }
2429 Ok(sum)
2430}
2431
2432fn add_reciprocal_edges(
2433 vectors: &SegmentVectors<'_>,
2434 adjacency: &mut Adjacency,
2435 owner: u32,
2436 neighbors: &[u32],
2437 params: GraphParams,
2438 alpha: f32,
2439 control: &mut BuildControl<'_>,
2440) -> Result<(), GraphBuildError> {
2441 for &neighbor in neighbors {
2442 control.step(1)?;
2443 vectors.validate_node(neighbor)?;
2444 let source = adjacency.neighbors(neighbor)?;
2445 let mut existing = Vec::new();
2446 let existing_bytes = reserve_vec(
2447 &mut existing,
2448 source.len().saturating_add(1),
2449 "reciprocal",
2450 control,
2451 )?;
2452 existing.extend_from_slice(source);
2453 if existing.contains(&owner) {
2454 drop(existing);
2455 control.release(existing_bytes)?;
2456 continue;
2457 }
2458 let mut separately_allocated = false;
2459 let updated = if existing.len() < usize::from(params.r_max()) {
2460 let mut updated = existing;
2461 updated.push(owner);
2462 updated
2463 } else {
2464 separately_allocated = true;
2465 let mut candidates = existing;
2466 candidates.push(owner);
2467 robust_prune(
2468 vectors,
2469 neighbor,
2470 &candidates,
2471 alpha,
2472 usize::from(params.r_target()),
2473 control,
2474 )?
2475 };
2476 adjacency.set_neighbors(neighbor, &updated)?;
2477 let updated_bytes = updated
2478 .capacity()
2479 .checked_mul(std::mem::size_of::<u32>())
2480 .ok_or_else(|| GraphBuildError::Geometry("reciprocal capacity overflow".to_owned()))?;
2481 drop(updated);
2482 control.release(existing_bytes)?;
2483 if separately_allocated {
2484 control.release(updated_bytes)?;
2485 }
2486 }
2487 Ok(())
2488}
2489
2490fn refined_entry_points(
2491 vectors: &SegmentVectors<'_>,
2492 shuffled: &[u32],
2493 control: &mut BuildControl<'_>,
2494) -> Result<Vec<u32>, GraphBuildError> {
2495 let sample = shuffled
2496 .get(..shuffled.len().min(MEDOID_SAMPLE_ROWS))
2497 .ok_or_else(|| GraphBuildError::Geometry("medoid sample is unavailable".to_owned()))?;
2498 let mut medoid_scores = Vec::new();
2499 let medoid_bytes = reserve_vec(&mut medoid_scores, sample.len(), "medoid", control)?;
2500 for candidate in sample {
2501 let mut distance = 0.0;
2502 for other in sample {
2503 distance += exact_row_distance_controlled(
2504 vectors.rescore,
2505 vectors.dimensions,
2506 *candidate,
2507 *other,
2508 control,
2509 )?;
2510 }
2511 medoid_scores.push(ScoredNode {
2512 node_id: *candidate,
2513 distance,
2514 });
2515 }
2516 sort_scored(&mut medoid_scores, scored_best_first, control)?;
2517 let medoid = medoid_scores
2518 .first()
2519 .map(|node| node.node_id)
2520 .ok_or_else(|| GraphBuildError::Geometry("medoid sample contains no rows".to_owned()))?;
2521 drop(medoid_scores);
2522 control.release(medoid_bytes)?;
2523 let mut entries = Vec::new();
2524 reserve_vec(&mut entries, ENTRY_POINT_COUNT, "entry", control)?;
2525 entries.push(medoid);
2526 while entries.len() < ENTRY_POINT_COUNT && entries.len() < sample.len() {
2527 let mut farthest = None::<ScoredNode>;
2528 for &candidate in sample {
2529 control.step(1)?;
2530 if entries.contains(&candidate) {
2531 continue;
2532 }
2533 let mut nearest = f64::INFINITY;
2534 for entry in &entries {
2535 nearest = nearest.min(exact_row_distance_controlled(
2536 vectors.rescore,
2537 vectors.dimensions,
2538 candidate,
2539 *entry,
2540 control,
2541 )?);
2542 }
2543 let scored = ScoredNode {
2544 node_id: candidate,
2545 distance: nearest,
2546 };
2547 if farthest.is_none_or(|current| {
2548 scored.distance > current.distance
2549 || (scored.distance == current.distance && scored.node_id < current.node_id)
2550 }) {
2551 farthest = Some(scored);
2552 }
2553 }
2554 let Some(next) = farthest else {
2555 break;
2556 };
2557 entries.push(next.node_id);
2558 }
2559 Ok(entries)
2560}
2561
2562fn move_entries_to_front(
2563 order: &mut Vec<u32>,
2564 entries: &[u32],
2565 control: &mut BuildControl<'_>,
2566) -> Result<(), GraphBuildError> {
2567 let mut index = 0;
2568 while index < order.len() {
2569 control.step(1)?;
2570 let node = *order
2571 .get(index)
2572 .ok_or_else(|| GraphBuildError::Geometry("build order row unavailable".to_owned()))?;
2573 if entries.contains(&node) {
2574 for _ in 0..order.len().saturating_sub(index).div_ceil(256) {
2575 control.step(1)?;
2576 }
2577 order.remove(index);
2578 } else {
2579 index += 1;
2580 }
2581 }
2582 for &entry in entries.iter().rev() {
2583 for _ in 0..order.len().div_ceil(256) {
2584 control.step(1)?;
2585 }
2586 order.insert(0, entry);
2587 }
2588 Ok(())
2589}
2590
2591fn encode_artifact(
2592 vectors: &SegmentVectors<'_>,
2593 params: GraphParams,
2594 entries: &[u32],
2595 adjacency: &Adjacency,
2596 memory: Option<AccountedCounter>,
2597) -> Result<GraphBuildArtifact, GraphBuildError> {
2598 encode_artifact_controlled(
2599 vectors,
2600 params,
2601 entries,
2602 adjacency,
2603 memory,
2604 &mut BuildControl::disabled(),
2605 )
2606}
2607
2608fn encode_artifact_controlled(
2609 vectors: &SegmentVectors<'_>,
2610 params: GraphParams,
2611 entries: &[u32],
2612 adjacency: &Adjacency,
2613 mut memory: Option<AccountedCounter>,
2614 control: &mut BuildControl<'_>,
2615) -> Result<GraphBuildArtifact, GraphBuildError> {
2616 let dimensions = u32::try_from(vectors.dimensions)
2617 .map_err(|_| GraphBuildError::Geometry("dimensions exceed u32".to_owned()))?;
2618 let padded_dimensions = dimensions
2619 .checked_add(127)
2620 .map(|value| value / 128 * 128)
2621 .ok_or_else(|| GraphBuildError::Geometry("padded dimensions overflow".to_owned()))?;
2622 let layout = GraphNodeLayout::new(dimensions, padded_dimensions, params.r_max())?;
2623 let padded_row_bytes = layout.code_bytes();
2624 let padded_length = (vectors.node_count as usize)
2625 .checked_mul(padded_row_bytes)
2626 .ok_or_else(|| GraphBuildError::Geometry("padded code bytes overflow".to_owned()))?;
2627 let encoded_length = (vectors.node_count as usize)
2628 .checked_mul(layout.stride() as usize)
2629 .and_then(|bytes| bytes.checked_add(NODE_BLOCK_TRAILER_LEN))
2630 .ok_or_else(|| GraphBuildError::Geometry("encoded graph bytes overflow".to_owned()))?;
2631 let node_inputs = (vectors.node_count as usize)
2632 .checked_mul(std::mem::size_of::<GraphNodeBlockInput<'_>>())
2633 .ok_or_else(|| GraphBuildError::Geometry("node input bytes overflow".to_owned()))?;
2634 let entry_bytes = entries
2635 .len()
2636 .checked_mul(std::mem::size_of::<u32>())
2637 .ok_or_else(|| GraphBuildError::Geometry("entry bytes overflow".to_owned()))?;
2638 if let Some(charge) = memory.as_mut() {
2639 let existing = usize::try_from(charge.bytes())
2640 .map_err(|_| GraphBuildError::Geometry("accounted bytes exceed usize".to_owned()))?;
2641 let peak = existing
2642 .checked_add(padded_length)
2643 .and_then(|bytes| bytes.checked_add(node_inputs))
2644 .and_then(|bytes| bytes.checked_add(encoded_length))
2645 .and_then(|bytes| bytes.checked_add(entry_bytes))
2646 .ok_or_else(|| GraphBuildError::Geometry("encode arena overflow".to_owned()))?;
2647 charge.set(peak)?;
2648 }
2649 let mut padded_codes = Vec::new();
2650 let padded_bytes = reserve_vec(&mut padded_codes, padded_length, "padded code", control)?;
2651 resize_vec_controlled(&mut padded_codes, padded_length, 0, "padded code", control)?;
2652 for node_id in 0..vectors.node_count {
2653 control.step(1)?;
2654 let source = vectors.code_row(node_id)?;
2655 let start = (node_id as usize)
2656 .checked_mul(padded_row_bytes)
2657 .ok_or_else(|| GraphBuildError::Geometry("padded code offset overflow".to_owned()))?;
2658 let end = start
2659 .checked_add(source.len())
2660 .ok_or_else(|| GraphBuildError::Geometry("padded code end overflow".to_owned()))?;
2661 let destination = padded_codes.get_mut(start..end).ok_or_else(|| {
2662 GraphBuildError::Geometry(format!("padded code row {node_id} is unavailable"))
2663 })?;
2664 for (output, input) in destination.chunks_mut(256).zip(source.chunks(256)) {
2665 control.step(input.len() as u64)?;
2666 output.copy_from_slice(input);
2667 }
2668 }
2669 let mut nodes = Vec::new();
2670 let nodes_bytes = reserve_vec(
2671 &mut nodes,
2672 vectors.node_count as usize,
2673 "node input",
2674 control,
2675 )?;
2676 for node_id in 0..vectors.node_count {
2677 control.step(1)?;
2678 let start = (node_id as usize)
2679 .checked_mul(padded_row_bytes)
2680 .ok_or_else(|| GraphBuildError::Geometry("encoded code offset overflow".to_owned()))?;
2681 let end = start
2682 .checked_add(padded_row_bytes)
2683 .ok_or_else(|| GraphBuildError::Geometry("encoded code end overflow".to_owned()))?;
2684 let codes = padded_codes.get(start..end).ok_or_else(|| {
2685 GraphBuildError::Geometry(format!("encoded code row {node_id} is unavailable"))
2686 })?;
2687 nodes.push(GraphNodeBlockInput {
2688 codes,
2689 factors: vectors.factor(node_id)?,
2690 flags: u8::from(entries.contains(&node_id)),
2691 neighbors: adjacency.neighbors(node_id)?,
2692 });
2693 }
2694 control.acquire(encoded_length)?;
2695 let encoded = encode_node_blocks_controlled(
2696 GraphNodeBlockBuild {
2697 layout,
2698 nodes: &nodes,
2699 },
2700 crate::graph::refine::RefinementPasses::default(),
2701 &mut |units| control.step(units).is_ok(),
2702 )?;
2703 control.reconcile(encoded_length, encoded.resident_bytes())?;
2704 let encoded_region = encoded.into_bytes();
2705 let mut entry_points = Vec::new();
2706 reserve_vec(&mut entry_points, entries.len(), "entry copy", control)?;
2707 for chunk in entries.chunks(256) {
2708 control.step(chunk.len() as u64)?;
2709 entry_points.extend_from_slice(chunk);
2710 }
2711 drop(nodes);
2712 control.release(nodes_bytes)?;
2713 drop(padded_codes);
2714 control.release(padded_bytes)?;
2715 if let Some(charge) = memory.as_mut() {
2716 let retained = encoded_region
2717 .capacity()
2718 .checked_add(
2719 entry_points
2720 .capacity()
2721 .checked_mul(std::mem::size_of::<u32>())
2722 .ok_or_else(|| {
2723 GraphBuildError::Geometry("retained entry bytes overflow".to_owned())
2724 })?,
2725 )
2726 .ok_or_else(|| GraphBuildError::Geometry("retained graph bytes overflow".to_owned()))?;
2727 charge.set(retained)?;
2728 }
2729 Ok(GraphBuildArtifact {
2730 encoded_region,
2731 entry_points,
2732 work_rows_completed: 0,
2733 _memory: memory,
2734 })
2735}
2736
2737fn shuffle_controlled(
2738 values: &mut [u32],
2739 random: &mut SplitMix64,
2740 control: &mut BuildControl<'_>,
2741) -> Result<(), GraphBuildError> {
2742 for upper in (1..values.len()).rev() {
2743 control.step(1)?;
2744 let modulus = u64::try_from(upper + 1)
2745 .map_err(|_| GraphBuildError::Geometry("shuffle width exceeds u64".to_owned()))?;
2746 let selected = usize::try_from(random.next_u64() % modulus)
2747 .map_err(|_| GraphBuildError::Geometry("shuffle index exceeds usize".to_owned()))?;
2748 values.swap(upper, selected);
2749 }
2750 Ok(())
2751}
2752
2753#[derive(Clone, Copy, Debug)]
2754struct SplitMix64 {
2755 state: u64,
2756}
2757
2758impl SplitMix64 {
2759 const fn new(seed: u64) -> Self {
2760 Self { state: seed }
2761 }
2762
2763 fn next_u64(&mut self) -> u64 {
2764 self.state = self.state.wrapping_add(0x9e37_79b9_7f4a_7c15);
2765 let mut value = self.state;
2766 value = (value ^ (value >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
2767 value = (value ^ (value >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
2768 value ^ (value >> 31)
2769 }
2770}
2771
2772#[cfg(test)]
2773#[allow(
2774 clippy::expect_used,
2775 clippy::indexing_slicing,
2776 clippy::panic,
2777 clippy::unwrap_used
2778)]
2779mod tests {
2780 use proptest::prelude::*;
2781 use proptest::test_runner::{Config, RngSeed, TestRunner};
2782 use rand::{Rng, RngCore};
2783 use xxhash_rust::xxh3::xxh3_64;
2784
2785 use crate::graph::block::decode_node_blocks;
2786 use crate::lifecycle::QueryCancellation;
2787 use crate::lifecycle::durability::{CommitTier, DurabilityMode, DurabilityPolicy};
2788 use crate::lifecycle::{
2789 CancelToken, ManualMonotonicClock, OpenOptions, QueryControl, Store, StoreTestDependencies,
2790 };
2791 use crate::meta::{AliveSet, ColumnStoreBuilder, Schema};
2792 use crate::quant::{Bit4Factors, quantize_bit4};
2793 use crate::segment::SegmentId;
2794 use crate::segment::reader::SegmentReader;
2795 use crate::segment::writer::{SegmentBuild, SegmentFactors, write_segment};
2796 use crate::vfs::{CountingVfs, StdVfs, Vfs};
2797 use std::sync::Arc;
2798
2799 use super::{
2800 Adjacency, CheckpointedGraphBuild, GraphBuildError, GraphBuildPasses, GraphBuildSession,
2801 SegmentVectors, build_graph, build_graph_checkpointed, decode_checkpoint,
2802 validate_graph_build_checkpoint,
2803 };
2804 use crate::graph::GraphParams;
2805
2806 fn fixture_reader(directory: &std::path::Path, rows: usize) -> SegmentReader {
2807 let dimensions = 128_usize;
2808 let mut random = crate::test_support::seeded_rng(
2809 "graph::build::vamana_build_is_deterministic_under_seed::fixture",
2810 );
2811 let mut rescore = Vec::with_capacity(rows * dimensions);
2812 let mut codes = Vec::with_capacity(rows * dimensions.div_ceil(2));
2813 let mut factors = Vec::<Bit4Factors>::with_capacity(rows);
2814 for _ in 0..rows {
2815 let row = (0..dimensions)
2816 .map(|_| random.random_range(-2.0_f32..2.0_f32))
2817 .collect::<Vec<_>>();
2818 let start = codes.len();
2819 codes.resize(start + dimensions.div_ceil(2), 0);
2820 factors.push(quantize_bit4(&row, &mut codes[start..]).expect("fixture row quantizes"));
2821 rescore.extend(row);
2822 }
2823 let mut columns = ColumnStoreBuilder::new(Schema::new(Vec::new()).expect("empty schema"));
2824 for row in 0..rows {
2825 columns
2826 .push_row(row as i64, &[])
2827 .expect("timestamp-only fixture row");
2828 }
2829 let columns = columns.finish().expect("fixture columns");
2830 let alive = AliveSet::new(rows as u32);
2831 let id = SegmentId::new(19, [7; 10]);
2832 let policy = DurabilityPolicy::new(DurabilityMode::Derived, CommitTier::None)
2833 .expect("derived policy");
2834 write_segment(
2835 &StdVfs,
2836 directory,
2837 SegmentBuild {
2838 id,
2839 scheme: 4,
2840 dims: dimensions as u32,
2841 codes: &codes,
2842 factors: SegmentFactors::Bit4(&factors),
2843 rescore: &rescore,
2844 columns: &columns,
2845 alive: &alive,
2846 },
2847 policy,
2848 )
2849 .expect("sealed Bit4 fixture");
2850 SegmentReader::open(&StdVfs, &directory.join(id.file_name()), id).expect("fixture reader")
2851 }
2852
2853 #[test]
2854 fn graph_checkpoint_io_goes_through_the_store_vfs() {
2855 let segment_directory = tempfile::tempdir().expect("checkpoint segment directory");
2856 let reader = fixture_reader(segment_directory.path(), 4);
2857 let store_directory = tempfile::tempdir().expect("checkpoint store directory");
2858 let checkpoint = store_directory.path().join("graph-build.checkpoint");
2859 let vfs = Arc::new(CountingVfs::new(StdVfs).with_read_path_filter(checkpoint.clone()));
2860 let store = Store::open_with_test_dependencies(
2861 store_directory.path(),
2862 OpenOptions::default(),
2863 StoreTestDependencies::new(vfs.clone(), Arc::new(ManualMonotonicClock::new())),
2864 )
2865 .expect("checkpoint store");
2866 let lease = store.snapshot().expect("checkpoint snapshot");
2867 let params = GraphParams::new(2, 3, 1.0, 1.0, 3, 1).expect("checkpoint params");
2868 let control = QueryControl::Cancel(CancelToken::new());
2869
2870 vfs.reset();
2871 let interrupted = build_graph_checkpointed(
2872 &store,
2873 &reader,
2874 CheckpointedGraphBuild::new(params, 0x06, GraphBuildPasses::One, &checkpoint, &control)
2875 .with_max_work_rows(1),
2876 &lease,
2877 );
2878 assert!(matches!(
2879 interrupted,
2880 Err(GraphBuildError::BudgetExhausted { rows_completed: 1 })
2881 ));
2882 build_graph_checkpointed(
2883 &store,
2884 &reader,
2885 CheckpointedGraphBuild::new(params, 0x06, GraphBuildPasses::One, &checkpoint, &control),
2886 &lease,
2887 )
2888 .expect("resumed graph build");
2889
2890 assert!(vfs.read_calls() >= 1, "checkpoint read bypassed Vfs");
2891 assert_eq!(
2892 vfs.filtered_read_calls(),
2893 1,
2894 "graph-build.checkpoint read bypassed Vfs"
2895 );
2896 assert!(vfs.write_calls() >= 1, "checkpoint write bypassed Vfs");
2897 assert!(vfs.rename_calls() >= 1, "checkpoint rename bypassed Vfs");
2898 assert!(vfs.delete_calls() >= 1, "checkpoint delete bypassed Vfs");
2899 drop(lease);
2900 store.close().expect("close checkpoint store");
2901 }
2902
2903 #[test]
2904 fn graph_checkpoint_publish_syncs_the_temp_file_and_its_directory() {
2905 let segment_directory = tempfile::tempdir().expect("checkpoint segment directory");
2906 let reader = fixture_reader(segment_directory.path(), 4);
2907 let store_directory = tempfile::tempdir().expect("checkpoint store directory");
2908 let checkpoint = store_directory.path().join("graph-build.checkpoint");
2909 let vfs = Arc::new(CountingVfs::new(StdVfs));
2910 let store = Store::open_with_test_dependencies(
2911 store_directory.path(),
2912 OpenOptions::default().with_durability(DurabilityMode::Durable, CommitTier::Durable),
2913 StoreTestDependencies::new(vfs.clone(), Arc::new(ManualMonotonicClock::new())),
2914 )
2915 .expect("durable checkpoint store");
2916 let lease = store.snapshot().expect("checkpoint snapshot");
2917 let params = GraphParams::new(2, 3, 1.0, 1.0, 3, 1).expect("checkpoint params");
2918 let control = QueryControl::Cancel(CancelToken::new());
2919
2920 vfs.reset();
2921 let interrupted = build_graph_checkpointed(
2922 &store,
2923 &reader,
2924 CheckpointedGraphBuild::new(params, 0x03, GraphBuildPasses::One, &checkpoint, &control)
2925 .with_max_work_rows(1),
2926 &lease,
2927 );
2928 assert!(matches!(
2929 interrupted,
2930 Err(GraphBuildError::BudgetExhausted { rows_completed: 1 })
2931 ));
2932 assert!(
2933 vfs.full_sync_calls() >= 2,
2934 "checkpoint publish skipped syncs"
2935 );
2936 assert!(vfs.rename_calls() >= 1, "checkpoint publish skipped rename");
2937 }
2938
2939 fn checkpoint_with_mutation(
2940 mut bytes: Vec<u8>,
2941 mutation: impl FnOnce(&mut Vec<u8>),
2942 ) -> Vec<u8> {
2943 mutation(&mut bytes);
2944 let checksum_start = bytes.len() - super::CHECKPOINT_CHECKSUM_BYTES;
2945 let checksum = xxh3_64(&bytes[..checksum_start]);
2946 bytes[checksum_start..].copy_from_slice(&checksum.to_le_bytes());
2947 bytes
2948 }
2949
2950 #[test]
2951 fn graph_build_errors_preserve_actionable_values_and_sources() {
2952 let io_error = GraphBuildError::CheckpointIo {
2953 path: std::path::PathBuf::from("blocked.checkpoint"),
2954 source: std::io::Error::new(std::io::ErrorKind::PermissionDenied, "blocked"),
2955 };
2956 assert!(io_error.to_string().contains("blocked.checkpoint"));
2957 assert!(std::error::Error::source(&io_error).is_some());
2958
2959 let errors = [
2960 GraphBuildError::CheckpointCorrupt("bad header".to_owned()),
2961 GraphBuildError::BudgetExhausted { rows_completed: 16 },
2962 GraphBuildError::Cancelled { partial: false },
2963 GraphBuildError::Timeout { partial: false },
2964 GraphBuildError::ReadCancelled { partial: false },
2965 GraphBuildError::Geometry("bad shape".to_owned()),
2966 GraphBuildError::NodeIdOutOfRange {
2967 node_id: 9,
2968 node_count: 3,
2969 },
2970 ];
2971 for error in errors {
2972 assert!(!error.to_string().is_empty());
2973 assert!(std::error::Error::source(&error).is_none());
2974 }
2975
2976 let sourced_errors = [
2977 GraphBuildError::from(crate::segment::SegmentError::Geometry(
2978 "bad segment".to_owned(),
2979 )),
2980 GraphBuildError::from(crate::quant::QuantError::EmptyVector),
2981 GraphBuildError::from(crate::graph::block::GraphNodeError::InvalidHeader(
2982 "bad node header".to_owned(),
2983 )),
2984 GraphBuildError::from(crate::lifecycle::StoreError::NotDirectory {
2985 path: std::path::PathBuf::from("not-a-store"),
2986 }),
2987 ];
2988 for error in sourced_errors {
2989 assert!(!error.to_string().is_empty());
2990 assert!(std::error::Error::source(&error).is_some());
2991 }
2992 }
2993
2994 #[test]
2995 fn built_graph_publishes_through_the_frozen_m2_writer() {
2996 let directory = tempfile::tempdir().expect("publication fixture directory");
2997 let reader = fixture_reader(directory.path(), 32);
2998 let params = GraphParams::new(8, 12, 1.0, 1.2, 24, 32).expect("publication params");
2999 let artifact = build_graph(
3000 &reader,
3001 params,
3002 0x0001_9000_3a71_fac7,
3003 GraphBuildPasses::Two,
3004 )
3005 .expect("publication graph");
3006 let output_id = SegmentId::new(19, [8; 10]);
3007 let policy = DurabilityPolicy::new(DurabilityMode::Derived, CommitTier::None)
3008 .expect("publication policy");
3009 let meta = artifact
3010 .write_segment_with_graph(&StdVfs, directory.path(), &reader, output_id, policy)
3011 .expect("published graph segment");
3012 assert_eq!(meta.id, output_id);
3013 let published = SegmentReader::open(
3014 &StdVfs,
3015 &directory.path().join(output_id.file_name()),
3016 output_id,
3017 )
3018 .expect("published reader");
3019 let graph = published
3020 .graph_node_blocks()
3021 .expect("published graph region");
3022 assert_eq!(graph.node_count(), 32);
3023 assert_eq!(graph.layout().max_degree(), 12);
3024 }
3025
3026 #[test]
3027 fn checked_adjacency_rejects_invalid_node_and_degree_inputs() {
3028 let directory = tempfile::tempdir().expect("checked vector directory");
3029 let reader = fixture_reader(directory.path(), 3);
3030 let vectors = SegmentVectors::new(&reader).expect("checked vectors");
3031 assert!(matches!(
3032 vectors.validate_node(3),
3033 Err(GraphBuildError::NodeIdOutOfRange { .. })
3034 ));
3035
3036 let mut adjacency = Adjacency::new(3, 2).expect("small adjacency");
3037 assert!(matches!(
3038 adjacency.neighbors(3),
3039 Err(GraphBuildError::NodeIdOutOfRange { .. })
3040 ));
3041 assert!(matches!(
3042 adjacency.set_neighbors(0, &[1, 2, 1]),
3043 Err(GraphBuildError::Geometry(_))
3044 ));
3045 assert!(matches!(
3046 adjacency.set_neighbors(0, &[3]),
3047 Err(GraphBuildError::NodeIdOutOfRange { .. })
3048 ));
3049 adjacency
3050 .set_neighbors(0, &[1, 2])
3051 .expect("valid adjacency row");
3052 assert_eq!(adjacency.neighbors(0).expect("active row"), &[1, 2]);
3053
3054 let mut self_edge = Adjacency::new(3, 2).expect("self-edge adjacency");
3055 self_edge.degrees[0] = 1;
3056 self_edge.slots[0] = 0;
3057 assert!(matches!(
3058 super::validate_checkpoint_adjacency(&self_edge),
3059 Err(GraphBuildError::CheckpointCorrupt(_))
3060 ));
3061
3062 let mut inactive_value = Adjacency::new(3, 2).expect("inactive-slot adjacency");
3063 inactive_value.slots[0] = 1;
3064 assert!(matches!(
3065 super::validate_checkpoint_adjacency(&inactive_value),
3066 Err(GraphBuildError::CheckpointCorrupt(_))
3067 ));
3068 assert!(matches!(
3069 super::decode_phase(3),
3070 Ok(super::BuildPhase::Complete)
3071 ));
3072 assert!(matches!(
3073 super::read_checkpoint_u8(&[], 0, "truncated"),
3074 Err(GraphBuildError::CheckpointCorrupt(_))
3075 ));
3076 }
3077
3078 #[test]
3079 fn checkpoint_parser_rejects_semantic_corruption_and_identity_drift() {
3080 let directory = tempfile::tempdir().expect("checkpoint parser directory");
3081 let reader = fixture_reader(directory.path(), 16);
3082 let checkpoint = directory.path().join("semantic.checkpoint");
3083 let params = GraphParams::new(4, 6, 1.0, 1.2, 8, 4).expect("checkpoint params");
3084 let store_directory = tempfile::tempdir().expect("checkpoint store directory");
3085 let store = Store::open(store_directory.path(), OpenOptions::default()).expect("store");
3086 let lease = store.snapshot().expect("checkpoint lease");
3087 let control = QueryControl::Cancel(CancelToken::new());
3088 let cancellation = QueryCancellation::new(&control, &lease);
3089 let policy = DurabilityPolicy::new(DurabilityMode::Derived, CommitTier::None)
3090 .expect("derived policy");
3091 let mut session = GraphBuildSession::new(
3092 &StdVfs,
3093 &reader,
3094 params,
3095 0x0001_9000_3c0d_ec01,
3096 GraphBuildPasses::Two,
3097 &checkpoint,
3098 policy,
3099 None,
3100 )
3101 .expect("checkpoint session");
3102 assert!(
3103 !session
3104 .advance_batch(&cancellation)
3105 .expect("checkpoint batch")
3106 );
3107 let valid = StdVfs.read(&checkpoint).expect("valid checkpoint bytes");
3108 validate_graph_build_checkpoint(&valid).expect("valid checkpoint parses");
3109
3110 let mut bad_checksum = valid.clone();
3111 bad_checksum[40] ^= 1;
3112 let corruptions = [
3113 bad_checksum,
3114 checkpoint_with_mutation(valid.clone(), |bytes| {
3115 bytes[8..10].copy_from_slice(&2_u16.to_le_bytes());
3116 }),
3117 checkpoint_with_mutation(valid.clone(), |bytes| bytes[10] = 9),
3118 checkpoint_with_mutation(valid.clone(), |bytes| bytes[11] = 9),
3119 checkpoint_with_mutation(valid.clone(), |bytes| bytes[14] = 1),
3120 checkpoint_with_mutation(valid.clone(), |bytes| bytes[12] = 0),
3121 checkpoint_with_mutation(valid.clone(), |bytes| bytes[20..24].fill(0)),
3122 checkpoint_with_mutation(valid.clone(), |bytes| {
3123 bytes[32..36].copy_from_slice(&17_u32.to_le_bytes());
3124 }),
3125 checkpoint_with_mutation(valid.clone(), |bytes| {
3126 bytes[48..52].copy_from_slice(&f32::NAN.to_bits().to_le_bytes());
3127 }),
3128 checkpoint_with_mutation(valid.clone(), |bytes| bytes[72..80].fill(0)),
3129 checkpoint_with_mutation(valid.clone(), |bytes| {
3130 bytes[80] = 1;
3131 bytes[96..100].copy_from_slice(&0_u32.to_le_bytes());
3132 }),
3133 checkpoint_with_mutation(valid.clone(), |bytes| {
3134 bytes[80] = 0;
3135 bytes[96..100].copy_from_slice(&1_u32.to_le_bytes());
3136 }),
3137 ];
3138 for corrupt in corruptions {
3139 assert!(matches!(
3140 validate_graph_build_checkpoint(&corrupt),
3141 Err(GraphBuildError::CheckpointCorrupt(_))
3142 ));
3143 }
3144
3145 let vectors = SegmentVectors::new(&reader).expect("checkpoint vectors");
3146 assert!(matches!(
3147 decode_checkpoint(
3148 &valid,
3149 &vectors,
3150 params,
3151 0x0001_9000_3c0d_ec02,
3152 GraphBuildPasses::Two,
3153 None,
3154 ),
3155 Err(GraphBuildError::CheckpointCorrupt(_))
3156 ));
3157 assert!(matches!(
3158 decode_checkpoint(
3159 &valid,
3160 &vectors,
3161 params,
3162 0x0001_9000_3c0d_ec01,
3163 GraphBuildPasses::One,
3164 None,
3165 ),
3166 Err(GraphBuildError::CheckpointCorrupt(_))
3167 ));
3168 let one_pass_refine = checkpoint_with_mutation(valid, |bytes| {
3169 bytes[10] = 2;
3170 bytes[11] = 1;
3171 });
3172 assert!(matches!(
3173 decode_checkpoint(
3174 &one_pass_refine,
3175 &vectors,
3176 params,
3177 0x0001_9000_3c0d_ec01,
3178 GraphBuildPasses::One,
3179 None,
3180 ),
3181 Err(GraphBuildError::CheckpointCorrupt(_))
3182 ));
3183 }
3184
3185 #[test]
3186 fn corrupt_checkpoint_refuses_once_clears_the_file_and_the_next_session_starts_fresh() {
3187 let directory = tempfile::tempdir().expect("corrupt checkpoint directory");
3188 let reader = fixture_reader(directory.path(), 16);
3189 let checkpoint = directory.path().join("corrupt.checkpoint");
3190 let temporary = super::checkpoint_temp_path(&checkpoint);
3191 let params = GraphParams::new(4, 6, 1.0, 1.2, 8, 4).expect("checkpoint params");
3192 let store_directory = tempfile::tempdir().expect("checkpoint store directory");
3193 let store = Store::open(store_directory.path(), OpenOptions::default()).expect("store");
3194 let lease = store.snapshot().expect("checkpoint lease");
3195 let control = QueryControl::Cancel(CancelToken::new());
3196 let cancellation = QueryCancellation::new(&control, &lease);
3197 let mut session = GraphBuildSession::new(
3198 &StdVfs,
3199 &reader,
3200 params,
3201 0x0001_9000_c077_0001,
3202 GraphBuildPasses::Two,
3203 &checkpoint,
3204 DurabilityPolicy::new(DurabilityMode::Derived, CommitTier::None)
3205 .expect("derived policy"),
3206 None,
3207 )
3208 .expect("checkpoint session");
3209 assert!(
3210 !session
3211 .advance_batch(&cancellation)
3212 .expect("checkpoint batch")
3213 );
3214 drop(session);
3215
3216 let mut corrupt = StdVfs.read(&checkpoint).expect("valid checkpoint bytes");
3217 corrupt[40] ^= 1;
3218 StdVfs
3219 .write(&checkpoint, &corrupt)
3220 .expect("write corrupt checkpoint");
3221 StdVfs
3222 .write(&temporary, b"stale temporary checkpoint")
3223 .expect("write stale temporary checkpoint");
3224
3225 assert!(matches!(
3226 GraphBuildSession::new(
3227 &StdVfs,
3228 &reader,
3229 params,
3230 0x0001_9000_c077_0001,
3231 GraphBuildPasses::Two,
3232 &checkpoint,
3233 DurabilityPolicy::new(DurabilityMode::Derived, CommitTier::None)
3234 .expect("derived policy"),
3235 None,
3236 ),
3237 Err(GraphBuildError::CheckpointCorrupt(_))
3238 ));
3239 assert_eq!(
3240 StdVfs
3241 .read(&checkpoint)
3242 .expect_err("corrupt checkpoint was retained")
3243 .kind(),
3244 std::io::ErrorKind::NotFound
3245 );
3246 assert_eq!(
3247 StdVfs
3248 .read(&temporary)
3249 .expect_err("stale checkpoint temporary was retained")
3250 .kind(),
3251 std::io::ErrorKind::NotFound
3252 );
3253
3254 let fresh = GraphBuildSession::new(
3255 &StdVfs,
3256 &reader,
3257 params,
3258 0x0001_9000_c077_0001,
3259 GraphBuildPasses::Two,
3260 &checkpoint,
3261 DurabilityPolicy::new(DurabilityMode::Derived, CommitTier::None)
3262 .expect("derived policy"),
3263 None,
3264 )
3265 .expect("next session starts fresh");
3266 assert_eq!(fresh.completed_work_rows().expect("fresh work rows"), 0);
3267 }
3268
3269 #[test]
3270 fn accounting_rejection_on_resume_preserves_the_valid_checkpoint() {
3271 let directory = tempfile::tempdir().expect("accounting checkpoint directory");
3272 let reader = fixture_reader(directory.path(), 16);
3273 let checkpoint = directory.path().join("accounting.checkpoint");
3274 let params = GraphParams::new(4, 6, 1.0, 1.2, 8, 4).expect("checkpoint params");
3275 let store_directory = tempfile::tempdir().expect("checkpoint store directory");
3276 let store = Store::open(store_directory.path(), OpenOptions::default()).expect("store");
3277 let lease = store.snapshot().expect("checkpoint lease");
3278 let control = QueryControl::Cancel(CancelToken::new());
3279 let cancellation = QueryCancellation::new(&control, &lease);
3280 let mut session = GraphBuildSession::new(
3281 &StdVfs,
3282 &reader,
3283 params,
3284 0x0001_9000_acc0_0001,
3285 GraphBuildPasses::Two,
3286 &checkpoint,
3287 DurabilityPolicy::new(DurabilityMode::Derived, CommitTier::None)
3288 .expect("derived policy"),
3289 None,
3290 )
3291 .expect("checkpoint session");
3292 assert!(
3293 !session
3294 .advance_batch(&cancellation)
3295 .expect("checkpoint batch")
3296 );
3297 drop(session);
3298 let valid = StdVfs.read(&checkpoint).expect("valid checkpoint bytes");
3299
3300 let refused_directory = tempfile::tempdir().expect("refused store directory");
3301 let refused = Store::open(
3302 refused_directory.path(),
3303 OpenOptions::default().with_max_temp_bytes(1),
3304 )
3305 .expect("refused store");
3306 assert!(matches!(
3307 GraphBuildSession::new(
3308 &StdVfs,
3309 &reader,
3310 params,
3311 0x0001_9000_acc0_0001,
3312 GraphBuildPasses::Two,
3313 &checkpoint,
3314 DurabilityPolicy::new(DurabilityMode::Derived, CommitTier::None)
3315 .expect("derived policy"),
3316 Some(&refused.accounting),
3317 ),
3318 Err(GraphBuildError::Store(_))
3319 ));
3320 assert_eq!(
3321 StdVfs
3322 .read(&checkpoint)
3323 .expect("accounting rejection removed the valid checkpoint"),
3324 valid
3325 );
3326 }
3327
3328 #[test]
3329 fn incomplete_session_and_checkpoint_io_fail_typed() {
3330 let directory = tempfile::tempdir().expect("typed failure directory");
3331 let reader = fixture_reader(directory.path(), 16);
3332 let params = GraphParams::new(4, 6, 1.0, 1.2, 8, 4).expect("typed failure params");
3333 let control = QueryControl::Cancel(CancelToken::new());
3334 let store_directory = tempfile::tempdir().expect("typed failure store");
3335 let store = Store::open(store_directory.path(), OpenOptions::default()).expect("store");
3336 let lease = store.snapshot().expect("typed failure lease");
3337 let cancellation = QueryCancellation::new(&control, &lease);
3338 let policy = DurabilityPolicy::new(DurabilityMode::Derived, CommitTier::None)
3339 .expect("derived policy");
3340
3341 let unfinished = GraphBuildSession::new(
3342 &StdVfs,
3343 &reader,
3344 params,
3345 0x0001_9000_3100_0001,
3346 GraphBuildPasses::Two,
3347 &directory.path().join("unfinished.checkpoint"),
3348 policy,
3349 None,
3350 )
3351 .expect("unfinished session");
3352 assert!(matches!(
3353 unfinished.artifact(),
3354 Err(GraphBuildError::Geometry(_))
3355 ));
3356
3357 assert!(matches!(
3358 GraphBuildSession::new(
3359 &StdVfs,
3360 &reader,
3361 params,
3362 0x0001_9000_3100_0001,
3363 GraphBuildPasses::Two,
3364 directory.path(),
3365 policy,
3366 None,
3367 ),
3368 Err(GraphBuildError::CheckpointIo { .. })
3369 ));
3370
3371 let missing_parent = directory.path().join("missing").join("write.checkpoint");
3372 let mut unwritable = GraphBuildSession::new(
3373 &StdVfs,
3374 &reader,
3375 params,
3376 0x0001_9000_3100_0001,
3377 GraphBuildPasses::Two,
3378 &missing_parent,
3379 policy,
3380 None,
3381 )
3382 .expect("unwritable session starts before checkpoint write");
3383 assert!(matches!(
3384 unwritable.advance_batch(&cancellation),
3385 Err(GraphBuildError::CheckpointIo { .. })
3386 ));
3387
3388 let timeout = QueryControl::Cancel(CancelToken::new());
3389 timeout.mark_timed_out();
3390 let timed_out = QueryCancellation::new(&timeout, &lease);
3391 assert!(matches!(
3392 super::check_cancellation(&timed_out),
3393 Err(GraphBuildError::Timeout { partial: false })
3394 ));
3395 }
3396
3397 #[test]
3398 fn vamana_build_is_deterministic_under_seed() {
3399 let directory = tempfile::tempdir().expect("fixture directory");
3400 let reader = fixture_reader(directory.path(), 48);
3401 let params = GraphParams::new(8, 12, 1.0, 1.2, 24, 48).expect("test params");
3402 let mut random = crate::test_support::seeded_rng(
3403 "graph::build::vamana_build_is_deterministic_under_seed",
3404 );
3405 let seed = random.next_u64();
3406 let first = build_graph(&reader, params, seed, GraphBuildPasses::Two)
3407 .expect("first deterministic build");
3408 let second = build_graph(&reader, params, seed, GraphBuildPasses::Two)
3409 .expect("second deterministic build");
3410 assert_eq!(first.encoded_region(), second.encoded_region());
3411
3412 let dimensions = 513;
3415 let mut rescore = vec![1.0_f32; dimensions * 2];
3416 rescore[0] = 134_217_728.0;
3417 rescore[dimensions..].fill(0.0);
3418 let code_stride = dimensions.div_ceil(2);
3419 let mut codes = vec![0_u8; code_stride * 2];
3420 let mut factors = Vec::new();
3421 for (row, output) in rescore
3422 .chunks_exact(dimensions)
3423 .zip(codes.chunks_exact_mut(code_stride))
3424 {
3425 factors.push(quantize_bit4(row, output).expect("513D arithmetic fixture quantizes"));
3426 }
3427 let vectors = SegmentVectors::from_native(dimensions, &codes, &factors, &rescore)
3428 .expect("513D arithmetic fixture geometry");
3429 let exact = super::exact_row_distance_controlled(
3430 &rescore,
3431 dimensions,
3432 0,
3433 1,
3434 &mut super::BuildControl::disabled(),
3435 )
3436 .expect("controlled exact row distance");
3437 let adjacency = Adjacency::new(2, 1).expect("two-row arithmetic adjacency");
3438 let candidates = super::search_candidates(
3439 &vectors,
3440 &adjacency,
3441 &[0, 1],
3442 0,
3443 &[1],
3444 1,
3445 seed,
3446 &mut [0; 2],
3447 &mut 0,
3448 &mut super::BuildControl::disabled(),
3449 )
3450 .expect("controlled candidate search");
3451 assert_eq!(candidates.len(), 1);
3452 assert_eq!(candidates[0].node_id, 1);
3453 assert_eq!(
3454 (exact.to_bits(), candidates[0].distance.to_bits()),
3455 (0x4350_0000_0000_0000, 0x4350_0000_0000_0000),
3456 "exact distance and candidate query norm retain sequential arithmetic",
3457 );
3458 #[cfg(feature = "allocation-audit")]
3459 {
3460 let mut allocation_observations = Vec::new();
3461 for ordinal in [1, 2] {
3462 let mut visited = [0; 2];
3463 let mut visit_epoch = 0;
3464 let mut control = super::BuildControl::disabled();
3465 let (result, fires) =
3466 crate::allocation_audit::fail_attributed_allocation(ordinal, || {
3467 super::search_candidates(
3468 &vectors,
3469 &adjacency,
3470 &[0, 1],
3471 0,
3472 &[1],
3473 1,
3474 seed,
3475 &mut visited,
3476 &mut visit_epoch,
3477 &mut control,
3478 )
3479 });
3480 allocation_observations.push((
3481 fires,
3482 matches!(
3483 result,
3484 Err(GraphBuildError::Store(
3485 crate::lifecycle::StoreError::AllocationFailed { needed: 513, .. }
3486 ))
3487 ),
3488 ));
3489 }
3490 assert_eq!(
3491 allocation_observations,
3492 [(1, true), (1, true)],
3493 "both actual query-code allocation failures retain their cause"
3494 );
3495 let restored = super::search_candidates(
3496 &vectors,
3497 &adjacency,
3498 &[0, 1],
3499 0,
3500 &[1],
3501 1,
3502 seed,
3503 &mut [0; 2],
3504 &mut 0,
3505 &mut super::BuildControl::disabled(),
3506 )
3507 .expect("restored query allocation control");
3508 assert_eq!(restored.len(), 1);
3509 assert_eq!(restored[0].node_id, 1);
3510 assert_eq!(restored[0].distance.to_bits(), 0x4350_0000_0000_0000);
3511 }
3512 }
3513
3514 #[test]
3515 fn cancelled_build_leaves_no_partial_artifact() {
3516 let directory = tempfile::tempdir().expect("sealed segment directory");
3517 let reader = fixture_reader(directory.path(), 48);
3518 let segment_path = directory.path().join(reader.meta().id.file_name());
3519 let before = StdVfs
3520 .read(&segment_path)
3521 .expect("sealed bytes before build");
3522 let checkpoint = directory.path().join("graph-build.checkpoint");
3523 let checkpoint_temp = directory.path().join("graph-build.checkpoint.tmp");
3524 let params = GraphParams::new(8, 12, 1.0, 1.2, 24, 8).expect("test params");
3525 let store_directory = tempfile::tempdir().expect("store directory");
3526 let store = Store::open(store_directory.path(), OpenOptions::default()).expect("store");
3527 let lease = store.snapshot().expect("snapshot lease");
3528 let token = CancelToken::new();
3529 let control = QueryControl::Cancel(token.clone());
3530 let cancellation = QueryCancellation::new(&control, &lease);
3531 let policy = DurabilityPolicy::new(DurabilityMode::Derived, CommitTier::None)
3532 .expect("derived policy");
3533 let mut session = GraphBuildSession::new(
3534 &StdVfs,
3535 &reader,
3536 params,
3537 0x19_0003,
3538 GraphBuildPasses::Two,
3539 &checkpoint,
3540 policy,
3541 None,
3542 )
3543 .expect("build session");
3544
3545 assert!(!session.advance_batch(&cancellation).expect("first batch"));
3546 token.cancel();
3547 assert!(matches!(
3548 session.advance_batch(&cancellation),
3549 Err(GraphBuildError::Cancelled { partial: false })
3550 ));
3551 assert_eq!(
3552 StdVfs
3553 .read(&segment_path)
3554 .expect("sealed bytes after cancel"),
3555 before
3556 );
3557 assert!(
3558 checkpoint.exists(),
3559 "completed batch checkpoint must remain resumable"
3560 );
3561 assert!(
3562 !checkpoint_temp.exists(),
3563 "checkpoint temp must not survive"
3564 );
3565 }
3566
3567 #[test]
3568 fn build_resumes_from_checkpoint_to_the_same_graph() {
3569 let directory = tempfile::tempdir().expect("sealed segment directory");
3570 let reader = fixture_reader(directory.path(), 64);
3571 let checkpoint = directory.path().join("resume.checkpoint");
3572 let params = GraphParams::new(8, 12, 1.0, 1.2, 24, 8).expect("test params");
3573 let store_directory = tempfile::tempdir().expect("store directory");
3574 let store = Store::open(store_directory.path(), OpenOptions::default()).expect("store");
3575 let lease = store.snapshot().expect("snapshot lease");
3576 let control = QueryControl::Cancel(CancelToken::new());
3577 let cancellation = QueryCancellation::new(&control, &lease);
3578 let policy = DurabilityPolicy::new(DurabilityMode::Derived, CommitTier::None)
3579 .expect("derived policy");
3580 let mut interrupted = GraphBuildSession::new(
3581 &StdVfs,
3582 &reader,
3583 params,
3584 0x19_0003_00c0_ffee,
3585 GraphBuildPasses::Two,
3586 &checkpoint,
3587 policy,
3588 None,
3589 )
3590 .expect("interrupted session");
3591 assert!(
3592 !interrupted
3593 .advance_batch(&cancellation)
3594 .expect("first batch")
3595 );
3596 drop(interrupted);
3597
3598 let mut resumed = GraphBuildSession::new(
3599 &StdVfs,
3600 &reader,
3601 params,
3602 0x19_0003_00c0_ffee,
3603 GraphBuildPasses::Two,
3604 &checkpoint,
3605 policy,
3606 None,
3607 )
3608 .expect("resumed session");
3609 while !resumed.advance_batch(&cancellation).expect("resumed batch") {}
3610 let resumed = resumed.artifact().expect("resumed artifact");
3611 let uninterrupted =
3612 build_graph(&reader, params, 0x19_0003_00c0_ffee, GraphBuildPasses::Two)
3613 .expect("uninterrupted artifact");
3614 assert_eq!(resumed.encoded_region(), uninterrupted.encoded_region());
3615 assert!(!checkpoint.exists(), "completed build removes checkpoint");
3616 }
3617
3618 #[test]
3619 fn checkpointed_build_reports_exact_work_for_each_declared_pass() {
3620 let directory = tempfile::tempdir().expect("work-count fixture directory");
3621 let rows = 24_usize;
3622 let reader = fixture_reader(directory.path(), rows);
3623 let params = GraphParams::new(6, 10, 1.0, 1.2, 16, 6).expect("work-count params");
3624 let store_directory = tempfile::tempdir().expect("work-count store directory");
3625 let store = Store::open(store_directory.path(), OpenOptions::default()).expect("store");
3626 let lease = store.snapshot().expect("work-count lease");
3627 let control = QueryControl::Cancel(CancelToken::new());
3628
3629 for (passes, expected_rows, suffix) in [
3630 (GraphBuildPasses::One, rows as u64, "one"),
3631 (GraphBuildPasses::Two, (rows * 2) as u64, "two"),
3632 ] {
3633 let checkpoint = directory
3634 .path()
3635 .join(format!("work-count-{suffix}.checkpoint"));
3636 let artifact = build_graph_checkpointed(
3637 &store,
3638 &reader,
3639 CheckpointedGraphBuild::new(params, 0x19_0003_c0a7, passes, &checkpoint, &control),
3640 &lease,
3641 )
3642 .expect("checkpointed build completes");
3643
3644 assert_eq!(artifact.work_rows_completed(), expected_rows);
3645 assert_eq!(
3646 decode_node_blocks(artifact.encoded_region())
3647 .expect("work-count graph decodes")
3648 .node_count(),
3649 rows as u32
3650 );
3651 assert!(
3652 !checkpoint.exists(),
3653 "completed build removes {suffix} checkpoint"
3654 );
3655 }
3656 }
3657
3658 #[test]
3659 fn graph_build_handles_the_minimum_shape_and_rejects_an_empty_segment() {
3660 let directory = tempfile::tempdir().expect("minimum-shape directory");
3661 let empty = fixture_reader(directory.path(), 0);
3662 assert!(matches!(
3663 SegmentVectors::new(&empty),
3664 Err(GraphBuildError::Geometry(detail)) if detail.contains("at least one row")
3665 ));
3666
3667 let one_row_directory = tempfile::tempdir().expect("one-row directory");
3668 let one_row = fixture_reader(one_row_directory.path(), 1);
3669 let params = GraphParams::new(1, 1, 1.0, 1.0, 1, 1).expect("one-row params");
3670 let artifact = build_graph(&one_row, params, 0x19_0003_0000_0001, GraphBuildPasses::One)
3671 .expect("one-row build succeeds");
3672 let graph = decode_node_blocks(artifact.encoded_region()).expect("one-row graph decodes");
3673 let node = graph.block(0).expect("row zero exists");
3674 assert_eq!(graph.node_count(), 1);
3675 assert_eq!(node.degree(), 0);
3676 assert_eq!(node.flags() & 1, 1);
3677 assert_eq!(artifact.entry_points(), &[0]);
3678 }
3679
3680 #[test]
3681 fn every_node_has_at_most_r_max_neighbours() {
3682 let mut seeded = crate::test_support::seeded_rng(
3683 "graph::build::every_node_has_at_most_r_max_neighbours",
3684 );
3685 let mut runner = TestRunner::new(Config {
3686 cases: 24,
3687 rng_seed: RngSeed::Fixed(seeded.next_u64()),
3688 ..Config::default()
3689 });
3690 let strategy = (6_usize..40, 2_u8..=8, any::<u64>());
3691 let result = runner.run(&strategy, |(rows, r_target, seed)| {
3692 let directory = tempfile::tempdir().expect("property directory");
3693 let reader = fixture_reader(directory.path(), rows);
3694 let r_max = r_target.saturating_add(4);
3695 let params = GraphParams::new(
3696 r_target,
3697 r_max,
3698 1.0,
3699 1.2,
3700 u16::from(r_max).saturating_add(8),
3701 rows as u32,
3702 )
3703 .expect("property params");
3704 let artifact =
3705 build_graph(&reader, params, seed, GraphBuildPasses::Two).expect("property graph");
3706 let graph = decode_node_blocks(artifact.encoded_region()).expect("decoded graph");
3707 prop_assert_eq!(graph.layout().max_degree(), r_max);
3708 for node_id in 0..graph.node_count() {
3709 let block = graph.block(node_id).expect("property node");
3710 prop_assert!(block.degree() <= r_max);
3711 }
3712 Ok(())
3713 });
3714 assert!(result.is_ok(), "property result: {result:?}");
3715 }
3716
3717 #[test]
3718 fn no_node_lists_itself_or_a_duplicate() {
3719 let mut seeded =
3720 crate::test_support::seeded_rng("graph::build::no_node_lists_itself_or_a_duplicate");
3721 let mut runner = TestRunner::new(Config {
3722 cases: 24,
3723 rng_seed: RngSeed::Fixed(seeded.next_u64()),
3724 ..Config::default()
3725 });
3726 let strategy = (6_usize..40, 2_u8..=8, any::<u64>());
3727 let result = runner.run(&strategy, |(rows, r_target, seed)| {
3728 let directory = tempfile::tempdir().expect("property directory");
3729 let reader = fixture_reader(directory.path(), rows);
3730 let r_max = r_target.saturating_add(4);
3731 let params = GraphParams::new(
3732 r_target,
3733 r_max,
3734 1.0,
3735 1.2,
3736 u16::from(r_max).saturating_add(8),
3737 rows as u32,
3738 )
3739 .expect("property params");
3740 let artifact =
3741 build_graph(&reader, params, seed, GraphBuildPasses::Two).expect("property graph");
3742 let graph = decode_node_blocks(artifact.encoded_region()).expect("decoded graph");
3743 for node_id in 0..graph.node_count() {
3744 let block = graph.block(node_id).expect("property node");
3745 let mut unique = std::collections::BTreeSet::new();
3746 for neighbor in block.neighbors_padded().take(usize::from(block.degree())) {
3747 prop_assert_ne!(neighbor, node_id);
3748 prop_assert!(unique.insert(neighbor));
3749 }
3750 }
3751 Ok(())
3752 });
3753 assert!(result.is_ok(), "property result: {result:?}");
3754 }
3755
3756 #[test]
3757 fn graph_is_connected_from_the_entry_points() {
3758 let directory = tempfile::tempdir().expect("connectivity directory");
3759 let reader = fixture_reader(directory.path(), 96);
3760 let params = GraphParams::new(8, 12, 1.0, 1.2, 24, 96).expect("connectivity params");
3761 let artifact = build_graph(
3762 &reader,
3763 params,
3764 0x0001_9000_3c01_1ec7,
3765 GraphBuildPasses::Two,
3766 )
3767 .expect("connectivity graph");
3768 let graph = decode_node_blocks(artifact.encoded_region()).expect("decoded graph");
3769 let mut reached = vec![false; graph.node_count() as usize];
3770 let mut pending = std::collections::VecDeque::new();
3771 for &entry in artifact.entry_points() {
3772 pending.push_back(entry);
3773 }
3774 while let Some(node_id) = pending.pop_front() {
3775 let node = node_id as usize;
3776 if reached[node] {
3777 continue;
3778 }
3779 reached[node] = true;
3780 let block = graph.block(node_id).expect("reachable node");
3781 for neighbor in block.neighbors_padded().take(usize::from(block.degree())) {
3782 if !reached[neighbor as usize] {
3783 pending.push_back(neighbor);
3784 }
3785 }
3786 }
3787 let reached_count = reached.iter().filter(|seen| **seen).count();
3788 assert_eq!(reached_count, graph.node_count() as usize);
3789 }
3790
3791 #[test]
3792 fn graph_build_memory_is_budgeted_and_exactly_accounted() {
3793 let directory = tempfile::tempdir().expect("accounting fixture directory");
3794 let reader = fixture_reader(directory.path(), 48);
3795 let params = GraphParams::new(8, 12, 1.0, 1.2, 24, 48).expect("accounting params");
3796 let control = QueryControl::Cancel(CancelToken::new());
3797
3798 let refused_directory = tempfile::tempdir().expect("refused store directory");
3799 let refused = Store::open(
3800 refused_directory.path(),
3801 OpenOptions::default().with_max_temp_bytes(1),
3802 )
3803 .expect("refused store");
3804 let refused_lease = refused.snapshot().expect("refused lease");
3805 let refused_result = build_graph_checkpointed(
3806 &refused,
3807 &reader,
3808 CheckpointedGraphBuild::new(
3809 params,
3810 0x19_0003_acc0,
3811 GraphBuildPasses::One,
3812 &directory.path().join("refused.checkpoint"),
3813 &control,
3814 ),
3815 &refused_lease,
3816 );
3817 assert!(matches!(refused_result, Err(GraphBuildError::Store(_))));
3818 assert_eq!(refused.stats().expect("refused stats").temporary_bytes, 0);
3819
3820 let store_directory = tempfile::tempdir().expect("accounted store directory");
3821 let store = Store::open(store_directory.path(), OpenOptions::default()).expect("store");
3822 let lease = store.snapshot().expect("accounted lease");
3823 let artifact = build_graph_checkpointed(
3824 &store,
3825 &reader,
3826 CheckpointedGraphBuild::new(
3827 params,
3828 0x19_0003_acc0,
3829 GraphBuildPasses::One,
3830 &directory.path().join("accounted.checkpoint"),
3831 &control,
3832 ),
3833 &lease,
3834 )
3835 .expect("accounted graph");
3836 let expected =
3837 artifact.encoded_region().len() + std::mem::size_of_val(artifact.entry_points());
3838 assert_eq!(
3839 store.stats().expect("live graph stats").temporary_bytes,
3840 expected as u64
3841 );
3842 drop(artifact);
3843 assert_eq!(
3844 store.stats().expect("released graph stats").temporary_bytes,
3845 0
3846 );
3847 }
3848
3849 #[test]
3850 fn arbitrary_checkpoint_bytes_fail_typed_without_panicking() {
3851 let mut seeded = crate::test_support::seeded_rng(
3852 "graph::build::arbitrary_checkpoint_bytes_fail_typed_without_panicking",
3853 );
3854 let mut runner = TestRunner::new(Config {
3855 cases: 64,
3856 rng_seed: RngSeed::Fixed(seeded.next_u64()),
3857 ..Config::default()
3858 });
3859 let result = runner.run(&proptest::collection::vec(any::<u8>(), 0..2048), |bytes| {
3860 prop_assert!(matches!(
3861 validate_graph_build_checkpoint(&bytes),
3862 Err(GraphBuildError::CheckpointCorrupt(_))
3863 ));
3864 Ok(())
3865 });
3866 assert!(result.is_ok(), "property result: {result:?}");
3867 }
3868}