hermes_core/structures/vector/
mod.rs1pub mod index;
26pub mod ivf;
27mod kmeans;
28pub mod quantization;
29
30#[cfg(target_pointer_width = "64")]
39pub(crate) const MAX_DENSE_ANN_PAYLOAD_BYTES: usize = 16 * 1024 * 1024 * 1024;
40#[cfg(not(target_pointer_width = "64"))]
41pub(crate) const MAX_DENSE_ANN_PAYLOAD_BYTES: usize = 512 * 1024 * 1024;
42
43#[cfg(target_pointer_width = "64")]
44const MAX_DENSE_ANN_DECODE_BUDGET: usize = 32 * 1024 * 1024 * 1024;
45#[cfg(not(target_pointer_width = "64"))]
46const MAX_DENSE_ANN_DECODE_BUDGET: usize = MAX_DENSE_ANN_PAYLOAD_BYTES;
47
48const ANN_DECODE_EXPANSION_FACTOR: usize = 8;
49const ANN_DECODE_FIXED_HEADROOM: usize = 16 * 1024 * 1024;
50
51fn decode_ann_with_limit<T: serde::de::DeserializeOwned, const LIMIT: usize>(
52 data: &[u8],
53) -> std::io::Result<(T, usize)> {
54 bincode::serde::decode_from_slice(data, bincode::config::standard().with_limit::<LIMIT>())
55 .map_err(|error| std::io::Error::new(std::io::ErrorKind::InvalidData, error))
56}
57
58#[cfg(target_pointer_width = "64")]
59fn decode_ann_with_relative_limit<T: serde::de::DeserializeOwned>(
60 data: &[u8],
61 budget: usize,
62) -> std::io::Result<(T, usize)> {
63 const MIB: usize = 1024 * 1024;
64 const GIB: usize = 1024 * MIB;
65 if budget <= 32 * MIB {
66 decode_ann_with_limit::<T, { 32 * MIB }>(data)
67 } else if budget <= 128 * MIB {
68 decode_ann_with_limit::<T, { 128 * MIB }>(data)
69 } else if budget <= 512 * MIB {
70 decode_ann_with_limit::<T, { 512 * MIB }>(data)
71 } else if budget <= GIB {
72 decode_ann_with_limit::<T, GIB>(data)
73 } else if budget <= 2 * GIB {
74 decode_ann_with_limit::<T, { 2 * GIB }>(data)
75 } else if budget <= 4 * GIB {
76 decode_ann_with_limit::<T, { 4 * GIB }>(data)
77 } else if budget <= 8 * GIB {
78 decode_ann_with_limit::<T, { 8 * GIB }>(data)
79 } else if budget <= 16 * GIB {
80 decode_ann_with_limit::<T, { 16 * GIB }>(data)
81 } else {
82 decode_ann_with_limit::<T, { 32 * GIB }>(data)
83 }
84}
85
86#[cfg(not(target_pointer_width = "64"))]
87fn decode_ann_with_relative_limit<T: serde::de::DeserializeOwned>(
88 data: &[u8],
89 budget: usize,
90) -> std::io::Result<(T, usize)> {
91 const MIB: usize = 1024 * 1024;
92 if budget <= 32 * MIB {
93 decode_ann_with_limit::<T, { 32 * MIB }>(data)
94 } else if budget <= 128 * MIB {
95 decode_ann_with_limit::<T, { 128 * MIB }>(data)
96 } else {
97 decode_ann_with_limit::<T, { 512 * MIB }>(data)
98 }
99}
100
101pub(crate) fn decode_ann_bincode_exact<T: serde::de::DeserializeOwned>(
104 data: &[u8],
105 description: &str,
106) -> std::io::Result<T> {
107 if data.len() > MAX_DENSE_ANN_PAYLOAD_BYTES {
108 return Err(std::io::Error::new(
109 std::io::ErrorKind::InvalidData,
110 format!(
111 "{description} payload is {} bytes, exceeding the {}-byte decode limit",
112 data.len(),
113 MAX_DENSE_ANN_PAYLOAD_BYTES
114 ),
115 ));
116 }
117 let budget = data
118 .len()
119 .checked_mul(ANN_DECODE_EXPANSION_FACTOR)
120 .and_then(|bytes| bytes.checked_add(ANN_DECODE_FIXED_HEADROOM))
121 .unwrap_or(MAX_DENSE_ANN_DECODE_BUDGET)
122 .min(MAX_DENSE_ANN_DECODE_BUDGET);
123 let (value, consumed) = decode_ann_with_relative_limit(data, budget)?;
124 if consumed != data.len() {
125 return Err(std::io::Error::new(
126 std::io::ErrorKind::InvalidData,
127 format!(
128 "{description} payload contains {} trailing bytes",
129 data.len() - consumed
130 ),
131 ));
132 }
133 Ok(value)
134}
135
136pub use ivf::{CoarseCentroids, CoarseConfig, IvfProbePlan, MultiAssignment, SoarConfig};
138
139pub use quantization::{DistanceTable, PQCodebook, PQConfig};
141
142pub use index::{
144 BinaryCoarseQuantizer, BinaryIvfConfig, BinaryIvfIndex, IVFPQConfig, IVFPQIndex, IvfPqQueryPlan,
145};
146
147#[cfg(test)]
148mod decode_tests {
149 use super::decode_ann_bincode_exact;
150 use super::{ANN_DECODE_EXPANSION_FACTOR, ANN_DECODE_FIXED_HEADROOM};
151 use super::{MAX_DENSE_ANN_DECODE_BUDGET, MAX_DENSE_ANN_PAYLOAD_BYTES};
152
153 #[test]
154 fn ann_bincode_decode_is_exact_and_rejects_large_claims_from_tiny_payloads() {
155 let encoded =
156 bincode::serde::encode_to_vec(vec![1u8, 2, 3], bincode::config::standard()).unwrap();
157 assert_eq!(
158 decode_ann_bincode_exact::<Vec<u8>>(&encoded, "test").unwrap(),
159 [1, 2, 3]
160 );
161
162 let mut trailing = encoded;
163 trailing.push(0);
164 assert!(decode_ann_bincode_exact::<Vec<u8>>(&trailing, "test").is_err());
165
166 let oversized_claim =
170 bincode::serde::encode_to_vec(64usize * 1024 * 1024, bincode::config::standard())
171 .unwrap();
172 assert!(decode_ann_bincode_exact::<Vec<u8>>(&oversized_claim, "test").is_err());
173 }
174
175 #[cfg(target_pointer_width = "64")]
176 #[test]
177 fn billion_scale_binary_payload_fits_checked_decode_budget() {
178 const OBSERVED_BINARY_IVF_BYTES: usize = 8_916_876_069;
182 const { assert!(OBSERVED_BINARY_IVF_BYTES < MAX_DENSE_ANN_PAYLOAD_BYTES) };
183 let budget = OBSERVED_BINARY_IVF_BYTES
184 .checked_mul(ANN_DECODE_EXPANSION_FACTOR)
185 .and_then(|bytes| bytes.checked_add(ANN_DECODE_FIXED_HEADROOM))
186 .unwrap_or(MAX_DENSE_ANN_DECODE_BUDGET)
187 .min(MAX_DENSE_ANN_DECODE_BUDGET);
188 assert_eq!(budget, 32 * 1024 * 1024 * 1024);
189 }
190}