diskann_quantization/alloc/
mod.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use std::{alloc::Layout, ptr::NonNull};
7
8mod aligned;
9mod aligned_slice;
10mod bump;
11mod poly;
12mod traits;
13
14pub use aligned::{AlignedAllocator, NotPowerOfTwo};
15pub use aligned_slice::{AlignedSlice, aligned_slice};
16pub use bump::BumpAllocator;
17pub use poly::{CompoundError, Poly, TrustedIter, poly};
18pub use traits::{Allocator, AllocatorCore, AllocatorError};
19
20/// A handle to Rust's global allocator. This type does not support allocations of size 0.
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub struct GlobalAllocator;
23
24// SAFETY: This is a simple wrapper around Rust's built-in allocation and deallocation
25// methods, augmented slightly to handle zero sized layouts by returning a dangling pointer.
26//
27// The returned slice from `allocate` always has the exact size and alignment as `layout`.
28unsafe impl AllocatorCore for GlobalAllocator {
29    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocatorError> {
30        if layout.size() == 0 {
31            return Err(AllocatorError);
32        }
33
34        // SAFETY: `layout` has a non-zero size.
35        let ptr = unsafe { std::alloc::alloc(layout) };
36        let ptr = std::ptr::slice_from_raw_parts_mut(ptr, layout.size());
37        NonNull::new(ptr).ok_or(AllocatorError)
38    }
39
40    unsafe fn deallocate(&self, ptr: NonNull<[u8]>, layout: Layout) {
41        // SAFETY: The caller has the responsibility to ensure that `ptr` and `layout`
42        // came from a previous allocation.
43        unsafe { std::alloc::dealloc(ptr.as_ptr().cast::<u8>(), layout) }
44    }
45}
46
47////////////
48// Scoped //
49////////////
50
51trait DebugAllocator: AllocatorCore + std::fmt::Debug {}
52impl<T> DebugAllocator for T where T: AllocatorCore + std::fmt::Debug {}
53
54/// A dynamic wrapper around an `AllocatorCore` that provides the guarantee that all
55/// allocated object are tied to a given scope.
56///
57/// Additionally, this can allow the use of an allocator that is not `Clone` in contexts
58/// where a clonable allocator is needed (provided the scoping limitations are acceptable).
59#[derive(Debug, Clone, Copy)]
60pub struct ScopedAllocator<'a> {
61    allocator: &'a dyn DebugAllocator,
62}
63
64impl<'a> ScopedAllocator<'a> {
65    /// Construct a new `ScopedAllocator` around the provided `allocator`.
66    pub const fn new<T>(allocator: &'a T) -> Self
67    where
68        T: AllocatorCore + std::fmt::Debug,
69    {
70        Self { allocator }
71    }
72}
73
74impl ScopedAllocator<'static> {
75    /// A convenience method for construcing a `ScopedAllocator` around the [`GlobalAllocator`]
76    /// for cases where a more specialized allocator is not needed.
77    pub const fn global() -> Self {
78        Self {
79            allocator: &GlobalAllocator,
80        }
81    }
82}
83
84// SAFETY: This allocator simply delegates to the underlying allocator.
85unsafe impl AllocatorCore for ScopedAllocator<'_> {
86    fn allocate(&self, layout: std::alloc::Layout) -> Result<NonNull<[u8]>, AllocatorError> {
87        self.allocator.allocate(layout)
88    }
89
90    unsafe fn deallocate(&self, ptr: NonNull<[u8]>, layout: std::alloc::Layout) {
91        // SAFETY: Inherited from caller.
92        unsafe { self.allocator.deallocate(ptr, layout) }
93    }
94}
95
96///////////////
97// Try Clone //
98///////////////
99
100/// A trait like [`Clone`] that allows graceful allocation failure.
101///
102/// # NOTE
103///
104/// Keep this `pub(crate)` for now because we do not want general users of the crate
105/// relying on the current implementations for [`Poly`]. In particular, the base case should
106/// be `Poly<T> where T: TryClone` instead of `Poly<T> where T: Clone`.
107pub(crate) trait TryClone: Sized {
108    /// Returns a duplicate of the value.
109    fn try_clone(&self) -> Result<Self, AllocatorError>;
110}
111
112///////////
113// Tests //
114///////////
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119
120    fn test_alloc<T>() {
121        let alloc = GlobalAllocator;
122
123        let layout = Layout::new::<T>();
124        let ptr = alloc.allocate(layout).unwrap();
125
126        assert_eq!(ptr.len(), layout.size());
127        assert_eq!(ptr.len(), std::mem::size_of::<T>());
128        assert_eq!((ptr.as_ptr().cast::<u8>() as usize) % layout.align(), 0);
129
130        // SAFETY: `ptr` was obtained from this allocator with the specified `layout`.
131        unsafe { alloc.deallocate(ptr, layout) };
132    }
133
134    #[test]
135    fn test_global_allocator() {
136        assert!(GlobalAllocator.allocate(Layout::new::<()>()).is_err());
137
138        test_alloc::<(u8,)>();
139        test_alloc::<(u8, u8)>();
140        test_alloc::<(u8, u8, u8)>();
141        test_alloc::<(u8, u8, u8, u8)>();
142        test_alloc::<(u8, u8, u8, u8, u8)>();
143        test_alloc::<(u8, u8, u8, u8, u8, u8)>();
144        test_alloc::<(u8, u8, u8, u8, u8, u8, u8)>();
145        test_alloc::<(u8, u8, u8, u8, u8, u8, u8, u8)>();
146        test_alloc::<(u8, u8, u8, u8, u8, u8, u8, u8, u8)>();
147
148        test_alloc::<(u16,)>();
149        test_alloc::<(u16, u16)>();
150        test_alloc::<(u16, u16, u16)>();
151        test_alloc::<(u16, u16, u16, u16)>();
152        test_alloc::<(u16, u16, u16, u16, u16)>();
153
154        test_alloc::<(u32,)>();
155        test_alloc::<(u32, u32)>();
156        test_alloc::<(u32, u32, u32)>();
157
158        test_alloc::<String>();
159    }
160}
diskann_quantization/alloc/mod.rs

diskann_quantization/alloc/
mod.rs
