diskann_quantization/
error.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6//! Formatting utilities for error chains.
7
8use std::{cell::UnsafeCell, marker::PhantomData, mem::MaybeUninit};
9
10#[derive(Debug)]
11pub enum Infallible {}
12
13impl std::fmt::Display for Infallible {
14    fn fmt(&self, _: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
15        unreachable!("Infallible is unconstructible");
16    }
17}
18
19impl std::error::Error for Infallible {}
20
21/// A utility for printing the whole [source tree](https://doc.rust-lang.org/std/error/trait.Error.html#method.source).
22/// of an error type.
23///
24/// ```rust
25/// use diskann_quantization::error::Format;
26/// use std::{error::Error, fmt::{Display, Formatter, Result}};
27///
28/// #[derive(Debug)]
29/// struct A;
30///
31/// impl Display for A {
32///     fn fmt(&self, f: &mut Formatter<'_>) -> Result {
33///         f.write_str("A")
34///     }
35/// }
36///
37/// impl Error for A {
38///     fn source(&self) -> Option<&(dyn Error + 'static)> {
39///         Some(&B)
40///     }
41/// }
42///
43/// #[derive(Debug)]
44/// struct B;
45///
46/// impl Display for B {
47///     fn fmt(&self, f: &mut Formatter<'_>) -> Result {
48///         f.write_str("B")
49///     }
50/// }
51///
52/// impl Error for B {}
53///
54/// assert_eq!(Format(A).to_string(), "A\n    caused by: B");
55/// assert_eq!(Format(&A).to_string(), "A\n    caused by: B");
56/// ```
57#[derive(Debug)]
58pub struct Format<T>(pub T);
59
60impl<T> std::fmt::Display for Format<T>
61where
62    T: std::error::Error,
63{
64    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
65        // Protect against pathological `Source` implementations that just return
66        // themselves.
67        const LIMIT: usize = 256;
68
69        // Cast wrap the walking of the source chain into something that behaves like an
70        // iterator.
71        struct Source<'a>(Option<&'a (dyn std::error::Error + 'static)>);
72        impl<'a> Iterator for Source<'a> {
73            type Item = &'a (dyn std::error::Error + 'static);
74            fn next(&mut self) -> Option<Self::Item> {
75                let current = self.0;
76                self.0 = match current {
77                    Some(current) => current.source(),
78                    None => None,
79                };
80                current
81            }
82        }
83
84        write!(f, "{}", self.0)?;
85        let mut itr = Source(self.0.source());
86        for source in itr.by_ref().take(LIMIT) {
87            write!(f, "\n    caused by: {}", source)?;
88        }
89
90        if itr.next().is_some() {
91            write!(f, "\n    ... (limit reached)")?;
92        }
93
94        Ok(())
95    }
96}
97
98/// Format the entire error chain for `err` by first calling `err.to_string()` and then
99/// by walking the error's
100/// [source tree](https://doc.rust-lang.org/std/error/trait.Error.html#method.source).
101pub fn format<E>(err: &E) -> String
102where
103    E: std::error::Error + ?Sized,
104{
105    Format(err).to_string()
106}
107
108/// An implementation of `Box<dyn std::error::Error>` that stores the error payload inline,
109/// avoiding dynamic memory allocation. This has several practical drawbacks:
110///
111/// 1. The size of the error payload must be at most `N` bytes.
112/// 2. The alignment of the error payload must be at most 8 bytes.
113///
114/// Both of these contraints are verified using post-monomorphization errors.
115///
116/// # Example
117///
118/// ```
119/// use diskann_quantization::error::InlineError;
120///
121/// let base_error = u32::try_from(u64::MAX).unwrap_err();
122/// let mut error = InlineError::<8>::new(base_error);
123/// assert_eq!(error.to_string(), base_error.to_string());
124///
125/// // Change the dynamic type of the contained error.
126/// error = InlineError::new(Box::new(base_error));
127/// assert_eq!(error.to_string(), base_error.to_string());
128/// ```
129#[repr(C)]
130pub struct InlineError<const N: usize = 16> {
131    // We place the vtable first to enable the niche-optimization.
132    vtable: &'static ErrorVTable,
133
134    // NOTE: We need to use `MaybeUninit` instead of `u8` to maintain the provenance of
135    // any pointers/references stored in the payload.
136    //
137    // Additionally, the `UnsafeCell` is needed because the payload may have interior
138    // mutability as a side-effect of API calls.
139    object: UnsafeCell<[MaybeUninit<u8>; N]>,
140}
141
142// SAFETY: We only allow error payloads that are `Send`.
143unsafe impl<const N: usize> Send for InlineError<N> {}
144
145// SAFETY: We only allow error payloads that are `Sync`.
146unsafe impl<const N: usize> Sync for InlineError<N> {}
147
148impl<const N: usize> InlineError<N> {
149    /// Construct a new `InlineError` around `error`.
150    ///
151    /// Fails to compile if:
152    ///
153    /// 1. `std::mem::align_of::<T>() > 8`: Objects of type `T` must be compatible with the
154    ///    inline storage buffer.
155    /// 2. `std::mem::size_of::<T>() > N`: Objects of type `T` must fit within a buffer of
156    ///    size `N`.
157    pub fn new<T>(error: T) -> Self
158    where
159        T: std::error::Error + Send + Sync + 'static,
160    {
161        const { assert!(std::mem::size_of::<T>() <= N, "error type is too big") };
162        const {
163            assert!(
164                std::mem::align_of::<T>() <= std::mem::align_of::<&'static ErrorVTable>(),
165                "error type has alignment stricter than 8"
166            )
167        };
168
169        let mut this = Self {
170            vtable: &ErrorVTable {
171                debug: error_debug::<T>,
172                display: error_display::<T>,
173                source: error_source::<T>,
174                drop: error_drop::<T>,
175            },
176            object: UnsafeCell::new([MaybeUninit::uninit(); N]),
177        };
178
179        // SAFETY: We have const assertions that the size and alignment of `T` are
180        // compatible with the buffer we created.
181        //
182        // Additionally, the memory we are writing to does not have a valid object stored,
183        // so using `ptr::write` will not leak memory.
184        unsafe { this.object.get_mut().as_mut_ptr().cast::<T>().write(error) };
185
186        this
187    }
188
189    // Return the base pointer of the inline storage in a type that propagates the lifetime
190    // of `self`. This allows the `.source()` implementation to propagate the correct
191    // lifetime.
192    fn ptr_ref(&self) -> Ref<'_> {
193        Ref {
194            ptr: self.object.get().cast::<MaybeUninit<u8>>(),
195            _lifetime: PhantomData,
196        }
197    }
198}
199
200impl<const N: usize> Drop for InlineError<N> {
201    fn drop(&mut self) {
202        // SAFETY: The constructor invariants of `InlineError` ensure that the vtable method
203        // is safe to call.
204        //
205        // Since the only place where the `drop` function is called is in the implementation
206        // of `Drop` for `InlineError`, we are guaranteed that the underlying object is
207        // valid.
208        unsafe { (self.vtable.drop)(self.object.get().cast::<MaybeUninit<u8>>()) }
209    }
210}
211
212impl<const N: usize> std::fmt::Display for InlineError<N> {
213    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
214        // SAFETY: The constructor invariants of `InlineError` ensure that the vtable method
215        // is safe to call.
216        unsafe { (self.vtable.display)(self.object.get().cast::<MaybeUninit<u8>>(), f) }
217    }
218}
219
220impl<const N: usize> std::fmt::Debug for InlineError<N> {
221    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
222        write!(f, "InlineError<{}> {{ object: ", N)?;
223        // SAFETY: The constructor invariants of `InlineError` ensure that the vtable method
224        // is safe to call.
225        unsafe { (self.vtable.debug)(self.object.get().cast::<MaybeUninit<u8>>(), f) }?;
226        write!(f, ", vtable: {:?} }}", self.vtable)
227    }
228}
229
230impl<const N: usize> std::error::Error for InlineError<N> {
231    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
232        // SAFETY: The constructor invariants of `InlineError` ensure that the vtable method
233        // is safe to call.
234        unsafe { (self.vtable.source)(self.ptr_ref()) }
235    }
236}
237
238#[derive(Debug)]
239struct ErrorVTable {
240    debug: unsafe fn(*const MaybeUninit<u8>, &mut std::fmt::Formatter<'_>) -> std::fmt::Result,
241    display: unsafe fn(*const MaybeUninit<u8>, &mut std::fmt::Formatter<'_>) -> std::fmt::Result,
242    source: unsafe fn(Ref<'_>) -> Option<&(dyn std::error::Error + 'static)>,
243    drop: unsafe fn(*mut MaybeUninit<u8>),
244}
245
246// SAFETY: `object` must point to a valid object of type `T`.
247unsafe fn error_debug<T>(
248    object: *const MaybeUninit<u8>,
249    f: &mut std::fmt::Formatter<'_>,
250) -> std::fmt::Result
251where
252    T: std::fmt::Debug,
253{
254    // SAFETY: Required of caller.
255    unsafe { &*object.cast::<T>() }.fmt(f)
256}
257
258// SAFETY: `object` must point to a valid object of type `T`.
259unsafe fn error_display<T>(
260    object: *const MaybeUninit<u8>,
261    f: &mut std::fmt::Formatter<'_>,
262) -> std::fmt::Result
263where
264    T: std::fmt::Display,
265{
266    // SAFETY: Required of caller.
267    unsafe { &*object.cast::<T>() }.fmt(f)
268}
269
270// SAFETY: A valid instance of type `T` must be stored in `object` beginning at the start
271// of the slice. Note that this implies that `std::mem::size_of::<T>() <= object.len()` and
272// that the start of the slice is properly aligned.
273unsafe fn error_source<T>(object: Ref<'_>) -> Option<&(dyn std::error::Error + 'static)>
274where
275    T: std::error::Error + 'static,
276{
277    // SAFETY: Required of caller.
278    unsafe { &*object.ptr.cast::<T>() }.source()
279}
280
281// A pointer with a tagged lifetime.
282struct Ref<'a> {
283    ptr: *const MaybeUninit<u8>,
284    _lifetime: PhantomData<&'a MaybeUninit<u8>>,
285}
286
287// SAFETY: `object` must point to a valid object of type `T`. As a side effect, the
288// pointed-to object will be dropped.
289unsafe fn error_drop<T>(object: *mut MaybeUninit<u8>) {
290    // SAFETY: Required of caller.
291    unsafe { std::ptr::drop_in_place::<T>(object.cast::<T>()) }
292}
293
294///////////
295// Tests //
296///////////
297
298#[cfg(test)]
299mod tests {
300    use std::sync::{
301        Arc, Mutex,
302        atomic::{AtomicUsize, Ordering},
303    };
304
305    use thiserror::Error;
306
307    use super::*;
308
309    #[derive(Error, Debug, Clone)]
310    #[error("error A")]
311    struct ErrorA;
312
313    #[derive(Error, Debug, Clone)]
314    #[error("error B with val {val}")]
315    struct ErrorB<Inner: std::error::Error> {
316        val: usize,
317        #[source]
318        source: Inner,
319    }
320
321    #[derive(Error, Debug)]
322    #[error("error C with message {message}")]
323    struct ErrorC<Inner: std::error::Error> {
324        message: String,
325        /// `thiserror` automatically marks this as the error source.
326        source: Inner,
327    }
328
329    #[test]
330    fn test_formatting() {
331        // No Nesting
332        let message = format(&ErrorA);
333        assert_eq!(message, "error A");
334
335        assert_eq!(Format(ErrorA).to_string(), "error A");
336        assert_eq!(Format(&ErrorA).to_string(), "error A");
337
338        // One Level of Nesting
339        let error = ErrorB {
340            val: 10,
341            source: ErrorA,
342        };
343
344        let expected = "error B with val 10\n    caused by: error A";
345        assert_eq!(format(&error), expected);
346        assert_eq!(Format(&error).to_string(), expected);
347        assert_eq!(Format(error.clone()).to_string(), expected);
348
349        // Multiple Levels of Nesting
350        let error = ErrorC {
351            message: "Hello World".to_string(),
352            source: error,
353        };
354        let expected = "error C with message Hello World\n    \
355                        caused by: error B with val 10\n    \
356                        caused by: error A";
357
358        assert_eq!(format(&error), expected);
359        assert_eq!(Format(&error).to_string(), expected);
360        assert_eq!(Format(error).to_string(), expected);
361    }
362
363    // A pathological error type that returns itself for `source` and thus ends up
364    // with an unlimited chain.
365    #[derive(Debug)]
366    struct Infinite;
367
368    impl std::fmt::Display for Infinite {
369        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
370            f.write_str("an unending source")
371        }
372    }
373
374    impl std::error::Error for Infinite {
375        fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
376            Some(self)
377        }
378    }
379
380    #[test]
381    fn test_infinite_detected() {
382        // Without a limit - this line will never finish (we'll either hit a stack overflow
383        // or run out of memory for the string).
384        let s = Format(Infinite).to_string();
385        assert!(s.contains("(limit reached)"));
386    }
387
388    ///////////
389    // Error //
390    ///////////
391
392    #[derive(Debug, Error)]
393    #[error("zero sized error")]
394    struct ZeroSizedError;
395
396    #[derive(Debug, Error)]
397    #[error("error with drop: {}", self.0.load(Ordering::Relaxed))]
398    struct ErrorWithDrop(Arc<AtomicUsize>);
399
400    impl Drop for ErrorWithDrop {
401        fn drop(&mut self) {
402            self.0.fetch_add(1, Ordering::Relaxed);
403        }
404    }
405
406    #[derive(Debug, Error)]
407    #[error("error with source")]
408    struct ErrorWithSource(#[from] ZeroSizedError);
409
410    // This tests (using Miri) that it's safe to contain error types with interior mutability.
411    struct ErrorWithInteriorMutability(Mutex<usize>);
412
413    impl std::fmt::Debug for ErrorWithInteriorMutability {
414        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
415            let current = {
416                let mut guard = self.0.lock().unwrap();
417                let current = *guard;
418                *guard += 1;
419                current
420            };
421
422            write!(f, "{}", current)
423        }
424    }
425
426    impl std::fmt::Display for ErrorWithInteriorMutability {
427        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
428            let current = {
429                let mut guard = self.0.lock().unwrap();
430                let current = *guard;
431                *guard += 1;
432                current
433            };
434
435            write!(f, "{}", current)
436        }
437    }
438
439    impl std::error::Error for ErrorWithInteriorMutability {
440        fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
441            *self.0.lock().unwrap() += 1;
442            None
443        }
444    }
445
446    #[test]
447    fn sizes_and_offsets() {
448        let ref_size = std::mem::size_of::<&'static ()>();
449        let ref_align = std::mem::align_of::<&'static ()>();
450
451        assert_eq!(std::mem::offset_of!(InlineError<0>, object), ref_size);
452        assert_eq!(std::mem::offset_of!(InlineError<8>, object), ref_size);
453        assert_eq!(std::mem::offset_of!(InlineError<16>, object), ref_size);
454
455        assert_eq!(std::mem::size_of::<InlineError<0>>(), ref_size);
456        assert_eq!(std::mem::size_of::<Option<InlineError<0>>>(), ref_size);
457        assert_eq!(std::mem::align_of::<InlineError<0>>(), ref_align);
458        assert_eq!(std::mem::align_of::<Option<InlineError<0>>>(), ref_align);
459
460        assert_eq!(std::mem::size_of::<InlineError<8>>(), ref_size + 8);
461        assert_eq!(std::mem::size_of::<Option<InlineError<8>>>(), ref_size + 8);
462        assert_eq!(std::mem::align_of::<InlineError<8>>(), ref_align);
463        assert_eq!(std::mem::align_of::<Option<InlineError<8>>>(), ref_align);
464
465        assert_eq!(std::mem::size_of::<InlineError<16>>(), ref_size + 16);
466        assert_eq!(
467            std::mem::size_of::<Option<InlineError<16>>>(),
468            ref_size + 16
469        );
470        assert_eq!(std::mem::align_of::<InlineError<16>>(), ref_align);
471        assert_eq!(std::mem::align_of::<Option<InlineError<16>>>(), ref_align);
472    }
473
474    #[test]
475    fn inline_error_zst() {
476        use std::error::Error;
477
478        let error = InlineError::<0>::new(ZeroSizedError);
479        assert_eq!(
480            std::mem::size_of_val(&error),
481            8,
482            "expected 8 bytes for the payload and 0-bytes for the vtable"
483        );
484        assert_eq!(error.to_string(), "zero sized error");
485
486        let debug = format!("{:?}", error);
487        assert!(
488            debug.starts_with(&format!("InlineError<0> {{ object: {:?}", ZeroSizedError)),
489            "debug message: {}",
490            debug
491        );
492
493        assert!(error.source().is_none());
494
495        // Move it into a box. This is mainly a Miri tests.
496        let _ = Box::new(error);
497    }
498
499    #[test]
500    fn inline_error_with_drop() {
501        use std::error::Error;
502
503        let count = Arc::new(AtomicUsize::new(10));
504        let mut error = InlineError::<8>::new(ErrorWithDrop(count.clone()));
505        assert_eq!(
506            std::mem::size_of_val(&error),
507            16,
508            "expected 8 bytes for the payload and 8-bytes for the vtable"
509        );
510        assert_eq!(error.to_string(), "error with drop: 10");
511        assert!(error.source().is_none());
512
513        // Move it into a box. This is mainly a Miri tests.
514        error = InlineError::new(ZeroSizedError);
515        assert_eq!(error.to_string(), "zero sized error");
516
517        assert_eq!(count.load(Ordering::Relaxed), 11, "failed to run \"drop\"");
518    }
519
520    #[test]
521    fn inline_error_with_interior_mutability() {
522        use std::error::Error;
523
524        // Use 64 bytes to accommodate larger mutex sizes on macOS
525        let error = InlineError::<64>::new(ErrorWithInteriorMutability(Mutex::new(0)));
526        assert_eq!(
527            std::mem::size_of_val(&error),
528            72,
529            "expected 64 bytes for the payload and 8-bytes for the vtable"
530        );
531        assert_eq!(error.to_string(), "0");
532        let debug = format!("{:?}", error);
533        assert!(debug.contains("object: 1"), "got {}", debug);
534        assert_eq!(error.to_string(), "2");
535
536        let debug = format!("{:?}", error);
537        assert!(debug.contains("object: 3"), "got {}", debug);
538
539        assert!(error.source().is_none());
540        assert_eq!(error.to_string(), "5");
541    }
542
543    #[test]
544    fn inline_error_with_source() {
545        use std::error::Error;
546
547        let error = InlineError::<8>::new(ErrorWithSource(ZeroSizedError));
548        assert_eq!(
549            std::mem::size_of_val(&error),
550            16,
551            "expected 8 bytes for the payload and 8-bytes for the vtable"
552        );
553        assert_eq!(error.to_string(), "error with source");
554        assert_eq!(error.source().unwrap().to_string(), "zero sized error");
555
556        // Move it into a box. This is mainly a Miri tests.
557        let _ = Box::new(error);
558    }
559}
diskann_quantization/error.rs

diskann_quantization/
error.rs
