1use alloc::{vec, vec::Vec};
2use core::{
3 cmp::{max, min},
4 mem,
5 ops::{BitAnd, Not, Range},
6};
7
8use primitive_types::U256;
9
10use crate::error::{ExitException, ExitFatal};
11
12#[derive(Clone, Debug)]
15pub struct Memory {
16 data: Vec<u8>,
17 effective_len: U256,
18 limit: usize,
19}
20
21impl Memory {
22 #[must_use]
24 pub fn new(limit: usize) -> Self {
25 Self {
26 data: Vec::new(),
27 effective_len: U256::zero(),
28 limit,
29 }
30 }
31
32 #[must_use]
34 pub const fn limit(&self) -> usize {
35 self.limit
36 }
37
38 #[must_use]
40 pub fn len(&self) -> usize {
41 self.data.len()
42 }
43
44 #[must_use]
46 pub const fn effective_len(&self) -> U256 {
47 self.effective_len
48 }
49
50 #[must_use]
52 pub fn is_empty(&self) -> bool {
53 self.len() == 0
54 }
55
56 #[must_use]
58 pub const fn data(&self) -> &Vec<u8> {
59 &self.data
60 }
61
62 pub(crate) fn swap_and_clear(&mut self, other: &mut Vec<u8>) {
63 mem::swap(&mut self.data, other);
64 self.data = Vec::new();
65 self.effective_len = U256::zero();
66 }
67
68 pub fn resize_offset(&mut self, offset: U256, len: U256) -> Result<(), ExitException> {
72 if len == U256::zero() {
73 return Ok(());
74 }
75
76 offset
77 .checked_add(len)
78 .map_or(Err(ExitException::InvalidRange), |end| self.resize_end(end))
79 }
80
81 pub fn resize_end(&mut self, end: U256) -> Result<(), ExitException> {
83 if end > self.effective_len {
84 let new_end = next_multiple_of_32(end).ok_or(ExitException::InvalidRange)?;
85 self.effective_len = new_end;
86 }
87
88 Ok(())
89 }
90
91 pub fn resize_to_range(&mut self, return_range: Range<U256>) {
93 let ret = if return_range.start > U256::from(usize::MAX) {
94 vec![0; (return_range.end - return_range.start).as_usize()]
95 } else if return_range.end > U256::from(usize::MAX) {
96 let mut ret = self.get(
97 return_range.start.as_usize(),
98 usize::MAX - return_range.start.as_usize(),
99 );
100 while ret.len() < (return_range.end - return_range.start).as_usize() {
101 ret.push(0);
102 }
103 ret
104 } else {
105 self.get(
106 return_range.start.as_usize(),
107 (return_range.end - return_range.start).as_usize(),
108 )
109 };
110 self.data = ret;
111 self.effective_len = return_range.end - return_range.start;
112 }
113
114 #[must_use]
121 pub fn get(&self, offset: usize, size: usize) -> Vec<u8> {
122 let mut ret = vec![0; size];
123
124 #[allow(clippy::needless_range_loop)]
125 for index in 0..size {
126 let position = if let Some(position) = offset.checked_add(index) {
127 position
128 } else {
129 break;
130 };
131
132 if position >= self.data.len() {
133 break;
134 }
135
136 ret[index] = self.data[position];
137 }
138
139 ret
140 }
141
142 pub fn set(
145 &mut self,
146 offset: usize,
147 value: &[u8],
148 target_size: Option<usize>,
149 ) -> Result<(), ExitFatal> {
150 let target_size = target_size.unwrap_or(value.len());
151 if target_size == 0 {
152 return Ok(());
153 }
154
155 if offset
156 .checked_add(target_size)
157 .map(|pos| pos > self.limit)
158 .unwrap_or(true)
159 {
160 return Err(ExitFatal::NotSupported);
161 }
162
163 if self.data.len() < offset + target_size {
164 self.data.resize(offset + target_size, 0);
165 }
166
167 if target_size > value.len() {
168 self.data[offset..((value.len()) + offset)].clone_from_slice(value);
169 for index in (value.len())..target_size {
170 self.data[offset + index] = 0;
171 }
172 } else {
173 self.data[offset..(target_size + offset)].clone_from_slice(&value[..target_size]);
174 }
175
176 Ok(())
177 }
178
179 pub fn copy_large(
181 &mut self,
182 memory_offset: U256,
183 data_offset: U256,
184 len: U256,
185 data: &[u8],
186 ) -> Result<(), ExitFatal> {
187 if len.is_zero() {
193 return Ok(());
194 }
195
196 let memory_offset = if memory_offset > U256::from(usize::MAX) {
197 return Err(ExitFatal::NotSupported);
198 } else {
199 memory_offset.as_usize()
200 };
201
202 let ulen = if len > U256::from(usize::MAX) {
203 return Err(ExitFatal::NotSupported);
204 } else {
205 len.as_usize()
206 };
207
208 let data: &[u8] = data_offset.checked_add(len).map_or(&[], |end| {
209 if end > U256::from(usize::MAX) {
210 &[]
211 } else {
212 let data_offset = data_offset.as_usize();
213 let end = end.as_usize();
214
215 if data_offset > data.len() {
216 &[]
217 } else {
218 &data[data_offset..min(end, data.len())]
219 }
220 }
221 });
222
223 self.set(memory_offset, data, Some(ulen))
224 }
225
226 pub fn copy(&mut self, dst: usize, src: usize, len: usize) {
228 let resize_offset = max(dst, src);
229 if self.data.len() < resize_offset + len {
230 self.data.resize(resize_offset + len, 0);
231 }
232 self.data.copy_within(src..src + len, dst);
233 }
234}
235
236#[inline]
238fn next_multiple_of_32(x: U256) -> Option<U256> {
239 let r = x.low_u32().bitand(31).not().wrapping_add(1).bitand(31);
240 x.checked_add(r.into())
241}
242
243#[cfg(test)]
244mod tests {
245 use super::{Memory, U256, next_multiple_of_32};
246
247 #[test]
248 fn test_next_multiple_of_32() {
249 for i in 0..32 {
251 let x = U256::from(i * 32);
252 assert_eq!(Some(x), next_multiple_of_32(x));
253 }
254
255 for x in 0..1024 {
257 if x % 32 == 0 {
258 continue;
259 }
260 let next_multiple = x + 32 - (x % 32);
261 assert_eq!(
262 Some(U256::from(next_multiple)),
263 next_multiple_of_32(x.into())
264 );
265 }
266
267 let last_multiple_of_32 = U256::MAX & !U256::from(31);
269 for i in 0..63 {
270 let x = U256::MAX - U256::from(i);
271 if x > last_multiple_of_32 {
272 assert_eq!(None, next_multiple_of_32(x));
273 } else {
274 assert_eq!(Some(last_multiple_of_32), next_multiple_of_32(x));
275 }
276 }
277 }
278
279 #[test]
280 fn test_memory_copy_works() {
281 let mut memory = Memory::new(100usize);
283
284 memory.set(3usize, &[1u8, 2u8, 3u8, 4u8], None).unwrap();
289 assert_eq!(memory.data(), &[0u8, 0u8, 0u8, 1u8, 2u8, 3u8, 4u8].to_vec());
290
291 memory.copy(0usize, 3usize, 1usize);
294
295 assert_eq!(memory.data(), &[1u8, 0u8, 0u8, 1u8, 2u8, 3u8, 4u8].to_vec());
297 }
298
299 #[test]
300 fn test_memory_copy_resize() {
301 let mut memory = Memory::new(100usize);
303
304 memory.set(3usize, &[1u8, 2u8, 3u8, 4u8], None).unwrap();
309 assert_eq!(memory.data(), &[0u8, 0u8, 0u8, 1u8, 2u8, 3u8, 4u8].to_vec());
310
311 memory.copy(3usize, 6usize, 2usize);
315
316 assert_eq!(
319 memory.data(),
320 &[0u8, 0u8, 0u8, 4u8, 0u8, 3u8, 4u8, 0u8].to_vec()
321 );
322 }
323}