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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
extern crate alloc;
use alloc::alloc::{alloc, alloc_zeroed, dealloc, handle_alloc_error};
use core::alloc::Layout;
use smallvec::SmallVec;
const WORK_ALIGNMENT: usize = 64;
#[derive(Debug)]
pub(super) struct FlatWork {
lanes: usize,
lane_len: usize,
ptr: *mut u8,
len: usize,
}
unsafe impl Send for FlatWork {}
unsafe impl Sync for FlatWork {}
impl FlatWork {
/// Create a new FlatWork with zeroed memory.
pub(super) fn new(lanes: usize, lane_len: usize) -> Self {
let len = lanes * lane_len;
let layout =
Layout::from_size_align(len, WORK_ALIGNMENT).expect("FlatWork layout overflow");
// SAFETY: layout is non-zero (lanes > 0, lane_len > 0).
let ptr = unsafe { alloc_zeroed(layout) };
if ptr.is_null() {
handle_alloc_error(layout);
}
Self {
lanes,
lane_len,
ptr,
len,
}
}
/// Create a new FlatWork WITHOUT zeroing (caller must write before read).
///
/// # Safety
/// All lanes must be fully written before any read. The encode path
/// satisfies this: data lanes get `copy_from_slice`, unused lanes get
/// `zero_trailing_lanes` or `fill(0)`.
pub(super) unsafe fn new_uninit(lanes: usize, lane_len: usize) -> Self {
let len = lanes * lane_len;
let layout =
Layout::from_size_align(len, WORK_ALIGNMENT).expect("FlatWork layout overflow");
// SAFETY: layout is non-zero (lanes > 0, lane_len > 0).
let ptr = unsafe { alloc(layout) };
if ptr.is_null() {
handle_alloc_error(layout);
}
Self {
lanes,
lane_len,
ptr,
len,
}
}
/// Reuse this buffer for a new encode call with the same dimensions.
/// Returns true if the buffer was reused, false if it needs replacement.
pub(super) fn can_reuse(&self, lanes: usize, lane_len: usize) -> bool {
self.lanes == lanes && self.lane_len == lane_len
}
pub(super) fn lanes(&self) -> usize {
self.lanes
}
pub(super) fn lane_len(&self) -> usize {
self.lane_len
}
pub(super) fn lane(&self, idx: usize) -> &[u8] {
debug_assert!(
idx < self.lanes,
"lane index {idx} out of bounds (lanes={})",
self.lanes
);
let start = idx * self.lane_len;
// SAFETY: start..start+lane_len is within allocated bounds, self.ptr is valid.
unsafe { core::slice::from_raw_parts(self.ptr.add(start), self.lane_len) }
}
pub(super) fn lane_mut(&mut self, idx: usize) -> &mut [u8] {
debug_assert!(
idx < self.lanes,
"lane index {idx} out of bounds (lanes={})",
self.lanes
);
let start = idx * self.lane_len;
// SAFETY: start..start+lane_len is within allocated bounds, self.ptr is valid,
// and we have exclusive access through &mut self.
unsafe { core::slice::from_raw_parts_mut(self.ptr.add(start), self.lane_len) }
}
/// Build lane view pointers into the cached `views` vector.
/// Returns a mutable slice of `&mut [u8]` views covering `size` bytes per lane.
pub(super) fn with_lane_views<R>(
&mut self,
lanes: usize,
size: usize,
f: impl FnOnce(&mut [&mut [u8]]) -> R,
) -> R {
debug_assert!(
size <= self.lane_len,
"view size {size} exceeds lane_len {}",
self.lane_len
);
debug_assert!(
lanes <= self.lanes,
"requested lanes {lanes} exceeds capacity {}",
self.lanes
);
// Build view pointers. We use a SmallVec on the stack because the views
// contain mutable references that cannot be cached across calls.
let mut views: SmallVec<[&mut [u8]; 96]> = (0..lanes)
.map(|i| {
let start = i * self.lane_len;
// SAFETY: each lane is at a distinct offset, no overlap.
unsafe {
let ptr = self.ptr.add(start);
&mut *core::ptr::slice_from_raw_parts_mut(ptr, size)
}
})
.collect();
f(&mut views)
}
}
impl Drop for FlatWork {
fn drop(&mut self) {
if self.len > 0 {
let layout = Layout::from_size_align(self.len, WORK_ALIGNMENT)
.expect("FlatWork layout overflow");
// SAFETY: self.ptr was allocated with the same layout.
unsafe { dealloc(self.ptr, layout) };
}
}
}