diff --git a/Cargo.lock b/Cargo.lock index b518311a4..dad861dbc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -789,7 +789,7 @@ dependencies = [ [[package]] name = "diskann-garnet" -version = "4.0.1" +version = "4.0.2" dependencies = [ "bytemuck", "crossbeam", diff --git a/diskann-garnet/Cargo.toml b/diskann-garnet/Cargo.toml index c0feea504..e895de001 100644 --- a/diskann-garnet/Cargo.toml +++ b/diskann-garnet/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "diskann-garnet" -version = "4.0.1" +version = "4.0.2" edition = "2024" authors.workspace = true license.workspace = true diff --git a/diskann-garnet/diskann-garnet.nuspec b/diskann-garnet/diskann-garnet.nuspec index aa2f0232e..a00c2edf5 100644 --- a/diskann-garnet/diskann-garnet.nuspec +++ b/diskann-garnet/diskann-garnet.nuspec @@ -2,7 +2,7 @@ diskann-garnet - 4.0.1 + 4.0.2 docs/README.md Microsoft https://github.com/microsoft/DiskANN diff --git a/diskann-garnet/src/provider.rs b/diskann-garnet/src/provider.rs index 5024b7b9d..012aa0455 100644 --- a/diskann-garnet/src/provider.rs +++ b/diskann-garnet/src/provider.rs @@ -212,8 +212,25 @@ impl GarnetProvider { return Err(GarnetProviderError::InvalidQuantizer); } - let quantizer = Box::new(quantization::MinMax8Bit::new(dim, metric_type)?) - as Box; + let quantizer = if let Some(quant_state) = + callbacks.read_varsize_iid::(&context.term(Term::Metadata), QUANT_STATE_KEY) + { + quantization::MinMax8Bit::new_from_bytes(metric_type, &quant_state)? + } else { + let quantizer = quantization::MinMax8Bit::new(dim, metric_type)?; + + if !callbacks.write_iid( + &context.term(Term::Metadata), + QUANT_STATE_KEY, + &quantizer.serialize()?, + ) { + return Err(GarnetError::Write.into()); + } + + quantizer + }; + + let quantizer = Box::new(quantizer) as Box; let canonical_bytes = quantizer.bytes(); // NOTE: Q8 needs no training, so it always starts with backfill @@ -2096,4 +2113,70 @@ mod tests { // quantized assert!(store.full_reads() < store.quant_reads()); } + + /// Test that restarts during phase three quant bootstrap work. + /// Phase three starts once backfill is complete and lasts for the remaining + /// life of the index. + #[test] + fn restart_q8() { + let store = Store::new(); + let ctx = Context::new(0); + let index = create_2d_f32_index(VectorQuantType::Q8, Metric::L2, &store, &ctx); + let provider = index.inner.provider(); + + let mut rng = rand::rng(); + + let mut first_insert = true; + for id in 0..10 { + let v = [rng.random(), rng.random()]; + + if first_insert { + provider.maybe_set_start_point(&ctx, &v).unwrap(); + first_insert = false; + } + + DynIndex::insert( + &index, + &ctx, + &GarnetId::from(bytemuck::bytes_of::(&id)), + bytemuck::cast_slice::(&v), + ) + .unwrap(); + } + + // Drop and re-create the index, keeping the same backing store + let index = create_2d_f32_index(VectorQuantType::Q8, Metric::L2, &store, &ctx); + let provider = index.inner.provider(); + + // There should be saved quant state. + assert!( + provider + .callbacks + .exists_iid(&ctx.term(Term::Metadata), QUANT_STATE_KEY), + "quant state missing" + ); + + // Vectors should serialize identically + + let quantizer = provider.quantizer.as_ref().expect("missing quantizer"); + let mut fv = vec![0f32; 2]; + let mut orig_qv = vec![0u8; quantizer.bytes()]; + let mut qv = vec![0u8; quantizer.bytes()]; + + assert!(provider.callbacks.read_single_iid( + &ctx.term(Term::Vector), + 1, + bytemuck::cast_slice_mut::(&mut fv), + )); + + quantizer.compress(&fv, &mut qv).unwrap(); + + assert!( + provider + .callbacks + .read_single_iid(&ctx.term(Term::Quantized), 1, &mut orig_qv) + ); + + assert_eq!(orig_qv, qv, "quant vectors mismatched"); + } } diff --git a/diskann-garnet/src/quantization.rs b/diskann-garnet/src/quantization.rs index 8faa876c1..ffc7a16cb 100644 --- a/diskann-garnet/src/quantization.rs +++ b/diskann-garnet/src/quantization.rs @@ -5,7 +5,7 @@ use diskann_quantization::{ CompressInto, algorithms::{Transform, TransformKind, transforms::NewTransformError}, alloc::{GlobalAllocator, Poly, ScopedAllocator}, - minmax, + minmax::{self, MinMaxQuantizer}, num::POSITIVE_ONE_F32, spherical::{ self, Data, PreScale, SphericalQuantizer, SupportedMetric, @@ -241,6 +241,7 @@ impl MinMax8Bit { Some(d) => d, None => return Err(GarnetQuantizerError::ZeroDim), }; + let mut rng = rand::rng(); let transform = Transform::new( TransformKind::DoubleHadamard { @@ -257,6 +258,15 @@ impl MinMax8Bit { inner: minmax::MinMaxQuantizer::new(transform, grid_scale), }) } + + pub(crate) fn new_from_bytes( + metric: Metric, + bytes: &[u8], + ) -> Result { + let inner = MinMaxQuantizer::try_deserialize(bytes) + .map_err(|e| GarnetQuantizerError::Deserialization(Box::new(e)))?; + Ok(Self { metric, inner }) + } } impl GarnetQuantizer for MinMax8Bit { @@ -306,7 +316,9 @@ impl GarnetQuantizer for MinMax8Bit { } fn serialize(&self) -> Result, GarnetQuantizerError> { - Err(GarnetQuantizerError::UnsupportedSerialization) + self.inner + .serialize(GlobalAllocator) + .map_err(|e| GarnetQuantizerError::Alloc(Box::new(e))) } fn deserialize(&self, _state: &[u8]) -> Result<(), GarnetQuantizerError> { diff --git a/diskann-quantization/build.rs b/diskann-quantization/build.rs index 14b74482a..7b06e23cb 100644 --- a/diskann-quantization/build.rs +++ b/diskann-quantization/build.rs @@ -34,7 +34,11 @@ mod compile { ) }); let out_dir = env::var_os("OUT_DIR").expect("This should be set by Cargo"); - let input_files = ["transforms_v001.fbs", "spherical_v001.fbs"]; + let input_files = [ + "transforms_v001.fbs", + "spherical_v001.fbs", + "minmax_v001.fbs", + ]; // Let `cargo` know that we want to rerun under the following conditions: // @@ -55,7 +59,7 @@ mod compile { // // We first remove the current contents of the generated `flatbuffers` directory // to ensure we don't leave cruft behind. - let output_dirs = ["transforms", "spherical"]; + let output_dirs = ["transforms", "spherical", "minmax"]; let generated_folder = "./src/flatbuffers"; diff --git a/diskann-quantization/schemas/minmax_v001.fbs b/diskann-quantization/schemas/minmax_v001.fbs new file mode 100644 index 000000000..e962ce2ad --- /dev/null +++ b/diskann-quantization/schemas/minmax_v001.fbs @@ -0,0 +1,12 @@ +include "transforms_v001.fbs"; + +namespace Minmax; + +file_identifier "mq00"; + +table Quantizer { + transform:Transforms.Transform (required); + grid_scale:float32; +} + +root_type Quantizer; diff --git a/diskann-quantization/src/flatbuffers.rs b/diskann-quantization/src/flatbuffers.rs index 0c57e2213..729d67953 100644 --- a/diskann-quantization/src/flatbuffers.rs +++ b/diskann-quantization/src/flatbuffers.rs @@ -78,6 +78,11 @@ pub(crate) mod spherical { import_schema!(quantizer_generated, "spherical/quantizer_generated.rs"); } +pub(crate) mod minmax { + use super::*; + import_schema!(quantizer_generated, "minmax/quantizer_generated.rs"); +} + /// Create a `FlatBufferBuilder` and pass it to the closure. /// /// After the closure runs, finish the serialization with the returned offset and pass diff --git a/diskann-quantization/src/flatbuffers/minmax/quantizer_generated.rs b/diskann-quantization/src/flatbuffers/minmax/quantizer_generated.rs new file mode 100644 index 000000000..2a770aff2 --- /dev/null +++ b/diskann-quantization/src/flatbuffers/minmax/quantizer_generated.rs @@ -0,0 +1,247 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +// automatically generated by the FlatBuffers compiler, do not modify +// @generated +extern crate alloc; +extern crate flatbuffers; +use self::flatbuffers::{EndianScalar, Follow}; +use super::*; +use alloc::boxed::Box; +use alloc::string::{String, ToString}; +use alloc::vec::Vec; +use core::cmp::Ordering; +use core::mem; +pub enum QuantizerOffset {} +#[derive(Copy, Clone, PartialEq)] + +pub struct Quantizer<'a> { + pub _tab: flatbuffers::Table<'a>, +} + +impl<'a> flatbuffers::Follow<'a> for Quantizer<'a> { + type Inner = Quantizer<'a>; + #[inline] + unsafe fn follow(buf: &'a [u8], loc: usize) -> Self::Inner { + Self { + _tab: flatbuffers::Table::new(buf, loc), + } + } +} + +impl<'a> Quantizer<'a> { + pub const VT_TRANSFORM: flatbuffers::VOffsetT = 4; + pub const VT_GRID_SCALE: flatbuffers::VOffsetT = 6; + + #[inline] + pub unsafe fn init_from_table(table: flatbuffers::Table<'a>) -> Self { + Quantizer { _tab: table } + } + #[allow(unused_mut)] + pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr, A: flatbuffers::Allocator + 'bldr>( + _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr, A>, + args: &'args QuantizerArgs<'args>, + ) -> flatbuffers::WIPOffset> { + let mut builder = QuantizerBuilder::new(_fbb); + builder.add_grid_scale(args.grid_scale); + if let Some(x) = args.transform { + builder.add_transform(x); + } + builder.finish() + } + + #[inline] + pub fn transform(&self) -> super::transforms::Transform<'a> { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { + self._tab + .get::>( + Quantizer::VT_TRANSFORM, + None, + ) + .unwrap() + } + } + #[inline] + pub fn grid_scale(&self) -> f32 { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { + self._tab + .get::(Quantizer::VT_GRID_SCALE, Some(0.0)) + .unwrap() + } + } +} + +impl flatbuffers::Verifiable for Quantizer<'_> { + #[inline] + fn run_verifier( + v: &mut flatbuffers::Verifier, + pos: usize, + ) -> Result<(), flatbuffers::InvalidFlatbuffer> { + use self::flatbuffers::Verifiable; + v.visit_table(pos)? + .visit_field::>( + "transform", + Self::VT_TRANSFORM, + true, + )? + .visit_field::("grid_scale", Self::VT_GRID_SCALE, false)? + .finish(); + Ok(()) + } +} +pub struct QuantizerArgs<'a> { + pub transform: Option>>, + pub grid_scale: f32, +} +impl<'a> Default for QuantizerArgs<'a> { + #[inline] + fn default() -> Self { + QuantizerArgs { + transform: None, // required field + grid_scale: 0.0, + } + } +} + +pub struct QuantizerBuilder<'a: 'b, 'b, A: flatbuffers::Allocator + 'a> { + fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a, A>, + start_: flatbuffers::WIPOffset, +} +impl<'a: 'b, 'b, A: flatbuffers::Allocator + 'a> QuantizerBuilder<'a, 'b, A> { + #[inline] + pub fn add_transform( + &mut self, + transform: flatbuffers::WIPOffset>, + ) { + self.fbb_ + .push_slot_always::>( + Quantizer::VT_TRANSFORM, + transform, + ); + } + #[inline] + pub fn add_grid_scale(&mut self, grid_scale: f32) { + self.fbb_ + .push_slot::(Quantizer::VT_GRID_SCALE, grid_scale, 0.0); + } + #[inline] + pub fn new(_fbb: &'b mut flatbuffers::FlatBufferBuilder<'a, A>) -> QuantizerBuilder<'a, 'b, A> { + let start = _fbb.start_table(); + QuantizerBuilder { + fbb_: _fbb, + start_: start, + } + } + #[inline] + pub fn finish(self) -> flatbuffers::WIPOffset> { + let o = self.fbb_.end_table(self.start_); + self.fbb_.required(o, Quantizer::VT_TRANSFORM, "transform"); + flatbuffers::WIPOffset::new(o.value()) + } +} + +impl core::fmt::Debug for Quantizer<'_> { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + let mut ds = f.debug_struct("Quantizer"); + ds.field("transform", &self.transform()); + ds.field("grid_scale", &self.grid_scale()); + ds.finish() + } +} +#[inline] +/// Verifies that a buffer of bytes contains a `Quantizer` +/// and returns it. +/// Note that verification is still experimental and may not +/// catch every error, or be maximally performant. For the +/// previous, unchecked, behavior use +/// `root_as_quantizer_unchecked`. +pub fn root_as_quantizer(buf: &[u8]) -> Result { + flatbuffers::root::(buf) +} +#[inline] +/// Verifies that a buffer of bytes contains a size prefixed +/// `Quantizer` and returns it. +/// Note that verification is still experimental and may not +/// catch every error, or be maximally performant. For the +/// previous, unchecked, behavior use +/// `size_prefixed_root_as_quantizer_unchecked`. +pub fn size_prefixed_root_as_quantizer( + buf: &[u8], +) -> Result { + flatbuffers::size_prefixed_root::(buf) +} +#[inline] +/// Verifies, with the given options, that a buffer of bytes +/// contains a `Quantizer` and returns it. +/// Note that verification is still experimental and may not +/// catch every error, or be maximally performant. For the +/// previous, unchecked, behavior use +/// `root_as_quantizer_unchecked`. +pub fn root_as_quantizer_with_opts<'b, 'o>( + opts: &'o flatbuffers::VerifierOptions, + buf: &'b [u8], +) -> Result, flatbuffers::InvalidFlatbuffer> { + flatbuffers::root_with_opts::>(opts, buf) +} +#[inline] +/// Verifies, with the given verifier options, that a buffer of +/// bytes contains a size prefixed `Quantizer` and returns +/// it. Note that verification is still experimental and may not +/// catch every error, or be maximally performant. For the +/// previous, unchecked, behavior use +/// `root_as_quantizer_unchecked`. +pub fn size_prefixed_root_as_quantizer_with_opts<'b, 'o>( + opts: &'o flatbuffers::VerifierOptions, + buf: &'b [u8], +) -> Result, flatbuffers::InvalidFlatbuffer> { + flatbuffers::size_prefixed_root_with_opts::>(opts, buf) +} +#[inline] +/// Assumes, without verification, that a buffer of bytes contains a Quantizer and returns it. +/// # Safety +/// Callers must trust the given bytes do indeed contain a valid `Quantizer`. +pub unsafe fn root_as_quantizer_unchecked(buf: &[u8]) -> Quantizer { + flatbuffers::root_unchecked::(buf) +} +#[inline] +/// Assumes, without verification, that a buffer of bytes contains a size prefixed Quantizer and returns it. +/// # Safety +/// Callers must trust the given bytes do indeed contain a valid size prefixed `Quantizer`. +pub unsafe fn size_prefixed_root_as_quantizer_unchecked(buf: &[u8]) -> Quantizer { + flatbuffers::size_prefixed_root_unchecked::(buf) +} +pub const QUANTIZER_IDENTIFIER: &str = "mq00"; + +#[inline] +pub fn quantizer_buffer_has_identifier(buf: &[u8]) -> bool { + flatbuffers::buffer_has_identifier(buf, QUANTIZER_IDENTIFIER, false) +} + +#[inline] +pub fn quantizer_size_prefixed_buffer_has_identifier(buf: &[u8]) -> bool { + flatbuffers::buffer_has_identifier(buf, QUANTIZER_IDENTIFIER, true) +} + +#[inline] +pub fn finish_quantizer_buffer<'a, 'b, A: flatbuffers::Allocator + 'a>( + fbb: &'b mut flatbuffers::FlatBufferBuilder<'a, A>, + root: flatbuffers::WIPOffset>, +) { + fbb.finish(root, Some(QUANTIZER_IDENTIFIER)); +} + +#[inline] +pub fn finish_size_prefixed_quantizer_buffer<'a, 'b, A: flatbuffers::Allocator + 'a>( + fbb: &'b mut flatbuffers::FlatBufferBuilder<'a, A>, + root: flatbuffers::WIPOffset>, +) { + fbb.finish_size_prefixed(root, Some(QUANTIZER_IDENTIFIER)); +} diff --git a/diskann-quantization/src/minmax/quantizer.rs b/diskann-quantization/src/minmax/quantizer.rs index d01df3983..4edf5a7eb 100644 --- a/diskann-quantization/src/minmax/quantizer.rs +++ b/diskann-quantization/src/minmax/quantizer.rs @@ -3,6 +3,9 @@ * Licensed under the MIT license. */ +#[cfg(feature = "flatbuffers")] +use thiserror::Error; + use super::vectors::{DataMutRef, FullQueryMut, MinMaxCompensation, MinMaxIP, MinMaxL2Squared}; use core::f32; @@ -15,6 +18,20 @@ use crate::{ num::Positive, scalar::{InputContainsNaN, bit_scale}, }; +#[cfg(feature = "flatbuffers")] +use crate::{ + algorithms::transforms::TransformError, + alloc::{Allocator, AllocatorError, Poly}, + flatbuffers as fb, +}; + +/// The base number of bytes to allocate when attempting to serialize a quantizer. +#[cfg(all(not(test), feature = "flatbuffers"))] +const DEFAULT_SERIALIZED_BYTES: usize = 1024; + +// When testing, use a small value so we trigger the reallocation logic. +#[cfg(all(test, feature = "flatbuffers"))] +const DEFAULT_SERIALIZED_BYTES: usize = 1; /// Recall that from the module-level documentation, MinMaxQuantizer, quantizes X /// into `n` bit vectors as follows - @@ -207,6 +224,90 @@ impl MinMaxQuantizer { }) } } + + /// Attempt to deserialize a FlatBuffer `fb::minmax::Quantizer` (as produced by [`MinMaxQuantizer::serialize`]) into [`MinMaxQuantizer`]. + #[cfg(feature = "flatbuffers")] + #[cfg_attr(docsrs, doc(cfg(feature = "flatbuffers")))] + pub fn try_deserialize(data: &[u8]) -> Result { + // Check that this is one of the known identifiers. + if !fb::minmax::quantizer_buffer_has_identifier(data) { + return Err(DeserializationError::InvalidIdentifier); + } + + let root = fb::minmax::root_as_quantizer(data)?; + + Ok(Self::new( + Transform::try_unpack(GlobalAllocator, root.transform())?, + Positive::new(root.grid_scale()) + .map_err(|_| DeserializationError::GridScaleNotPositive)?, + )) + } + + #[cfg(feature = "flatbuffers")] + pub fn serialize(&self, allocator: A) -> Result, AllocatorError> + where + A: Allocator + std::panic::UnwindSafe, + { + use flatbuffers::FlatBufferBuilder; + + let mut buf = FlatBufferBuilder::new_in(Poly::broadcast( + 0u8, + DEFAULT_SERIALIZED_BYTES, + allocator.clone(), + )?); + + let transform = &self.transform; + + let (root, mut buf) = match std::panic::catch_unwind(move || { + let offset = transform.pack(&mut buf); + + let root = fb::minmax::Quantizer::create( + &mut buf, + &fb::minmax::QuantizerArgs { + transform: Some(offset), + grid_scale: self.grid_scale.into_inner(), + }, + ); + (root, buf) + }) { + Ok(ret) => ret, + Err(err) => match err.downcast_ref::() { + Some(msg) => { + if msg.contains("AllocatorError") { + return Err(AllocatorError); + } else { + std::panic::resume_unwind(err); + } + } + None => std::panic::resume_unwind(err), + }, + }; + + // Finish serializing and then copy out the finished data into a newly allocated buffer. + fb::minmax::finish_quantizer_buffer(&mut buf, root); + Poly::from_iter(buf.finished_data().iter().copied(), allocator) + } +} + +//////////////// +// Flatbuffer // +//////////////// + +#[cfg(feature = "flatbuffers")] +#[cfg_attr(docsrs, doc(cfg(feature = "flatbuffers")))] +#[derive(Debug, Clone, Error)] +#[non_exhaustive] +pub enum DeserializationError { + #[error("unhandled file identifier in flatbuffer")] + InvalidIdentifier, + #[error(transparent)] + InvalidFlatBuffer(#[from] flatbuffers::InvalidFlatbuffer), + #[error(transparent)] + AllocatorError(#[from] AllocatorError), + #[error(transparent)] + TransformError(#[from] TransformError), + #[error("grid-scale is missing or is not positive")] + GridScaleNotPositive, } ///////////////// @@ -670,4 +771,164 @@ mod minmax_quantizer_tests { // This should panic due to assertion in compress_into let _ = quantizer.compress_into(&wrong_vector, encoded.reborrow_mut()); } + + #[cfg(feature = "flatbuffers")] + mod serialization { + use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }; + + use super::*; + use crate::{ + algorithms::{TransformKind, transforms::TargetDim}, + alloc::{AllocatorCore, AllocatorError}, + }; + + /// Build a quantizer backed by a random `DoubleHadamard` transform so that + /// serialization has to round-trip non-trivial transform state (the random + /// sign flips), not just the scalar grid-scale. + fn make_hadamard_quantizer(dim: usize, grid_scale: f32, seed: u64) -> MinMaxQuantizer { + let mut rng = StdRng::seed_from_u64(seed); + let transform = Transform::new( + TransformKind::DoubleHadamard { + target_dim: TargetDim::Same, + }, + NonZeroUsize::new(dim).unwrap(), + Some(&mut rng), + GlobalAllocator, + ) + .unwrap(); + + MinMaxQuantizer::new(transform, Positive::new(grid_scale).unwrap()) + } + + /// Compress `vector` into a freshly allocated canonical byte buffer. + fn compress_canonical( + quantizer: &MinMaxQuantizer, + vector: &[f32], + ) -> Vec + where + Unsigned: Representation, + MinMaxQuantizer: + for<'a, 'b> CompressInto<&'a [f32], DataMutRef<'b, NBITS>, Output = L2Loss>, + { + let dim = quantizer.output_dim(); + let mut bytes = vec![0u8; Data::::canonical_bytes(dim)]; + let data = DataMutRef::::from_canonical_front_mut(&mut bytes, dim).unwrap(); + quantizer.compress_into(vector, data).unwrap(); + bytes + } + + fn test_roundtrip_inner(seed: u64) + where + Unsigned: Representation, + MinMaxQuantizer: + for<'a, 'b> CompressInto<&'a [f32], DataMutRef<'b, NBITS>, Output = L2Loss>, + { + let dim = 16; + let quantizer = make_hadamard_quantizer(dim, 0.9, seed); + + let serialized = quantizer.serialize(GlobalAllocator).unwrap(); + let deserialized = MinMaxQuantizer::try_deserialize(&serialized).unwrap(); + + // Structural properties survive. + assert_eq!(deserialized.dim(), quantizer.dim()); + assert_eq!(deserialized.output_dim(), quantizer.output_dim()); + + // The random transform and grid-scale survive exactly: identical inputs must + // produce byte-identical compressed output after a round-trip. + let distribution = Uniform::new_inclusive::(-1.0, 1.0).unwrap(); + let mut rng = StdRng::seed_from_u64(seed ^ 0x9e37_79b9_7f4a_7c15); + for _ in 0..8 { + let vector: Vec = (&distribution).sample_iter(&mut rng).take(dim).collect(); + assert_eq!( + compress_canonical::(&quantizer, &vector), + compress_canonical::(&deserialized, &vector), + "compressed output differs after serialization round-trip", + ); + } + } + + macro_rules! roundtrip_test { + ($name:ident, $nbits:literal, $seed:literal) => { + #[test] + fn $name() { + test_roundtrip_inner::<$nbits>($seed); + } + }; + } + roundtrip_test!(serialize_roundtrip_1bit, 1, 0x1b3f_8c21_6d90_a54e); + roundtrip_test!(serialize_roundtrip_2bit, 2, 0x77c2_e0a4_35bd_1f08); + roundtrip_test!(serialize_roundtrip_4bit, 4, 0x0a9d_4471_c8e2_9b6f); + roundtrip_test!(serialize_roundtrip_8bit, 8, 0xd4e1_5aa9_3f7c_2210); + + /// Deserializing a buffer without the expected file identifier is rejected + /// (rather than panicking or silently misparsing). + #[test] + fn deserialize_rejects_invalid_identifier() { + // A zero-filled buffer long enough to hold an identifier slot has the wrong + // identifier and must be rejected. + assert!(matches!( + MinMaxQuantizer::try_deserialize(&[0u8; 32]), + Err(DeserializationError::InvalidIdentifier), + )); + + // Corrupting the file identifier of an otherwise-valid buffer is also rejected. + let quantizer = make_hadamard_quantizer(16, 1.0, 0x51ed_9c33); + let mut serialized = quantizer.serialize(GlobalAllocator).unwrap().to_vec(); + // The file identifier occupies bytes 4..8 of a (non size-prefixed) flatbuffer. + serialized[4] ^= 0xff; + assert!(matches!( + MinMaxQuantizer::try_deserialize(&serialized), + Err(DeserializationError::InvalidIdentifier), + )); + } + + /// An allocator that succeeds on its first allocation but fails on every + /// subsequent one, used to exercise the reallocation path in `serialize`. + #[derive(Debug, Clone)] + struct FlakyAllocator { + have_allocated: Arc, + } + + impl FlakyAllocator { + fn new(have_allocated: Arc) -> Self { + Self { have_allocated } + } + } + + // SAFETY: This wraps `GlobalAllocator` and only changes when allocation fails. + unsafe impl AllocatorCore for FlakyAllocator { + fn allocate( + &self, + layout: std::alloc::Layout, + ) -> Result, AllocatorError> { + if self.have_allocated.swap(true, Ordering::Relaxed) { + Err(AllocatorError) + } else { + GlobalAllocator.allocate(layout) + } + } + + unsafe fn deallocate(&self, ptr: std::ptr::NonNull<[u8]>, layout: std::alloc::Layout) { + // SAFETY: Inherited from caller. + unsafe { GlobalAllocator.deallocate(ptr, layout) } + } + } + + /// Serialization must surface an [`AllocatorError`] instead of panicking when the + /// buffer needs to grow but reallocation fails. Under `cfg(test)` the initial + /// buffer is intentionally tiny, so growth (and thus the second allocation) always + /// happens. + #[test] + fn serialize_allocation_failure_does_not_panic() { + let quantizer = make_hadamard_quantizer(16, 1.0, 0x7a11_36bd); + let have_allocated = Arc::new(AtomicBool::new(false)); + let _: AllocatorError = quantizer + .serialize(FlakyAllocator::new(have_allocated.clone())) + .unwrap_err(); + assert!(have_allocated.load(Ordering::Relaxed)); + } + } }