1 files changed, 47 insertions, 5 deletions

diff --git a/src/mem.rs b/src/mem.rs
index 0b2336e..c5e93f8 100644
--- a/src/mem.rs
+++ b/src/mem.rs
@@ -203,7 +203,7 @@ impl Ram {
         })
     }
 
-    const VERSION_CHUNK_SIZE: usize = 64;
+    pub const VERSION_CHUNK_SIZE: usize = 64;
 
     pub fn buf_mut(&mut self) -> &mut [u8] {
         self.buf.as_mut()
@@ -216,7 +216,7 @@ impl Ram {
     /// It must also be known that the contents of RAM are made up of naturally
     /// aligned valid instances of T.
     #[inline]
-    unsafe fn buf_transmuted<T>(&self) -> &[T] {
+    pub(crate) unsafe fn buf_transmuted<T>(&self) -> &[T] {
         debug_assert!(self.buf.len().is_multiple_of(std::mem::size_of::<T>()));
         unsafe {
             std::slice::from_raw_parts(
@@ -412,11 +412,38 @@ impl Ram {
         let chunk_counter = self.version_counters.get(chunk_id)?;
         Some(RamVersionClaim::claim_even(&chunk_counter))
     }
+
+    // Tries to create a claim for a specified chunk id with a specific version
+    // Outer Option represents whether the chunk id exists
+    // Inner Option represents whether the claim succeeded
+    pub fn claim_expected<'a>(
+        &'a self,
+        chunk_id: usize,
+        expected: u32,
+    ) -> Option<Option<RamVersionClaim<'a>>> {
+        self.version_counters
+            .get(chunk_id)
+            .map(|chunk_counter| RamVersionClaim::claim_expected(chunk_counter, expected))
+    }
+
+    /// Waits for a specific address to have an even (ready) version
+    /// number and returns the version chunk id and version
+    pub fn wait_for_even_version(&self, addr: u64) -> Option<(usize, u32)> {
+        let chunk_id = addr as usize / Self::VERSION_CHUNK_SIZE;
+        let chunk_counter = self.version_counters.get(chunk_id)?;
+
+        loop {
+            let current_version = chunk_counter.load(Ordering::Acquire);
+            if current_version.is_multiple_of(2) {
+                return Some((chunk_id, current_version));
+            }
+        }
+    }
 }
 
 pub struct RamVersionClaim<'a> {
     version_counter: &'a AtomicU32,
-    intial_version: u32,
+    initial_version: u32,
 }
 
 impl<'a> RamVersionClaim<'a> {
@@ -438,17 +465,32 @@ impl<'a> RamVersionClaim<'a> {
             if let Ok(initial_version) = res {
                 return RamVersionClaim {
                     version_counter: counter,
-                    intial_version: initial_version,
+                    initial_version,
                 };
             }
         }
     }
+
+    pub fn claim_expected(counter: &'a AtomicU32, expected: u32) -> Option<RamVersionClaim<'a>> {
+        counter
+            .compare_exchange(
+                expected,
+                expected.wrapping_add(1),
+                Ordering::AcqRel,
+                Ordering::Acquire,
+            )
+            .ok()
+            .map(|initial_version| RamVersionClaim {
+                version_counter: counter,
+                initial_version,
+            })
+    }
 }
 
 impl<'a> Drop for RamVersionClaim<'a> {
     fn drop(&mut self) {
         self.version_counter
-            .store(self.intial_version.wrapping_add(2), Ordering::Release);
+            .store(self.initial_version.wrapping_add(2), Ordering::Release);
     }
 }