1use std::{error::Error, fmt::Display, num::TryFromIntError, sync::Mutex};
2#[allow(deprecated)]
3use std::{
4 hash::{BuildHasher, SipHasher},
5 io::{Read, Write},
6};
7
8use crate::{
9 bucket::Bucket,
10 config::{
11 BitCount, CounterConfig, CuckooConfiguration, LruAgingStrategy, LruConfig, TtlConfig,
12 },
13};
14
15#[derive(Debug)]
17pub enum ImportError {
18 Io(std::io::Error),
20 Config(crate::config::ConfigError),
22 NonZeroError(TryFromIntError),
24 InvalidHasherName {
26 expected: String,
28 found: String,
30 },
31 InvalidHasherHeader {
33 found: String,
35 },
36}
37
38#[derive(Debug)]
40pub enum ExportError {
41 Io(std::io::Error),
43}
44
45impl Display for ImportError {
46 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
47 match self {
48 ImportError::Io(error) => write!(f, "IO error while importing: {}", error),
49 ImportError::Config(error) => {
50 write!(f, "config error on imported data: {}", error)
51 }
52 ImportError::NonZeroError(error) => {
53 write!(f, "conversion error in imported data: {}", error)
54 }
55 ImportError::InvalidHasherName { expected, found } => write!(
56 f,
57 "Invalid hasher name found in header: \"{}\". Expected: \"{}\"",
58 found, expected
59 ),
60 ImportError::InvalidHasherHeader { found } => write!(
61 f,
62 "Invalid hasher header. Expected \"{}\" at the end, but found \"{}\"",
63 HEADER_END, found
64 ),
65 }
66 }
67}
68
69impl Error for ImportError {
70 fn cause(&self) -> Option<&dyn Error> {
71 match self {
72 ImportError::Io(error) => Some(error),
73 ImportError::Config(error) => Some(error),
74 ImportError::NonZeroError(error) => Some(error),
75 ImportError::InvalidHasherName { .. } => None,
76 ImportError::InvalidHasherHeader { .. } => None,
77 }
78 }
79}
80
81impl From<std::io::Error> for ImportError {
82 fn from(value: std::io::Error) -> Self {
83 Self::Io(value)
84 }
85}
86
87impl From<crate::config::ConfigError> for ImportError {
88 fn from(value: crate::config::ConfigError) -> Self {
89 Self::Config(value)
90 }
91}
92
93impl From<TryFromIntError> for ImportError {
94 fn from(value: TryFromIntError) -> Self {
95 Self::NonZeroError(value)
96 }
97}
98
99impl Display for ExportError {
100 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
101 match self {
102 ExportError::Io(error) => write!(f, "IO error while exporting: {}", error),
103 }
104 }
105}
106
107impl Error for ExportError {
108 fn cause(&self) -> Option<&dyn Error> {
109 match self {
110 ExportError::Io(error) => Some(error),
111 }
112 }
113}
114
115impl From<std::io::Error> for ExportError {
116 fn from(value: std::io::Error) -> Self {
117 Self::Io(value)
118 }
119}
120
121const HEADER: &str = "cuckoo-clock:";
122const HEADER_END: &str = ":header-end\n";
123
124pub trait Exportable: Clone {
128 fn write_to(&self, writer: impl Write) -> Result<(), ExportError>;
130
131 fn read_from(reader: impl Read) -> Result<Self, ImportError>;
133}
134
135pub trait ExportableBuildHasher: Exportable + BuildHasher {
139 const NAME: &str;
142}
143
144#[derive(Clone)]
146pub struct ExportableRandomState {
147 k0: u64,
148 k1: u64,
149}
150
151impl ExportableRandomState {
152 pub(crate) const fn new() -> Self {
153 Self { k0: 0, k1: 0 }
154 }
155
156 #[must_use]
158 pub fn new_random() -> Self {
159 Self {
160 k0: rand::random(),
161 k1: rand::random(),
162 }
163 }
164}
165
166impl Default for ExportableRandomState {
167 fn default() -> Self {
168 Self::new()
169 }
170}
171
172#[allow(deprecated)]
173impl BuildHasher for ExportableRandomState {
174 type Hasher = SipHasher;
175
176 fn build_hasher(&self) -> Self::Hasher {
177 SipHasher::new_with_keys(self.k0, self.k1)
178 }
179}
180
181impl Exportable for ExportableRandomState {
182 fn write_to(&self, mut writer: impl Write) -> Result<(), ExportError> {
183 let k0 = self.k0.to_be_bytes();
184 let k1 = self.k1.to_be_bytes();
185 writer.write_all(&k0)?;
186 writer.write_all(&k1)?;
187 Ok(())
188 }
189
190 fn read_from(mut reader: impl Read) -> Result<Self, ImportError> {
191 let mut buf = [0u8; 8];
192 reader.read_exact(&mut buf)?;
193 let k0 = u64::from_be_bytes(buf);
194 reader.read_exact(&mut buf)?;
195 let k1 = u64::from_be_bytes(buf);
196 Ok(Self { k0, k1 })
197 }
198}
199
200impl ExportableBuildHasher for ExportableRandomState {
201 const NAME: &str = "cuckoo_clock::exporter::ExportableRandomState";
202}
203
204pub struct CuckooFilterExporter<'a, H: ExportableBuildHasher> {
208 hasher: &'a H,
209 buckets: &'a Vec<Mutex<Bucket>>,
210 config: &'a CuckooConfiguration,
211}
212
213impl<'a, H: ExportableBuildHasher> CuckooFilterExporter<'a, H> {
214 pub(crate) const fn new(
215 hasher: &'a H,
216 buckets: &'a Vec<Mutex<Bucket>>,
217 config: &'a CuckooConfiguration,
218 ) -> Self {
219 Self {
220 hasher,
221 buckets,
222 config,
223 }
224 }
225
226 pub fn write_to(&self, mut writer: impl Write) -> Result<(), ExportError> {
233 writer.write_all(format!("{}:{}:", HEADER, H::NAME).as_bytes())?;
234 self.hasher.write_to(&mut writer)?;
235 writer.write_all(HEADER_END.as_bytes())?;
236
237 self.config.write_to(&mut writer)?;
238 writer.write_all(b"\n")?;
239
240 for b in self.buckets.iter() {
241 #[expect(clippy::unwrap_used)]
242 let bucket = b.lock().unwrap();
243 bucket.export(&mut writer)?;
244 }
245
246 Ok(())
247 }
248
249 pub fn snapshot(&self) -> Result<Vec<u8>, ExportError> {
254 let mut result = Vec::new();
255 self.write_to(&mut result)?;
256 Ok(result)
257 }
258}
259
260pub(crate) fn read_hasher_from<H: ExportableBuildHasher + BuildHasher>(
261 mut reader: impl Read,
262) -> Result<H, ImportError> {
263 let mut header_prefix = vec![0; HEADER.len() + H::NAME.len() + 2];
264 reader.read_exact(&mut header_prefix)?;
265 if header_prefix != format!("{}:{}:", HEADER, H::NAME).as_bytes() {
266 return Err(ImportError::InvalidHasherName {
267 expected: H::NAME.to_string(),
268 found: String::from_utf8(header_prefix)
269 .map(|s| {
270 s.trim_start_matches(&format!("{}:", HEADER))
271 .trim_end_matches(":")
272 .to_string()
273 })
274 .unwrap_or("<non UTF-8 data>".to_string()),
275 });
276 };
277 let hasher = H::read_from(&mut reader)?;
278 let mut header_suffix = vec![0; HEADER_END.len()];
279 reader.read_exact(&mut header_suffix)?;
280 if header_suffix != HEADER_END.as_bytes() {
281 return Err(ImportError::InvalidHasherHeader {
282 found: String::from_utf8(header_suffix).unwrap_or("<non UTF-8 data>".to_string()),
283 });
284 };
285 Ok(hasher)
286}
287
288impl Exportable for CuckooConfiguration {
289 fn write_to(&self, mut writer: impl std::io::Write) -> Result<(), crate::ExportError> {
290 writer.write_all(&[0, 0, 2, 4])?;
292 #[allow(clippy::expect_used)]
295 writer.write_all(
297 &u8::try_from((self.fingerprint_field_config.value_mask() as usize + 1).ilog2())
298 .expect("Fingeprint bits can't be higher than 32")
299 .to_be_bytes(),
300 )?;
301 writer.write_all(&(self.bucket_size as u64).to_be_bytes())?;
302 writer.write_all(&((self.bucket_count * self.bucket_size) as u64).to_be_bytes())?;
304 writer.write_all(&(self.max_kicks as u64).to_be_bytes())?;
305 if let Some((lru, _)) = &self.lru_field_config {
306 writer.write_all(&[1])?;
307 lru.write_to(&mut writer)?;
308 }
309 if let Some((ttl, _)) = &self.ttl_field_config {
310 writer.write_all(&[2])?;
311 ttl.write_to(&mut writer)?;
312 }
313 if let Some((counter, _)) = &self.counter_field_config {
314 writer.write_all(&[3])?;
315 counter.write_to(&mut writer)?;
316 }
317 Ok(())
318 }
319
320 fn read_from(mut reader: impl std::io::Read) -> Result<Self, crate::ImportError> {
321 let mut u8_buf = [0u8; 1];
322 let mut u64_buf = [0u8; 8];
323 reader.read_exact(&mut u8_buf)?;
324 let mut version: usize = 0;
325 if u8::from_be_bytes(u8_buf) == 0 {
326 reader.read_exact(&mut u8_buf)?;
328 version += 1_000_000 * usize::from(u8::from_be_bytes(u8_buf));
329 reader.read_exact(&mut u8_buf)?;
330 version += 1_000 * usize::from(u8::from_be_bytes(u8_buf));
331 reader.read_exact(&mut u8_buf)?;
332 version += usize::from(u8::from_be_bytes(u8_buf));
333 reader.read_exact(&mut u8_buf)?;
334 } else {
335 version = 1_000;
337 }
338 let fp_bits: BitCount = usize::from(u8::from_be_bytes(u8_buf)).try_into()?;
339 reader.read_exact(&mut u64_buf)?;
340 let bucket_size = u64::from_be_bytes(u64_buf);
341 reader.read_exact(&mut u64_buf)?;
342 let max_entries = u64::from_be_bytes(u64_buf);
343 reader.read_exact(&mut u64_buf)?;
344 let max_kicks = u64::from_be_bytes(u64_buf);
345 let mut builder = CuckooConfiguration::builder(max_entries.try_into()?)
346 .fingerprint_bits(fp_bits)
347 .bucket_size(usize::try_from(bucket_size)?.try_into()?)
348 .max_kicks(max_kicks.try_into()?);
349 while let Ok(()) = reader.read_exact(&mut u8_buf) {
350 let conf_type = u8::from_be_bytes(u8_buf);
351 match conf_type {
352 1 => {
353 builder = builder.with_lru(LruConfig::read_from(&mut reader, version)?);
354 }
355 2 => {
356 builder = builder.with_ttl(TtlConfig::read_from(&mut reader, version)?);
357 }
358 3 => {
359 builder = builder.with_counter(CounterConfig::read_from(&mut reader, version)?);
360 }
361 _ => break,
364 }
365 }
366 Ok(builder.build()?)
367 }
368}
369
370pub(crate) trait ExportableConfig: Clone {
372 fn write_to(&self, writer: impl Write) -> Result<(), ExportError>;
374
375 fn read_from(reader: impl Read, version: usize) -> Result<Self, ImportError>;
377}
378
379impl ExportableConfig for LruConfig {
380 fn write_to(&self, mut writer: impl Write) -> Result<(), ExportError> {
381 writer.write_all(&u8::from(self.counter_bits).to_be_bytes())?;
382 writer.write_all(&u8::from(self.remove_on_zero).to_be_bytes())?;
383 self.aging_strategy.write_to(&mut writer)?;
384 writer.write_all(&self.starting_value.to_be_bytes())?;
385 writer.write_all(&self.increment.to_be_bytes())?;
386 Ok(())
387 }
388
389 fn read_from(mut reader: impl Read, version: usize) -> Result<Self, ImportError> {
390 let mut u8_buf = [0u8; 1];
391 let mut u32_buf = [0u8; 4];
392 reader.read_exact(&mut u8_buf)?;
393 let bits = u8::from_be_bytes(u8_buf);
394 let mut remove_on_zero = false;
395 let mut aging_strategy = LruAgingStrategy::default();
396 let mut increment = 1;
397 let mut starting_value = 1;
398 if version >= 2_000 {
399 reader.read_exact(&mut u8_buf)?;
400 remove_on_zero = u8::from_be_bytes(u8_buf) != 0;
401 }
402 if version >= 2_002 {
403 aging_strategy = LruAgingStrategy::read_from(&mut reader, version)?;
404 reader.read_exact(&mut u32_buf)?;
405 starting_value = u32::from_be_bytes(u32_buf);
406 reader.read_exact(&mut u32_buf)?;
407 increment = u32::from_be_bytes(u32_buf);
408 }
409 Ok(Self {
410 counter_bits: (bits as usize).try_into()?,
411 aging_strategy,
412 starting_value,
413 remove_on_zero,
414 increment,
415 })
416 }
417}
418
419impl ExportableConfig for LruAgingStrategy {
420 fn write_to(&self, mut writer: impl Write) -> Result<(), ExportError> {
421 match self {
422 LruAgingStrategy::Halving => writer.write_all(&1u8.to_be_bytes())?,
423 LruAgingStrategy::Decrement(value) => {
424 writer.write_all(&2u8.to_be_bytes())?;
425 writer.write_all(&value.to_be_bytes())?;
426 }
427 }
428 Ok(())
429 }
430
431 fn read_from(mut reader: impl Read, _version: usize) -> Result<Self, ImportError> {
432 let mut u8_buf = [0u8; 1];
433 let mut u32_buf = [0u8; 4];
434 reader.read_exact(&mut u8_buf)?;
435 let strategy = u8::from_be_bytes(u8_buf);
436 match strategy {
437 1 => Ok(Self::Halving),
438 2 => {
439 reader.read_exact(&mut u32_buf)?;
440 let value = u32::from_be_bytes(u32_buf);
441 Ok(Self::Decrement(value))
442 }
443 _ => Ok(Self::default()),
445 }
446 }
447}
448
449impl ExportableConfig for TtlConfig {
450 fn write_to(&self, mut writer: impl Write) -> Result<(), ExportError> {
451 writer.write_all(&self.ttl.get().to_be_bytes())?;
452 writer.write_all(&u8::from(self.ttl_bits).to_be_bytes())?;
453 Ok(())
454 }
455
456 fn read_from(mut reader: impl Read, _version: usize) -> Result<Self, ImportError> {
457 let mut u8_buf = [0u8; 1];
458 let mut u32_buf = [0u8; 4];
459 reader.read_exact(&mut u32_buf)?;
460 let ttl = u32::from_be_bytes(u32_buf);
461 reader.read_exact(&mut u8_buf)?;
462 let bits = u8::from_be_bytes(u8_buf);
463 Ok(Self {
464 ttl: ttl.try_into()?,
465 ttl_bits: (bits as usize).try_into()?,
466 })
467 }
468}
469
470impl ExportableConfig for CounterConfig {
471 fn write_to(&self, mut writer: impl Write) -> Result<(), ExportError> {
472 writer.write_all(&u8::from(self.counter_bits).to_be_bytes())?;
473 writer.write_all(&self.change_on_insert.to_be_bytes())?;
474 writer.write_all(&self.change_on_lookup.to_be_bytes())?;
475 Ok(())
476 }
477
478 fn read_from(mut reader: impl Read, _version: usize) -> Result<Self, ImportError> {
479 let mut u8_buf = [0u8; 1];
480 let mut u32_buf = [0u8; 4];
481 reader.read_exact(&mut u8_buf)?;
482 let bits = u8::from_be_bytes(u8_buf);
483 reader.read_exact(&mut u32_buf)?;
484 let change_on_insert = i32::from_be_bytes(u32_buf);
485 reader.read_exact(&mut u32_buf)?;
486 let change_on_lookup = i32::from_be_bytes(u32_buf);
487 Ok(Self {
488 counter_bits: (bits as usize).try_into()?,
489 change_on_insert,
490 change_on_lookup,
491 })
492 }
493}