1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
use std::ops::{Deref, RangeTo};
use aligned_vec::{AVec, RuntimeAlign};
use crate::common::zeros::ZEROS;
/// Like [`Vec<u8>`], but `&buf[i]` is aligned as if it lived at
/// `file_offset + i` in a file.
///
/// # Motivational example
///
/// Suppose a file that contains fields with mixed types and alignments. We want
/// to read 12-byte entry starting from `file_offset = 0xABC3` into a buffer.
///
/// ```text
/// file_offset=0xABC3 file_offset=0xABC8
/// ↓ ↓
/// ───┬───┬───┬───┬───┬───┬───┬───┬───┬───┬───┬───┬───┬───┬───┬───┬──
/// │ │ 0 │ 1 │ 2 │ 3 │ 4 │ 5 │ 6 │ 7 │ 8 │ 9 │10 │11 │12 │ │
/// ───┴───┼───┴───┼───┴───┴───┼───┴───┴───┴───┴───┴───┴───┴───┼───┴──
/// │ val_1 │ padding │ val_2 │
/// │[u8; 2]│ [u8; 3] │ u64 │
///
/// ```
///
/// For `val_2: u64`, its file offset is aligned by 8 bytes, but its position in
/// the buffer is not. So, to make it work, `&buf[5]` should be aligned.
/// So, the invariant (for some `align` that is large enough):
///
/// `ptr_of(&buf[buf_pos]) % align == (file_offset + buf_pos) % align`
///
/// Files read through mmaps satisfy this invariant: mmaps' natural assignment is
/// 4096 bytes (usually), which is more than enough. This structure mimics that
/// fraction of mmap's power.
pub struct AlignedBuf {
buf: AVec<u8, RuntimeAlign>,
offset: usize,
}
impl AlignedBuf {
/// Creates [`AlignedBuf`]. It will hold the alignment invariant for `align`,
/// assuming we will read starting from `file_offset`.
pub fn new_for_offset(file_offset: u64, align: usize) -> Self {
assert!(align.is_power_of_two(), "Alignment must be a power of two");
Self {
buf: AVec::new(align),
offset: (file_offset & (align as u64 - 1)) as usize,
}
}
/// Just like [`Vec::extend_from_slice`].
pub fn extend_from_slice(&mut self, data: &[u8]) {
if self.buf.is_empty() && self.offset > 0 {
self.buf.reserve(self.offset + data.len());
self.buf.extend_from_slice(&ZEROS[..self.offset]);
}
self.buf.extend_from_slice(data);
}
/// Drop the first `count` bytes from the front, like `vec.drain(..count)`.
///
/// Preserves alignment:
/// `&buf[i]` was aligned before => `&buf[i - count]` will be aligned after.
pub fn remove_prefix(&mut self, range: RangeTo<usize>) {
let RangeTo { end: count } = range;
self.offset += count;
let to_drop = self.offset & !(self.buf.alignment() - 1);
self.offset &= self.buf.alignment() - 1;
if to_drop > 0 {
self.buf.copy_within(to_drop.., 0);
self.buf.truncate(self.buf.len() - to_drop);
}
}
}
impl Deref for AlignedBuf {
type Target = [u8];
fn deref(&self) -> &[u8] {
if self.buf.is_empty() {
&[]
} else {
&self.buf[self.offset..]
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_aligned_buf() {
let digits = (0..123).collect::<Vec<u8>>();
let ptr_of = |x: &u8| x as *const _ as usize;
let mut buf = AlignedBuf::new_for_offset(0xABC3, align_of::<u128>());
assert_eq!(buf.buf.capacity(), 0, "The constructor shouldn't allocate");
assert_eq!(&buf[..], [] as [u8; 0]);
buf.extend_from_slice(&digits[0..12]);
assert_eq!(&buf[..], &digits[0..12]);
assert_eq!(ptr_of(&buf[5]) % align_of::<u64>(), 0); // 0xABC3 + 5
buf.extend_from_slice(&digits[12..123]);
assert_eq!(&buf[..], &digits[0..123]);
assert_eq!(ptr_of(&buf[5]) % align_of::<u64>(), 0); // 0xABC3 + 5
buf.remove_prefix(..26);
assert_eq!(&buf[..], &digits[26..123]);
assert_eq!(ptr_of(&buf[3]) % align_of::<u64>(), 0); // 0xABC3 + 26 + 3
}
}