3 files changed, 40 insertions, 32 deletions

diff --git a/src/core.rs b/src/core.rs
index 3b59782..9e8e8dd 100644
--- a/src/core.rs
+++ b/src/core.rs
@@ -191,33 +191,35 @@ pub fn lhu(state: &mut CpuState, rd: usize, rs1: usize, imm: u32) {
     state.registers[rd] = ((bytes[1] as u16) << 8 | bytes[0] as u16) as u32;
 }
 pub fn sb(state: &mut CpuState, rs1: usize, rs2: usize, imm: u32) {
+    let addr = (state.registers[rs1] + imm) as usize;
     state
         .memory
-        .write(
-            (state.registers[rs1] + imm) as usize,
-            99,
-            state.registers[rs2] as u8,
-        )
+        .write(addr, 99, state.registers[rs2] as u8)
         .unwrap();
+    state.release_reservations(&[addr]);
 }
 
 pub fn sh(state: &mut CpuState, rs1: usize, rs2: usize, imm: u32) {
+    let addr = (state.registers[rs1] + imm) as usize;
+
     let halfword = state.registers[rs2] & 0xFFFF;
-    let index = (state.registers[rs1] + imm) as usize;
 
     state
         .memory
-        .write_multiple_bytes(index, &[(halfword & 0xFF) as u8, (halfword >> 8) as u8], 99)
+        .write_multiple_bytes(addr, &[(halfword & 0xFF) as u8, (halfword >> 8) as u8], 99)
         .unwrap();
+
+    state.release_reservations(&[addr, addr + 1]);
 }
 
 pub fn sw(state: &mut CpuState, rs1: usize, rs2: usize, imm: u32) {
-    let index = (state.registers[rs1] + imm) as usize;
+    let addr = (state.registers[rs1] + imm) as usize;
 
     state
         .memory
-        .write_multiple_bytes(index, &state.registers[rs2].to_le_bytes(), 99)
+        .write_multiple_bytes(addr, &state.registers[rs2].to_le_bytes(), 99)
         .unwrap();
+    state.release_reservations(&[addr, addr + 1, addr + 2, addr + 3]);
 }
 
 pub fn beq(state: &mut CpuState, rs1: usize, rs2: usize, imm: u32) {
@@ -271,12 +273,12 @@ pub fn lui(state: &mut CpuState, rd: usize, imm: u32) {
 pub fn lr_w(state: &mut CpuState, rd: usize, rs1: usize) {
     let addr = state.registers[rs1] as usize;
     state.registers[rd] = load_word(state, addr);
-    state.acquire_reservations(&[addr, addr + 1, addr + 2, addr + 3]);
+    state.acquire_reservations(addr);
 }
 
 pub fn sc_w(state: &mut CpuState, rd: usize, rs1: usize, rs2: usize) {
     let addr = state.registers[rs1] as usize;
-    if state.reservations.contains(&addr) {
+    if state.is_addr_reserved(addr) {
         store_word(state, addr, state.registers[rs2]);
         state.registers[rd] = 0;
     } else {
diff --git a/src/main.rs b/src/main.rs
index a3571c7..3d8008d 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -19,7 +19,10 @@ fn main() {
 
     let mut state = CpuState::new(memory);
 
-    state.memory.write_multiple_bytes(RO_TARGET as usize, f_dat.as_slice(), 99).unwrap();
+    state
+        .memory
+        .write_multiple_bytes(RO_TARGET as usize, f_dat.as_slice(), 99)
+        .unwrap();
     state.memory.write(UART_ADDR as usize, 99, 0).unwrap();
     state.pc = RO_TARGET;
 
@@ -134,18 +137,9 @@ fn handle_instruction(state: &mut CpuState) {
             let addr = (state.registers[rs1] + imm) as usize;
 
             match f3 {
-                0 => {
-                    sb(state, rs1, rs2, imm);
-                    state.release_reservations(&[addr]);
-                }
-                1 => {
-                    sh(state, rs1, rs2, imm);
-                    state.release_reservations(&[addr, addr + 1]);
-                }
-                2 => {
-                    sw(state, rs1, rs2, imm);
-                    state.release_reservations(&[addr, addr + 1, addr + 2, addr + 3]);
-                }
+                0 => sb(state, rs1, rs2, imm),
+                1 => sh(state, rs1, rs2, imm),
+                2 => sw(state, rs1, rs2, imm),
                 _ => panic!("Illegal instruction: {:x}", word),
             };
         }
diff --git a/src/state.rs b/src/state.rs
index 8a6c332..6b1fb81 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -6,7 +6,7 @@ pub struct CpuState {
     pub registers: [u32; 32],
     pub f_registers: [f32; 32],
     pub pc: u32,
-    pub reservations: HashSet<usize>,
+    pub reserved_addr: Option<usize>,
 }
 
 impl CpuState {
@@ -16,18 +16,30 @@ impl CpuState {
             registers: [0u32; 32],
             f_registers: [0f32; 32],
             pc: 0,
-            reservations: HashSet::new(),
+            reserved_addr: None,
         }
     }
 
-    pub fn release_reservations(&mut self, indices: &[usize]) {
-        for index in indices {
-            self.reservations.remove(&index);
+    pub fn is_addr_reserved(&self, addr: usize) -> bool {
+        if let Some(reservation) = self.reserved_addr
+            && (reservation..reservation + 4).contains(&addr)
+        {
+            return true
         }
+        false
     }
-    pub fn acquire_reservations(&mut self, indices: &[usize]) {
-        for index in indices {
-            self.reservations.insert(*index);
+
+    pub fn release_reservations(&mut self, indices: &[usize]) {
+        if let Some(addr) = self.reserved_addr {
+            for index in indices {
+                if (addr..addr + 4).contains(index) {
+                    self.reserved_addr = None;
+                    return;
+                }
+            }
         }
     }
+    pub fn acquire_reservations(&mut self, addr: usize) {
+        self.reserved_addr = Some(addr);
+    }
 }