Skip to main content

rustfs_erasure_codec/galois_8/
aligned.rs

1extern crate alloc;
2
3use alloc::alloc::{alloc_zeroed, dealloc, handle_alloc_error};
4use alloc::vec;
5use alloc::vec::Vec;
6use core::alloc::Layout;
7use core::fmt;
8use core::iter::FromIterator;
9use core::ops::{Deref, DerefMut};
10use core::ptr::NonNull;
11use core::slice;
12
13use crate::ShardSlot;
14
15pub const SHARD_ALIGNMENT: usize = 64;
16
17pub struct AlignedShard {
18    ptr: NonNull<u8>,
19    len: usize,
20}
21
22impl AlignedShard {
23    pub fn new_zeroed(len: usize) -> Self {
24        if len == 0 {
25            return Self {
26                ptr: NonNull::dangling(),
27                len: 0,
28            };
29        }
30
31        let layout = Layout::from_size_align(len, SHARD_ALIGNMENT)
32            .expect("aligned shard layout must be valid: len or alignment overflow");
33        // SAFETY: `layout` is valid (checked above). `alloc_zeroed` returns a
34        // uniquely owned, zero-filled allocation or null on OOM. The returned
35        // pointer is valid for `layout.size()` bytes.
36        let ptr = unsafe { alloc_zeroed(layout) };
37        let ptr = NonNull::new(ptr).unwrap_or_else(|| handle_alloc_error(layout));
38
39        Self { ptr, len }
40    }
41
42    pub fn from_slice(data: &[u8]) -> Self {
43        let mut shard = Self::new_zeroed(data.len());
44        shard.as_mut().copy_from_slice(data);
45        shard
46    }
47
48    pub fn len(&self) -> usize {
49        self.len
50    }
51
52    pub fn is_empty(&self) -> bool {
53        self.len == 0
54    }
55
56    pub fn as_ptr(&self) -> *const u8 {
57        self.ptr.as_ptr()
58    }
59
60    pub fn as_mut_ptr(&mut self) -> *mut u8 {
61        self.ptr.as_ptr()
62    }
63}
64
65impl Clone for AlignedShard {
66    fn clone(&self) -> Self {
67        Self::from_slice(self.as_ref())
68    }
69}
70
71impl fmt::Debug for AlignedShard {
72    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
73        f.debug_struct("AlignedShard")
74            .field("len", &self.len)
75            .field("alignment", &SHARD_ALIGNMENT)
76            .finish()
77    }
78}
79
80impl Drop for AlignedShard {
81    fn drop(&mut self) {
82        if self.len == 0 {
83            return;
84        }
85
86        let layout = Layout::from_size_align(self.len, SHARD_ALIGNMENT)
87            .expect("aligned shard layout must be valid: len or alignment overflow");
88        // SAFETY: `self.ptr` was allocated from `alloc_zeroed` with the same
89        // layout in `new_zeroed`. This type owns the allocation uniquely
90        // (no cloning without explicit Clone impl that creates a new allocation).
91        unsafe {
92            dealloc(self.ptr.as_ptr(), layout);
93        }
94    }
95}
96
97impl Deref for AlignedShard {
98    type Target = [u8];
99
100    fn deref(&self) -> &Self::Target {
101        self.as_ref()
102    }
103}
104
105impl DerefMut for AlignedShard {
106    fn deref_mut(&mut self) -> &mut Self::Target {
107        self.as_mut()
108    }
109}
110
111impl AsRef<[u8]> for AlignedShard {
112    fn as_ref(&self) -> &[u8] {
113        // SAFETY: `self.ptr` points to `self.len` bytes allocated via `alloc_zeroed`
114        // with SHARD_ALIGNMENT, or is `NonNull::dangling()` when `self.len == 0`
115        // (which produces a valid empty slice). The allocation is uniquely owned
116        // by this value and outlives the returned reference.
117        unsafe { slice::from_raw_parts(self.ptr.as_ptr(), self.len) }
118    }
119}
120
121impl AsMut<[u8]> for AlignedShard {
122    fn as_mut(&mut self) -> &mut [u8] {
123        // SAFETY: same as `as_ref`. `&mut self` guarantees unique mutable access
124        // — no other reference to the same memory can exist concurrently.
125        unsafe { slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len) }
126    }
127}
128
129impl FromIterator<u8> for AlignedShard {
130    fn from_iter<T: IntoIterator<Item = u8>>(iter: T) -> Self {
131        let bytes: Vec<u8> = iter.into_iter().collect();
132        Self::from_slice(&bytes)
133    }
134}
135
136// SAFETY: `AlignedShard` owns a heap allocation of plain `u8` bytes with no
137// interior mutability. Moving it across threads transfers ownership and cannot
138// create aliased mutable references. The `NonNull<u8>` pointer is not shared
139// across threads — it follows the value.
140unsafe impl Send for AlignedShard {}
141// SAFETY: Shared `&AlignedShard` only exposes immutable `[u8]` slices via
142// `as_ref()`. Mutable access requires `&mut self`, which the borrow checker
143// ensures is exclusive. No interior mutability or shared mutable state exists.
144unsafe impl Sync for AlignedShard {}
145
146pub fn alloc_aligned_shards(total_shards: usize, shard_len: usize) -> Vec<AlignedShard> {
147    (0..total_shards)
148        .map(|_| AlignedShard::new_zeroed(shard_len))
149        .collect()
150}
151
152pub fn alloc_shard_slots(total_shards: usize, shard_len: usize) -> Vec<ShardSlot<Vec<u8>>> {
153    (0..total_shards)
154        .map(|_| ShardSlot::new_missing(vec![0u8; shard_len]))
155        .collect()
156}
157
158pub fn shards_to_slots<T: Clone>(shards: &[T]) -> Vec<ShardSlot<T>> {
159    shards.iter().cloned().map(ShardSlot::new_present).collect()
160}
161
162pub fn mark_missing_slots<T>(slots: &mut [ShardSlot<T>], missing_indices: &[usize]) {
163    for &idx in missing_indices {
164        if let Some(slot) = slots.get_mut(idx) {
165            slot.mark_missing();
166        }
167    }
168}
169
170impl crate::ReedSolomon<super::Field> {
171    pub fn alloc_aligned(&self, shard_len: usize) -> Vec<AlignedShard> {
172        alloc_aligned_shards(self.total_shard_count(), shard_len)
173    }
174
175    pub fn alloc_shard_slots(&self, shard_len: usize) -> Vec<ShardSlot<Vec<u8>>> {
176        alloc_shard_slots(self.total_shard_count(), shard_len)
177    }
178}