dynamo_memory/actions.rs
1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Storage actions.
5
6use super::{MemoryDescription, StorageError};
7
8/// Extension trait for storage types that support memory setting operations
9pub trait Memset: MemoryDescription {
10 /// Sets a region of memory to a specific value
11 ///
12 /// # Arguments
13 /// * `value` - The value to set (will be truncated to u8)
14 /// * `offset` - Offset in bytes from the start of the storage
15 /// * `size` - Number of bytes to set
16 ///
17 /// # Safety
18 /// The caller must ensure:
19 /// - offset + size <= self.size()
20 /// - No other references exist to the memory region being set
21 fn memset(&mut self, value: u8, offset: usize, size: usize) -> Result<(), StorageError>;
22}
23
24/// Extension trait for storage types that support slicing operations
25pub trait Slice: MemoryDescription + 'static {
26 /// Returns an immutable byte slice view of the entire storage region
27 ///
28 /// # Safety
29 /// This is an unsafe method. The caller must ensure:
30 /// - The memory region remains valid for the lifetime of the returned slice
31 /// - The memory region is properly initialized
32 /// - No concurrent mutable access occurs while the slice is in use
33 /// - The memory backing this storage remains valid (implementors with owned
34 /// memory satisfy this, but care must be taken with unowned memory regions)
35 unsafe fn as_slice(&self) -> Result<&[u8], StorageError>;
36
37 /// Returns an immutable byte slice view of a subregion
38 ///
39 /// # Arguments
40 /// * `offset` - Offset in bytes from the start of the storage
41 /// * `len` - Number of bytes to slice
42 ///
43 /// # Safety
44 /// The caller must ensure:
45 /// - offset + len <= self.size()
46 /// - The memory region is valid and initialized
47 /// - No concurrent mutable access occurs while the slice is in use
48 fn slice(&self, offset: usize, len: usize) -> Result<&[u8], StorageError> {
49 // SAFETY: Caller guarantees memory validity per trait's safety contract
50 let slice = unsafe { self.as_slice()? };
51
52 // validate offset and len
53 if offset.saturating_add(len) > slice.len() {
54 return Err(StorageError::Unsupported("slice out of bounds".into()));
55 }
56
57 slice
58 .get(offset..offset.saturating_add(len))
59 .ok_or_else(|| StorageError::Unsupported("slice out of bounds".into()))
60 }
61
62 /// Returns a typed immutable slice view of the entire storage region
63 ///
64 /// # Safety
65 /// The caller must ensure:
66 /// - The memory region is valid and initialized
67 /// - The memory is properly aligned for type T
68 /// - The size is a multiple of `size_of::<T>()`
69 /// - No concurrent mutable access occurs while the slice is in use
70 /// - The data represents valid values of type T
71 fn as_slice_typed<T: Sized>(&self) -> Result<&[T], StorageError> {
72 // SAFETY: Caller guarantees memory validity per trait's safety contract
73 let bytes = unsafe { self.as_slice()? };
74 let ptr = bytes.as_ptr() as *const T;
75 let elem_size = std::mem::size_of::<T>();
76 if elem_size == 0 {
77 return Err(StorageError::Unsupported(
78 "zero-sized types are not supported".into(),
79 ));
80 }
81 let len = bytes.len() / elem_size;
82
83 if !(bytes.as_ptr() as usize).is_multiple_of(std::mem::align_of::<T>()) {
84 return Err(StorageError::Unsupported(format!(
85 "memory not aligned for type (required alignment: {})",
86 std::mem::align_of::<T>()
87 )));
88 }
89
90 if bytes.len() % elem_size != 0 {
91 return Err(StorageError::Unsupported(format!(
92 "size {} is not a multiple of type size {}",
93 bytes.len(),
94 elem_size
95 )));
96 }
97
98 // SAFETY: Caller guarantees memory is valid, aligned, and properly initialized for T
99 Ok(unsafe { std::slice::from_raw_parts(ptr, len) })
100 }
101
102 /// Returns a typed immutable slice view of a subregion
103 ///
104 /// # Arguments
105 /// * `offset` - Offset in bytes from the start of the storage
106 /// * `len` - Number of elements of type T to slice
107 ///
108 /// # Safety
109 /// The caller must ensure:
110 /// - offset + (len * size_of::<T>()) <= self.size()
111 /// - offset is properly aligned for type T
112 /// - The memory region is valid and initialized
113 /// - No concurrent mutable access occurs while the slice is in use
114 /// - The data represents valid values of type T
115 fn slice_typed<T: Sized>(&self, offset: usize, len: usize) -> Result<&[T], StorageError> {
116 let type_size = std::mem::size_of::<T>();
117 let byte_len = len
118 .checked_mul(type_size)
119 .ok_or_else(|| StorageError::Unsupported("length overflow".into()))?;
120
121 let bytes = self.slice(offset, byte_len)?;
122 let ptr = bytes.as_ptr() as *const T;
123
124 if !(bytes.as_ptr() as usize).is_multiple_of(std::mem::align_of::<T>()) {
125 return Err(StorageError::Unsupported(format!(
126 "memory not aligned for type (required alignment: {})",
127 std::mem::align_of::<T>()
128 )));
129 }
130
131 // SAFETY: Caller guarantees memory is valid, aligned, and properly initialized for T
132 Ok(unsafe { std::slice::from_raw_parts(ptr, len) })
133 }
134}
135
136pub trait SliceMut: MemoryDescription + 'static {
137 /// Returns a mutable byte slice view of the entire storage region
138 ///
139 /// # Safety
140 /// This is an unsafe method. The caller must ensure:
141 /// - The memory region remains valid for the lifetime of the returned slice
142 /// - The memory region is valid and accessible
143 /// - No other references (mutable or immutable) exist to this memory region
144 /// - The memory backing this storage remains valid (implementors with owned
145 /// memory satisfy this, but care must be taken with unowned memory regions)
146 unsafe fn as_slice_mut(&mut self) -> Result<&mut [u8], StorageError>;
147
148 /// Returns a mutable byte slice view of a subregion
149 ///
150 /// # Arguments
151 /// * `offset` - Offset in bytes from the start of the storage
152 /// * `len` - Number of bytes to slice
153 ///
154 /// # Safety
155 /// The caller must ensure:
156 /// - offset + len <= self.size()
157 /// - The memory region is valid
158 /// - No other references (mutable or immutable) exist to this memory region
159 fn slice_mut(&mut self, offset: usize, len: usize) -> Result<&mut [u8], StorageError> {
160 // SAFETY: Caller guarantees memory validity per trait's safety contract
161 let slice = unsafe { self.as_slice_mut()? };
162
163 // validate offset and len
164 if offset.saturating_add(len) > slice.len() {
165 return Err(StorageError::Unsupported("slice out of bounds".into()));
166 }
167
168 slice
169 .get_mut(offset..offset.saturating_add(len))
170 .ok_or_else(|| StorageError::Unsupported("slice out of bounds".into()))
171 }
172
173 /// Returns a typed mutable slice view of the entire storage region
174 ///
175 /// # Safety
176 /// The caller must ensure:
177 /// - The memory region is valid
178 /// - The memory is properly aligned for type T
179 /// - The size is a multiple of `size_of::<T>()`
180 /// - No other references (mutable or immutable) exist to this memory region
181 fn as_slice_typed_mut<T: Sized>(&mut self) -> Result<&mut [T], StorageError> {
182 // SAFETY: Caller guarantees memory validity per trait's safety contract
183 let bytes = unsafe { self.as_slice_mut()? };
184 let ptr = bytes.as_mut_ptr() as *mut T;
185 let len = bytes.len() / std::mem::size_of::<T>();
186
187 if !(bytes.as_ptr() as usize).is_multiple_of(std::mem::align_of::<T>()) {
188 return Err(StorageError::Unsupported(format!(
189 "memory not aligned for type (required alignment: {})",
190 std::mem::align_of::<T>()
191 )));
192 }
193
194 if bytes.len() % std::mem::size_of::<T>() != 0 {
195 return Err(StorageError::Unsupported(format!(
196 "size {} is not a multiple of type size {}",
197 bytes.len(),
198 std::mem::size_of::<T>()
199 )));
200 }
201
202 // SAFETY: Caller guarantees memory is valid, aligned, and no aliasing
203 Ok(unsafe { std::slice::from_raw_parts_mut(ptr, len) })
204 }
205
206 /// Returns a typed mutable slice view of a subregion
207 ///
208 /// # Arguments
209 /// * `offset` - Offset in bytes from the start of the storage
210 /// * `len` - Number of elements of type T to slice
211 ///
212 /// # Safety
213 /// The caller must ensure:
214 /// - offset + (len * size_of::<T>()) <= self.size()
215 /// - offset is properly aligned for type T
216 /// - The memory region is valid
217 /// - No other references (mutable or immutable) exist to this memory region
218 fn slice_typed_mut<T: Sized>(
219 &mut self,
220 offset: usize,
221 len: usize,
222 ) -> Result<&mut [T], StorageError> {
223 let type_size = std::mem::size_of::<T>();
224 let byte_len = len
225 .checked_mul(type_size)
226 .ok_or_else(|| StorageError::Unsupported("length overflow".into()))?;
227
228 let bytes = self.slice_mut(offset, byte_len)?;
229 let ptr = bytes.as_mut_ptr() as *mut T;
230
231 if !(bytes.as_ptr() as usize).is_multiple_of(std::mem::align_of::<T>()) {
232 return Err(StorageError::Unsupported(format!(
233 "memory not aligned for type (required alignment: {})",
234 std::mem::align_of::<T>()
235 )));
236 }
237
238 // SAFETY: Caller guarantees memory is valid, aligned, and no aliasing
239 Ok(unsafe { std::slice::from_raw_parts_mut(ptr, len) })
240 }
241}