From b464b4185c0e5399076c5e3beb045e34b1e6c362 Mon Sep 17 00:00:00 2001
From: clabby <ben@clab.by>
Date: Sat, 17 Feb 2024 18:18:12 -0500
Subject: [PATCH] =?UTF-8?q?=E2=99=BB=20Clean=20up=20`OracleReader`?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 crates/preimage/src/key.rs    |  2 +-
 crates/preimage/src/oracle.rs | 69 ++++++++---------------------------
 2 files changed, 16 insertions(+), 55 deletions(-)

diff --git a/crates/preimage/src/key.rs b/crates/preimage/src/key.rs
index c20a3b7a6..d297d2b01 100644
--- a/crates/preimage/src/key.rs
+++ b/crates/preimage/src/key.rs
@@ -63,7 +63,7 @@ mod test {
     use super::*;
 
     #[test]
-    fn test_local_key() {
+    fn test_preimage_keys() {
         let types = [PreimageKeyType::Local, PreimageKeyType::Keccak256];
 
         for key_type in types {
diff --git a/crates/preimage/src/oracle.rs b/crates/preimage/src/oracle.rs
index ccb9cc002..f14df03cb 100644
--- a/crates/preimage/src/oracle.rs
+++ b/crates/preimage/src/oracle.rs
@@ -6,85 +6,48 @@ use anyhow::{bail, Result};
 #[derive(Debug, Clone, Copy)]
 pub struct OracleReader {
     pipe_handle: PipeHandle,
-    key: Option<PreimageKey>,
-    length: usize,
-    cursor: usize,
 }
 
 impl OracleReader {
     /// Create a new [OracleReader] from a [PipeHandle].
     pub fn new(pipe_handle: PipeHandle) -> Self {
-        Self {
-            pipe_handle,
-            key: None,
-            length: 0,
-            cursor: 0,
-        }
-    }
-
-    /// Return the current key stored in the global oracle reader
-    pub fn key(&self) -> Option<PreimageKey> {
-        self.key
-    }
-
-    /// Return the length of the current pre-image
-    pub fn length(&self) -> usize {
-        self.length
-    }
-
-    /// Current position of the read cursor within the current pre-image
-    pub fn cursor(&self) -> usize {
-        self.cursor
+        Self { pipe_handle }
     }
 
-    /// Get the data corresponding to the currently set key from the host. Return the data in a new heap allocated `Vec<u8>`
-    ///
-    /// Internally this reads self.length bytes from the ReadPreimage file descriptor into a new heap allocated `Vec<u8>` and returns it.
-    /// This is a high level way to interact with the preimage oracle but may not be the best way if heap allocations are not desirable.
+    /// Get the data corresponding to the currently set key from the host. Return the data in a new heap allocated
+    /// `Vec<u8>`
     pub fn get(&mut self, key: PreimageKey) -> Result<Vec<u8>> {
-        self.set_key(key)?;
-        let mut data_buffer = alloc::vec![0; self.length];
+        let length = self.write_key(key)?;
+        let mut data_buffer = alloc::vec![0; length];
 
         // Grab a read lock on the preimage pipe to read the data.
-        self.cursor += self.pipe_handle.read_exact(&mut data_buffer)? as usize;
+        self.pipe_handle.read_exact(&mut data_buffer)? as usize;
 
         Ok(data_buffer)
     }
 
     /// Get the data corresponding to the currently set key from the host. Write the data into the provided buffer
-    ///
-    /// # Panics
-    /// This will panic if the size of the buffer is not equal to the size of the preimage as reported by the host
     pub fn get_exact(&mut self, key: PreimageKey, buf: &mut [u8]) -> Result<()> {
-        self.set_key(key)?;
+        // Write the key to the host and read the length of the preimage.
+        let length = self.write_key(key)?;
 
         // Ensure the buffer is the correct size.
-        if buf.len() != self.length {
+        if buf.len() != length {
             bail!(
                 "Buffer size {} does not match preimage size {}",
                 buf.len(),
-                self.length
+                length
             );
         }
 
-        // Grab a read lock on the preimage pipe to read the data.
-        self.cursor += self.pipe_handle.read_exact(buf)? as usize;
+        self.pipe_handle.read_exact(buf)?;
 
         Ok(())
     }
 
-    /// Set the preimage key for the global oracle reader. This will overwrite any existing key, and block until all
-    /// data has been read from the host.
-    ///
-    /// Internally this sends the 32 bytes of the key to the host by writing into the WritePreimage file descriptor.
-    /// This may require several writes as the host may only accept a few bytes at a time. Once 32 bytes have been written
-    /// successfully the key is considered set. If it fails to write 32 bytes it will return an error.
-    /// Once it has written the key it will read the first 8 bytes of the ReadPreimage file descriptor which is the length
-    /// encoded as a big endian u64. This is stored in the oracle reader along with the read cursor position.
-    fn set_key(&mut self, key: PreimageKey) -> Result<()> {
-        // Set the active key.
-        self.key = Some(key);
-
+    /// Set the preimage key for the global oracle reader. This will overwrite any existing key, and block until the 
+    /// host has prepared the preimage and responded with the length of the preimage.
+    fn write_key(&mut self, key: PreimageKey) -> Result<usize> {
         // Write the key to the host so that it can prepare the preimage.
         let key_bytes: [u8; 32] = key.into();
         self.pipe_handle.write(&key_bytes)?;
@@ -92,9 +55,7 @@ impl OracleReader {
         // Read the length prefix and reset the cursor.
         let mut length_buffer = [0u8; 8];
         self.pipe_handle.read_exact(&mut length_buffer)?;
-        self.length = u64::from_be_bytes(length_buffer) as usize;
-        self.cursor = 0;
-        Ok(())
+        Ok(u64::from_be_bytes(length_buffer) as usize)
     }
 }