From f10e66835ae5ba9298786acd9510297be4c18ad2 Mon Sep 17 00:00:00 2001
From: Nathan Lee <me@nwlee.tech>
Date: Sat, 6 Jun 2026 14:44:55 -0500
Subject: implement slightly better atomic semantics

---
 src/main.rs | 57 +++++++++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 43 insertions(+), 14 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index a351d48..4ccd0a1 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,12 +1,13 @@
 pub mod core;
 
+use crate::core::*;
 use elf::ElfBytes;
 use elf::endian::AnyEndian;
+use std::collections::HashSet;
 use std::io;
 use std::io::Write;
 use std::process::exit;
 use std::sync::{LazyLock, Mutex};
-use crate::core::*;
 
 const MEMORY_SIZE: usize = 64 * 1024 * 1024;
 const REGISTER_NAME_WIDTH: u32 = 0b11111;
@@ -17,7 +18,7 @@ struct CpuState<'a> {
     memory: &'a mut Box<[u8]>,
     registers: &'a mut [u32; 32],
     pc: &'a mut u32,
-    reservations: &'a mut Vec<usize>,
+    reservations: &'a mut HashSet<usize>,
 }
 
 fn main() {
@@ -43,7 +44,7 @@ fn main() {
 
     let mut registers: [u32; 32] = [0; 32];
     let mut pc: u32 = 0;
-    let mut reservations: Vec<usize> = vec![];
+    let mut reservations: HashSet<usize> = HashSet::new();
 
     let str_tab = stab.expect("Cannot read strtab");
     if let Some(headers) = sh {
@@ -70,7 +71,14 @@ fn main() {
     }
 }
 
-fn handle_instruction(CpuState { memory, registers, pc, reservations }: CpuState) {
+fn handle_instruction(
+    CpuState {
+        memory,
+        registers,
+        pc,
+        reservations,
+    }: CpuState,
+) {
     let word = u32::from_le_bytes([
         memory[*pc as usize],
         memory[*pc as usize + 1],
@@ -167,10 +175,21 @@ fn handle_instruction(CpuState { memory, registers, pc, reservations }: CpuState
 
             let imm = ((((imm_high << 5) | imm_low) as i32) << 20 >> 20) as u32;
 
+            let addr = (registers[rs1] + imm) as usize;
+
             match f3 {
-                0 => sb(registers, memory, rs1, rs2, imm), // sb
-                1 => sh(registers, memory, rs1, rs2, imm),
-                2 => sw(registers, memory, rs1, rs2, imm),
+                0 => {
+                    sb(registers, memory, rs1, rs2, imm);
+                    release_reservations(reservations, &[addr]);
+                }
+                1 => {
+                    sh(registers, memory, rs1, rs2, imm);
+                    release_reservations(reservations, &[addr, addr + 1]);
+                }
+                2 => {
+                    sw(registers, memory, rs1, rs2, imm);
+                    release_reservations(reservations, &[addr, addr + 1, addr + 2, addr + 3]);
+                }
                 _ => panic!("Illegal instruction: {:x}", word),
             };
         }
@@ -199,11 +218,10 @@ fn handle_instruction(CpuState { memory, registers, pc, reservations }: CpuState
         }
         0b1101111 => {
             let rd = (word >> 7) & REGISTER_NAME_WIDTH;
-            let imm =
-                (((word >> 31) & 0x1) << 20) |
-                    (((word >> 12) & 0xFF) << 12) |
-                    (((word >> 20) & 0x1) << 11) |
-                    (((word >> 21) & 0x3FF) << 1);
+            let imm = (((word >> 31) & 0x1) << 20)
+                | (((word >> 12) & 0xFF) << 12)
+                | (((word >> 20) & 0x1) << 11)
+                | (((word >> 21) & 0x3FF) << 1);
             let imm = ((imm as i32) << 11 >> 11) as u32;
             jal(registers, pc, rd as usize, imm);
             return;
@@ -265,7 +283,7 @@ fn handle_instruction(CpuState { memory, registers, pc, reservations }: CpuState
                     let addr = registers[rs1] as usize;
                     let number: u32 = load_word(&memory, addr);
                     registers[rd] = number;
-                    (*reservations).push(addr);
+                    acquire_reservations(reservations, &[addr, addr + 1, addr + 2, addr + 3]);
                 } // lr.w
                 3 => {
                     let addr = registers[rs1] as usize;
@@ -340,10 +358,21 @@ fn handle_instruction(CpuState { memory, registers, pc, reservations }: CpuState
     *pc += 4;
 }
 
+fn release_reservations(reservations: &mut HashSet<usize>, indices: &[usize]) {
+    for index in indices {
+        reservations.remove(&index);
+    }
+}
+fn acquire_reservations(reservations: &mut HashSet<usize>, indices: &[usize]) {
+    for index in indices {
+        reservations.insert(*index);
+    }
+}
+
 fn load_word(memory: &&mut Box<[u8]>, addr: usize) -> u32 {
     let slice = memory.get(addr..addr + 4).unwrap();
     u32::from_le_bytes(slice.try_into().unwrap())
 }
 fn store_word(memory: &mut Box<[u8]>, addr: usize, value: u32) {
     memory[addr..addr + 4].copy_from_slice(&value.to_le_bytes());
-}
\ No newline at end of file
+}
-- 
cgit v1.2.3