Skip to main content

wave_emu/
memory.rs

1// Copyright 2026 Ojima Abraham
2// SPDX-License-Identifier: Apache-2.0
3
4//! Memory subsystem implementation. `DeviceMemory` is global shared storage across
5//!
6//! all workgroups. `LocalMemory` is per-workgroup scratch space. Both support byte,
7//! half, word, and double-word loads/stores with bounds checking.
8
9use crate::EmulatorError;
10
11#[derive(Debug)]
12pub struct DeviceMemory {
13    data: Vec<u8>,
14}
15
16impl DeviceMemory {
17    pub fn new(size: usize) -> Self {
18        Self {
19            data: vec![0; size],
20        }
21    }
22
23    pub fn size(&self) -> usize {
24        self.data.len()
25    }
26
27    fn check_bounds(&self, addr: u64, size: usize) -> Result<usize, EmulatorError> {
28        let addr = addr as usize;
29        if addr
30            .checked_add(size)
31            .is_none_or(|end| end > self.data.len())
32        {
33            return Err(EmulatorError::MemoryOutOfBounds {
34                address: addr as u64,
35            });
36        }
37        Ok(addr)
38    }
39
40    pub fn read_u8(&self, addr: u64) -> Result<u8, EmulatorError> {
41        let addr = self.check_bounds(addr, 1)?;
42        Ok(self.data[addr])
43    }
44
45    pub fn read_u16(&self, addr: u64) -> Result<u16, EmulatorError> {
46        let addr = self.check_bounds(addr, 2)?;
47        Ok(u16::from_le_bytes([self.data[addr], self.data[addr + 1]]))
48    }
49
50    pub fn read_u32(&self, addr: u64) -> Result<u32, EmulatorError> {
51        let addr = self.check_bounds(addr, 4)?;
52        Ok(u32::from_le_bytes([
53            self.data[addr],
54            self.data[addr + 1],
55            self.data[addr + 2],
56            self.data[addr + 3],
57        ]))
58    }
59
60    pub fn read_u64(&self, addr: u64) -> Result<u64, EmulatorError> {
61        let addr = self.check_bounds(addr, 8)?;
62        Ok(u64::from_le_bytes([
63            self.data[addr],
64            self.data[addr + 1],
65            self.data[addr + 2],
66            self.data[addr + 3],
67            self.data[addr + 4],
68            self.data[addr + 5],
69            self.data[addr + 6],
70            self.data[addr + 7],
71        ]))
72    }
73
74    pub fn read_u128(&self, addr: u64) -> Result<u128, EmulatorError> {
75        let addr = self.check_bounds(addr, 16)?;
76        let mut bytes = [0u8; 16];
77        bytes.copy_from_slice(&self.data[addr..addr + 16]);
78        Ok(u128::from_le_bytes(bytes))
79    }
80
81    pub fn write_u8(&mut self, addr: u64, value: u8) -> Result<(), EmulatorError> {
82        let addr = self.check_bounds(addr, 1)?;
83        self.data[addr] = value;
84        Ok(())
85    }
86
87    pub fn write_u16(&mut self, addr: u64, value: u16) -> Result<(), EmulatorError> {
88        let addr = self.check_bounds(addr, 2)?;
89        let bytes = value.to_le_bytes();
90        self.data[addr..addr + 2].copy_from_slice(&bytes);
91        Ok(())
92    }
93
94    pub fn write_u32(&mut self, addr: u64, value: u32) -> Result<(), EmulatorError> {
95        let addr = self.check_bounds(addr, 4)?;
96        let bytes = value.to_le_bytes();
97        self.data[addr..addr + 4].copy_from_slice(&bytes);
98        Ok(())
99    }
100
101    pub fn write_u64(&mut self, addr: u64, value: u64) -> Result<(), EmulatorError> {
102        let addr = self.check_bounds(addr, 8)?;
103        let bytes = value.to_le_bytes();
104        self.data[addr..addr + 8].copy_from_slice(&bytes);
105        Ok(())
106    }
107
108    pub fn write_u128(&mut self, addr: u64, value: u128) -> Result<(), EmulatorError> {
109        let addr = self.check_bounds(addr, 16)?;
110        let bytes = value.to_le_bytes();
111        self.data[addr..addr + 16].copy_from_slice(&bytes);
112        Ok(())
113    }
114
115    pub fn write_slice(&mut self, offset: u64, data: &[u8]) -> Result<(), EmulatorError> {
116        let addr = self.check_bounds(offset, data.len())?;
117        self.data[addr..addr + data.len()].copy_from_slice(data);
118        Ok(())
119    }
120
121    pub fn atomic_add(&mut self, addr: u64, value: u32) -> Result<u32, EmulatorError> {
122        let old = self.read_u32(addr)?;
123        self.write_u32(addr, old.wrapping_add(value))?;
124        Ok(old)
125    }
126
127    pub fn atomic_sub(&mut self, addr: u64, value: u32) -> Result<u32, EmulatorError> {
128        let old = self.read_u32(addr)?;
129        self.write_u32(addr, old.wrapping_sub(value))?;
130        Ok(old)
131    }
132
133    pub fn atomic_min(&mut self, addr: u64, value: u32) -> Result<u32, EmulatorError> {
134        let old = self.read_u32(addr)?;
135        self.write_u32(addr, old.min(value))?;
136        Ok(old)
137    }
138
139    pub fn atomic_max(&mut self, addr: u64, value: u32) -> Result<u32, EmulatorError> {
140        let old = self.read_u32(addr)?;
141        self.write_u32(addr, old.max(value))?;
142        Ok(old)
143    }
144
145    pub fn atomic_and(&mut self, addr: u64, value: u32) -> Result<u32, EmulatorError> {
146        let old = self.read_u32(addr)?;
147        self.write_u32(addr, old & value)?;
148        Ok(old)
149    }
150
151    pub fn atomic_or(&mut self, addr: u64, value: u32) -> Result<u32, EmulatorError> {
152        let old = self.read_u32(addr)?;
153        self.write_u32(addr, old | value)?;
154        Ok(old)
155    }
156
157    pub fn atomic_xor(&mut self, addr: u64, value: u32) -> Result<u32, EmulatorError> {
158        let old = self.read_u32(addr)?;
159        self.write_u32(addr, old ^ value)?;
160        Ok(old)
161    }
162
163    pub fn atomic_exchange(&mut self, addr: u64, value: u32) -> Result<u32, EmulatorError> {
164        let old = self.read_u32(addr)?;
165        self.write_u32(addr, value)?;
166        Ok(old)
167    }
168
169    pub fn atomic_cas(
170        &mut self,
171        addr: u64,
172        expected: u32,
173        desired: u32,
174    ) -> Result<u32, EmulatorError> {
175        let old = self.read_u32(addr)?;
176        if old == expected {
177            self.write_u32(addr, desired)?;
178        }
179        Ok(old)
180    }
181
182    pub fn atomic_add_i32(&mut self, addr: u64, value: i32) -> Result<i32, EmulatorError> {
183        let old = self.read_u32(addr)? as i32;
184        self.write_u32(addr, old.wrapping_add(value) as u32)?;
185        Ok(old)
186    }
187
188    pub fn atomic_min_i32(&mut self, addr: u64, value: i32) -> Result<i32, EmulatorError> {
189        let old = self.read_u32(addr)? as i32;
190        self.write_u32(addr, old.min(value) as u32)?;
191        Ok(old)
192    }
193
194    pub fn atomic_max_i32(&mut self, addr: u64, value: i32) -> Result<i32, EmulatorError> {
195        let old = self.read_u32(addr)? as i32;
196        self.write_u32(addr, old.max(value) as u32)?;
197        Ok(old)
198    }
199
200    pub fn atomic_add_f32(&mut self, addr: u64, value: f32) -> Result<f32, EmulatorError> {
201        let old_bits = self.read_u32(addr)?;
202        let old = f32::from_bits(old_bits);
203        let new = old + value;
204        self.write_u32(addr, new.to_bits())?;
205        Ok(old)
206    }
207}
208
209#[derive(Debug)]
210pub struct LocalMemory {
211    data: Vec<u8>,
212}
213
214impl LocalMemory {
215    pub fn new(size: usize) -> Self {
216        Self {
217            data: vec![0; size],
218        }
219    }
220
221    pub fn size(&self) -> usize {
222        self.data.len()
223    }
224
225    fn check_bounds(&self, addr: u32, size: usize) -> Result<usize, EmulatorError> {
226        let addr = addr as usize;
227        if addr
228            .checked_add(size)
229            .is_none_or(|end| end > self.data.len())
230        {
231            return Err(EmulatorError::MemoryOutOfBounds {
232                address: addr as u64,
233            });
234        }
235        Ok(addr)
236    }
237
238    pub fn read_u8(&self, addr: u32) -> Result<u8, EmulatorError> {
239        let addr = self.check_bounds(addr, 1)?;
240        Ok(self.data[addr])
241    }
242
243    pub fn read_u16(&self, addr: u32) -> Result<u16, EmulatorError> {
244        let addr = self.check_bounds(addr, 2)?;
245        Ok(u16::from_le_bytes([self.data[addr], self.data[addr + 1]]))
246    }
247
248    pub fn read_u32(&self, addr: u32) -> Result<u32, EmulatorError> {
249        let addr = self.check_bounds(addr, 4)?;
250        Ok(u32::from_le_bytes([
251            self.data[addr],
252            self.data[addr + 1],
253            self.data[addr + 2],
254            self.data[addr + 3],
255        ]))
256    }
257
258    pub fn read_u64(&self, addr: u32) -> Result<u64, EmulatorError> {
259        let addr = self.check_bounds(addr, 8)?;
260        Ok(u64::from_le_bytes([
261            self.data[addr],
262            self.data[addr + 1],
263            self.data[addr + 2],
264            self.data[addr + 3],
265            self.data[addr + 4],
266            self.data[addr + 5],
267            self.data[addr + 6],
268            self.data[addr + 7],
269        ]))
270    }
271
272    pub fn write_u8(&mut self, addr: u32, value: u8) -> Result<(), EmulatorError> {
273        let addr = self.check_bounds(addr, 1)?;
274        self.data[addr] = value;
275        Ok(())
276    }
277
278    pub fn write_u16(&mut self, addr: u32, value: u16) -> Result<(), EmulatorError> {
279        let addr = self.check_bounds(addr, 2)?;
280        let bytes = value.to_le_bytes();
281        self.data[addr..addr + 2].copy_from_slice(&bytes);
282        Ok(())
283    }
284
285    pub fn write_u32(&mut self, addr: u32, value: u32) -> Result<(), EmulatorError> {
286        let addr = self.check_bounds(addr, 4)?;
287        let bytes = value.to_le_bytes();
288        self.data[addr..addr + 4].copy_from_slice(&bytes);
289        Ok(())
290    }
291
292    pub fn write_u64(&mut self, addr: u32, value: u64) -> Result<(), EmulatorError> {
293        let addr = self.check_bounds(addr, 8)?;
294        let bytes = value.to_le_bytes();
295        self.data[addr..addr + 8].copy_from_slice(&bytes);
296        Ok(())
297    }
298
299    pub fn atomic_add(&mut self, addr: u32, value: u32) -> Result<u32, EmulatorError> {
300        let old = self.read_u32(addr)?;
301        self.write_u32(addr, old.wrapping_add(value))?;
302        Ok(old)
303    }
304
305    pub fn atomic_sub(&mut self, addr: u32, value: u32) -> Result<u32, EmulatorError> {
306        let old = self.read_u32(addr)?;
307        self.write_u32(addr, old.wrapping_sub(value))?;
308        Ok(old)
309    }
310
311    pub fn atomic_min(&mut self, addr: u32, value: u32) -> Result<u32, EmulatorError> {
312        let old = self.read_u32(addr)?;
313        self.write_u32(addr, old.min(value))?;
314        Ok(old)
315    }
316
317    pub fn atomic_max(&mut self, addr: u32, value: u32) -> Result<u32, EmulatorError> {
318        let old = self.read_u32(addr)?;
319        self.write_u32(addr, old.max(value))?;
320        Ok(old)
321    }
322
323    pub fn atomic_and(&mut self, addr: u32, value: u32) -> Result<u32, EmulatorError> {
324        let old = self.read_u32(addr)?;
325        self.write_u32(addr, old & value)?;
326        Ok(old)
327    }
328
329    pub fn atomic_or(&mut self, addr: u32, value: u32) -> Result<u32, EmulatorError> {
330        let old = self.read_u32(addr)?;
331        self.write_u32(addr, old | value)?;
332        Ok(old)
333    }
334
335    pub fn atomic_xor(&mut self, addr: u32, value: u32) -> Result<u32, EmulatorError> {
336        let old = self.read_u32(addr)?;
337        self.write_u32(addr, old ^ value)?;
338        Ok(old)
339    }
340
341    pub fn atomic_exchange(&mut self, addr: u32, value: u32) -> Result<u32, EmulatorError> {
342        let old = self.read_u32(addr)?;
343        self.write_u32(addr, value)?;
344        Ok(old)
345    }
346
347    pub fn atomic_cas(
348        &mut self,
349        addr: u32,
350        expected: u32,
351        desired: u32,
352    ) -> Result<u32, EmulatorError> {
353        let old = self.read_u32(addr)?;
354        if old == expected {
355            self.write_u32(addr, desired)?;
356        }
357        Ok(old)
358    }
359}
360
361#[cfg(test)]
362mod tests {
363    use super::*;
364
365    #[test]
366    fn test_device_memory_read_write_u32() {
367        let mut mem = DeviceMemory::new(1024);
368        mem.write_u32(0x100, 0xDEADBEEF).unwrap();
369        assert_eq!(mem.read_u32(0x100).unwrap(), 0xDEADBEEF);
370    }
371
372    #[test]
373    fn test_device_memory_read_write_u64() {
374        let mut mem = DeviceMemory::new(1024);
375        mem.write_u64(0x100, 0x123456789ABCDEF0).unwrap();
376        assert_eq!(mem.read_u64(0x100).unwrap(), 0x123456789ABCDEF0);
377    }
378
379    #[test]
380    fn test_device_memory_bounds_check() {
381        let mem = DeviceMemory::new(1024);
382        assert!(mem.read_u32(1024).is_err());
383        assert!(mem.read_u32(1021).is_err());
384    }
385
386    #[test]
387    fn test_device_memory_atomic_add() {
388        let mut mem = DeviceMemory::new(1024);
389        mem.write_u32(0x100, 10).unwrap();
390
391        let old = mem.atomic_add(0x100, 5).unwrap();
392        assert_eq!(old, 10);
393        assert_eq!(mem.read_u32(0x100).unwrap(), 15);
394    }
395
396    #[test]
397    fn test_device_memory_atomic_cas() {
398        let mut mem = DeviceMemory::new(1024);
399        mem.write_u32(0x100, 42).unwrap();
400
401        let old = mem.atomic_cas(0x100, 42, 100).unwrap();
402        assert_eq!(old, 42);
403        assert_eq!(mem.read_u32(0x100).unwrap(), 100);
404
405        let old = mem.atomic_cas(0x100, 42, 200).unwrap();
406        assert_eq!(old, 100);
407        assert_eq!(mem.read_u32(0x100).unwrap(), 100);
408    }
409
410    #[test]
411    fn test_local_memory_read_write() {
412        let mut mem = LocalMemory::new(1024);
413        mem.write_u32(0x100, 0xCAFEBABE).unwrap();
414        assert_eq!(mem.read_u32(0x100).unwrap(), 0xCAFEBABE);
415    }
416
417    #[test]
418    fn test_local_memory_atomic_ops() {
419        let mut mem = LocalMemory::new(1024);
420        mem.write_u32(0, 100).unwrap();
421
422        let old = mem.atomic_sub(0, 30).unwrap();
423        assert_eq!(old, 100);
424        assert_eq!(mem.read_u32(0).unwrap(), 70);
425
426        let old = mem.atomic_min(0, 50).unwrap();
427        assert_eq!(old, 70);
428        assert_eq!(mem.read_u32(0).unwrap(), 50);
429    }
430}