diskann_disk/search/provider/aligned_file_reader/
aligned_read.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use std::marker::PhantomData;
7
8use diskann::{ANNError, ANNResult};
9use diskann_quantization::num::PowerOfTwo;
10
11/// Type-level memory-alignment witness for [`AlignedRead`]. Each implementor is
12/// a unit type carrying a single `PowerOfTwo` value.
13///
14/// Custom readers can define their own marker (e.g. `A4096`) by adding a unit
15/// type and an `impl Alignment` with the desired `VALUE`.
16pub trait Alignment {
17    /// The alignment, in bytes.
18    const VALUE: PowerOfTwo;
19}
20
21macro_rules! alignment_marker {
22    ($name:ident, $value:expr) => {
23        #[doc = concat!("Alignment witness for ", stringify!($value), " bytes.")]
24        #[derive(Debug, Clone, Copy)]
25        pub struct $name;
26        impl Alignment for $name {
27            const VALUE: PowerOfTwo = $value;
28        }
29    };
30}
31
32alignment_marker!(A1, PowerOfTwo::V1);
33alignment_marker!(A512, PowerOfTwo::V512);
34
35/// Disk-IO read request, parameterized by its required memory alignment `A`.
36///
37/// Three constraints govern a read:
38/// 1. Disk offset alignment.
39/// 2. Buffer length alignment.
40/// 3. Buffer pointer alignment in memory.
41///
42/// All three are checked against `A::VALUE` at construction time by
43/// [`AlignedRead::new`]. A typed `AlignedRead<T, A>` is therefore a witness
44/// that the request satisfies its declared alignment, and the file reader's
45/// `read` method can rely on it without re-checking.
46#[derive(Debug)]
47pub struct AlignedRead<'a, T, A: Alignment = A1> {
48    offset: u64,
49    aligned_buf: &'a mut [T],
50    _alignment: PhantomData<A>,
51}
52
53impl<'a, T, A: Alignment> AlignedRead<'a, T, A> {
54    /// Build an `AlignedRead` after validating that `offset`, the buffer
55    /// length (in bytes), and the buffer pointer all satisfy `A::VALUE`.
56    pub fn new(offset: u64, aligned_buf: &'a mut [T]) -> ANNResult<Self> {
57        Self::assert_is_aligned(aligned_buf.as_ptr() as usize, "buffer pointer")?;
58        Self::assert_is_aligned(std::mem::size_of_val(aligned_buf), "buffer length")?;
59        Self::assert_is_aligned(offset as usize, "offset")?;
60        Ok(Self {
61            offset,
62            aligned_buf,
63            _alignment: PhantomData,
64        })
65    }
66
67    fn assert_is_aligned(val: usize, kind: &str) -> ANNResult<()> {
68        let align = A::VALUE.raw();
69        if val.is_multiple_of(align) {
70            Ok(())
71        } else {
72            Err(ANNError::log_disk_io_request_alignment_error(format!(
73                "{kind} {val} not aligned to {align}",
74            )))
75        }
76    }
77
78    pub fn offset(&self) -> u64 {
79        self.offset
80    }
81
82    pub fn aligned_buf(&self) -> &[T] {
83        self.aligned_buf
84    }
85
86    pub fn aligned_buf_mut(&mut self) -> &mut [T] {
87        self.aligned_buf
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94    use diskann::ANNErrorKind;
95    use diskann_quantization::alloc::{AlignedAllocator, Poly};
96
97    fn aligned_512(len: usize) -> Poly<[u8], AlignedAllocator> {
98        Poly::broadcast(0u8, len, AlignedAllocator::A512).unwrap()
99    }
100
101    #[test]
102    fn aligned_read_carries_offset_and_buffer() {
103        let mut buffer = vec![0u8; 512];
104        let read = AlignedRead::<u8, A1>::new(512, &mut buffer).unwrap();
105        assert_eq!(read.offset(), 512);
106        assert_eq!(read.aligned_buf().len(), 512);
107    }
108
109    #[test]
110    fn aligned_read_buffer_access() {
111        let mut buffer = vec![42u8; 512];
112        let mut read = AlignedRead::<u8, A1>::new(0, &mut buffer).unwrap();
113        assert_eq!(read.aligned_buf()[0], 42);
114        read.aligned_buf_mut()[0] = 100;
115        assert_eq!(read.aligned_buf()[0], 100);
116    }
117
118    #[test]
119    fn a512_accepts_fully_aligned_request() {
120        let mut buf = aligned_512(512);
121        AlignedRead::<u8, A512>::new(0, &mut buf).expect("aligned request should pass");
122    }
123
124    #[test]
125    fn a1_default_accepts_anything() {
126        let mut buffer = vec![0u8; 100];
127        AlignedRead::<u8, A1>::new(1, &mut buffer).expect("A1 alignment should accept any request");
128    }
129
130    #[test]
131    fn rejects_unaligned_buffer_pointer() {
132        let mut buf = aligned_512(1024);
133        let slice = &mut buf[1..513]; // ptr offset by 1; length 512 ✓; offset 0 ✓
134        let err = AlignedRead::<u8, A512>::new(0, slice)
135            .expect_err("misaligned buffer pointer should be rejected");
136        assert_eq!(err.kind(), ANNErrorKind::DiskIOAlignmentError);
137    }
138
139    #[test]
140    fn rejects_unaligned_buffer_length() {
141        let mut buf = aligned_512(1024);
142        let slice = &mut buf[..100]; // ptr ✓; length 100 ✗; offset 0 ✓
143        let err = AlignedRead::<u8, A512>::new(0, slice)
144            .expect_err("buffer length 100 (not a multiple of 512) should be rejected");
145        assert_eq!(err.kind(), ANNErrorKind::DiskIOAlignmentError);
146    }
147
148    #[test]
149    fn rejects_unaligned_offset() {
150        let mut buf = aligned_512(1024);
151        let slice = &mut buf[..512]; // ptr ✓; length 512 ✓; offset 1 ✗
152        let err = AlignedRead::<u8, A512>::new(1, slice)
153            .expect_err("offset 1 (not a multiple of 512) should be rejected");
154        assert_eq!(err.kind(), ANNErrorKind::DiskIOAlignmentError);
155    }
156}
diskann_disk/search/provider/aligned_file_reader/aligned_read.rs

diskann_disk/search/provider/aligned_file_reader/
aligned_read.rs
