diskann_quantization/alloc/
aligned_slice.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use super::Poly;
7use crate::alloc::{AlignedAllocator, AllocatorError};
8
9use crate::num::PowerOfTwo;
10
11/// Type alias for an aligned, heap-allocated slice
12///
13/// Shorthand for `Poly<[T], AlignedAllocator>` which is intended
14/// for allocations requiring specific alignment (e.g. cache-line or disk-sector alignment)
15pub type AlignedSlice<T> = Poly<[T], AlignedAllocator>;
16
17/// Create a new [`AlignedSlice`] with the given capacity and alignment
18/// initialized to `T::default()`
19pub fn aligned_slice<T: Default>(
20    capacity: usize,
21    alignment: PowerOfTwo,
22) -> Result<AlignedSlice<T>, AllocatorError> {
23    let allocator = AlignedAllocator::new(alignment);
24    Poly::from_iter((0..capacity).map(|_| T::default()), allocator)
25}
26
27#[cfg(test)]
28mod tests {
29    use super::*;
30
31    #[test]
32    fn aligned_slice_alignment_32() {
33        let data = aligned_slice::<f32>(1000, PowerOfTwo::new(32).unwrap()).unwrap();
34        assert_eq!(data.len(), 1000);
35        assert_eq!(data.as_ptr() as usize % 32, 0);
36    }
37
38    #[test]
39    fn aligned_slice_alignment_256() {
40        let data = aligned_slice::<u8>(500, PowerOfTwo::new(256).unwrap()).unwrap();
41        assert_eq!(data.len(), 500);
42        assert_eq!(data.as_ptr() as usize % 256, 0);
43    }
44
45    #[test]
46    fn aligned_slice_alignment_512() {
47        let data = aligned_slice::<u8>(4096, PowerOfTwo::new(512).unwrap()).unwrap();
48        assert_eq!(data.len(), 4096);
49        assert_eq!(data.as_ptr() as usize % 512, 0);
50    }
51
52    #[test]
53    fn aligned_slice_zero_length() {
54        // Zero-length aligned slices should succeed: `Poly::from_iter`
55        // special-cases empty iterators and returns an empty slice.
56        let data = aligned_slice::<f32>(0, PowerOfTwo::new(16).unwrap()).unwrap();
57        assert!(data.is_empty());
58        assert_eq!(data.len(), 0);
59    }
60
61    #[test]
62    fn aligned_slice_default_initialized() {
63        let data = aligned_slice::<f32>(100, PowerOfTwo::new(64).unwrap()).unwrap();
64        for &val in data.iter() {
65            assert_eq!(val, 0.0);
66        }
67    }
68
69    #[test]
70    fn aligned_slice_deref_mut() {
71        let mut data = aligned_slice::<f32>(4, PowerOfTwo::new(32).unwrap()).unwrap();
72        data[0] = 1.0;
73        data[1] = 2.0;
74        assert_eq!(&data[..2], &[1.0, 2.0]);
75    }
76}
diskann_quantization/alloc/aligned_slice.rs

diskann_quantization/alloc/
aligned_slice.rs
