1use crate::{ExitError, ExitFatal};
2use alloc::vec::Vec;
3use core::cmp::{max, min};
4use core::ops::{BitAnd, Not};
5use primitive_types::U256;
6
7#[derive(Clone, Debug)]
10pub struct Memory {
11 data: Vec<u8>,
12 effective_len: U256,
13 limit: usize,
14}
15
16impl Memory {
17 pub fn new(limit: usize) -> Self {
19 Self {
20 data: Vec::new(),
21 effective_len: U256::zero(),
22 limit,
23 }
24 }
25
26 pub fn limit(&self) -> usize {
28 self.limit
29 }
30
31 pub fn len(&self) -> usize {
33 self.data.len()
34 }
35
36 pub fn effective_len(&self) -> U256 {
38 self.effective_len
39 }
40
41 pub fn is_empty(&self) -> bool {
43 self.len() == 0
44 }
45
46 pub fn data(&self) -> &Vec<u8> {
48 &self.data
49 }
50
51 pub fn resize_offset(&mut self, offset: U256, len: U256) -> Result<(), ExitError> {
54 if len == U256::zero() {
55 return Ok(());
56 }
57
58 if let Some(end) = offset.checked_add(len) {
59 self.resize_end(end)
60 } else {
61 Err(ExitError::InvalidRange)
62 }
63 }
64
65 pub fn resize_end(&mut self, end: U256) -> Result<(), ExitError> {
67 if end > self.effective_len {
68 let new_end = next_multiple_of_32(end).ok_or(ExitError::InvalidRange)?;
69 self.effective_len = new_end;
70 }
71
72 Ok(())
73 }
74
75 #[allow(clippy::slow_vector_initialization)]
82 pub fn get(&self, offset: usize, size: usize) -> Vec<u8> {
85 let mut ret = Vec::new();
86 ret.resize(size, 0);
87
88 #[allow(clippy::needless_range_loop)]
89 for index in 0..size {
90 let position = offset + index;
91 if position >= self.data.len() {
92 break;
93 }
94
95 ret[index] = self.data[position];
96 }
97
98 ret
99 }
100
101 pub fn set(
104 &mut self,
105 offset: usize,
106 value: &[u8],
107 target_size: Option<usize>,
108 ) -> Result<(), ExitFatal> {
109 let target_size = target_size.unwrap_or(value.len());
110 if target_size == 0 {
111 return Ok(());
112 }
113
114 if offset
115 .checked_add(target_size)
116 .map(|pos| pos > self.limit)
117 .unwrap_or(true)
118 {
119 return Err(ExitFatal::NotSupported);
120 }
121
122 if self.data.len() < offset + target_size {
123 self.data.resize(offset + target_size, 0);
124 }
125
126 if target_size > value.len() {
127 self.data[offset..((value.len()) + offset)].clone_from_slice(value);
128 for index in (value.len())..target_size {
129 self.data[offset + index] = 0;
130 }
131 } else {
132 self.data[offset..(target_size + offset)].clone_from_slice(&value[..target_size]);
133 }
134
135 Ok(())
136 }
137
138 pub fn copy_large(
140 &mut self,
141 memory_offset: U256,
142 data_offset: U256,
143 len: U256,
144 data: &[u8],
145 ) -> Result<(), ExitFatal> {
146 if len.is_zero() {
152 return Ok(());
153 }
154
155 let memory_offset = if memory_offset > U256::from(usize::MAX) {
156 return Err(ExitFatal::NotSupported);
157 } else {
158 memory_offset.as_usize()
159 };
160
161 let ulen = if len > U256::from(usize::MAX) {
162 return Err(ExitFatal::NotSupported);
163 } else {
164 len.as_usize()
165 };
166
167 let data = if let Some(end) = data_offset.checked_add(len) {
168 if end > U256::from(usize::MAX) {
169 &[]
170 } else {
171 let data_offset = data_offset.as_usize();
172 let end = end.as_usize();
173
174 if data_offset > data.len() {
175 &[]
176 } else {
177 &data[data_offset..min(end, data.len())]
178 }
179 }
180 } else {
181 &[]
182 };
183
184 self.set(memory_offset, data, Some(ulen))
185 }
186
187 pub fn copy(&mut self, dst: usize, src: usize, len: usize) {
189 let resize_offset = max(dst, src);
190 if self.data.len() < resize_offset + len {
191 self.data.resize(resize_offset + len, 0);
192 }
193 self.data.copy_within(src..src + len, dst);
194 }
195}
196
197#[inline]
199fn next_multiple_of_32(x: U256) -> Option<U256> {
200 let r = x.low_u32().bitand(31).not().wrapping_add(1).bitand(31);
201 x.checked_add(r.into())
202}
203
204#[cfg(test)]
205mod tests {
206 use super::{next_multiple_of_32, Memory, U256};
207
208 #[test]
209 fn test_next_multiple_of_32() {
210 for i in 0..32 {
212 let x = U256::from(i * 32);
213 assert_eq!(Some(x), next_multiple_of_32(x));
214 }
215
216 for x in 0..1024 {
218 if x % 32 == 0 {
219 continue;
220 }
221 let next_multiple = x + 32 - (x % 32);
222 assert_eq!(
223 Some(U256::from(next_multiple)),
224 next_multiple_of_32(x.into())
225 );
226 }
227
228 let last_multiple_of_32 = U256::MAX & !U256::from(31);
230 for i in 0..63 {
231 let x = U256::MAX - U256::from(i);
232 if x > last_multiple_of_32 {
233 assert_eq!(None, next_multiple_of_32(x));
234 } else {
235 assert_eq!(Some(last_multiple_of_32), next_multiple_of_32(x));
236 }
237 }
238 }
239
240 #[test]
241 fn test_memory_copy_works() {
242 let mut memory = Memory::new(100usize);
244
245 memory.set(3usize, &[1u8, 2u8, 3u8, 4u8], None).unwrap();
250 assert_eq!(memory.data(), &[0u8, 0u8, 0u8, 1u8, 2u8, 3u8, 4u8].to_vec());
251
252 memory.copy(0usize, 3usize, 1usize);
255
256 assert_eq!(memory.data(), &[1u8, 0u8, 0u8, 1u8, 2u8, 3u8, 4u8].to_vec());
258 }
259
260 #[test]
261 fn test_memory_copy_resize() {
262 let mut memory = Memory::new(100usize);
264
265 memory.set(3usize, &[1u8, 2u8, 3u8, 4u8], None).unwrap();
270 assert_eq!(memory.data(), &[0u8, 0u8, 0u8, 1u8, 2u8, 3u8, 4u8].to_vec());
271
272 memory.copy(3usize, 6usize, 2usize);
276
277 assert_eq!(
280 memory.data(),
281 &[0u8, 0u8, 0u8, 4u8, 0u8, 3u8, 4u8, 0u8].to_vec()
282 );
283 }
284}