address PR comments

arbitrator/prover/src/lib.rs

@@ -28,13 +28,13 @@ use std::{
     path::Path,
     sync::{
         atomic::{self, AtomicU8},
-        Arc, RwLock,
+        Arc, Mutex,
     },
 };
 use utils::{Bytes32, CBytes};
 
 lazy_static::lazy_static! {
-    static ref CACHE_PREIMAGE_REHASH_CHECK: RwLock<LruCache<(Bytes32, PreimageType), bool>> = RwLock::new(LruCache::new(NonZeroUsize::new(1024).unwrap()));
+    static ref BLOBHASH_PREIMAGE_CACHE: Mutex<LruCache<Bytes32, CBytes>> = Mutex::new(LruCache::new(NonZeroUsize::new(12).unwrap()));
 }
 
 #[repr(C)]
@@ -296,45 +296,47 @@ pub struct ResolvedPreimage {
     pub len: isize, // negative if not found
 }
 
+macro_rules! handle_preimage_resolution {
+    ($context:expr, $ty:expr, $hash:expr, $resolver:expr) => {{
+        let res = $resolver($context, $ty.into(), $hash.as_ptr());
+        if res.len < 0 {
+            return None;
+        }
+        let data = CBytes::from_raw_parts(res.ptr, res.len as usize);
+        // Check if preimage rehashes to the provided hash
+        match crate::utils::hash_preimage(&data, $ty) {
+            Ok(have_hash) if have_hash.as_slice() == *$hash => {}
+            Ok(got_hash) => panic!(
+                "Resolved incorrect data for hash {} (rehashed to {})",
+                $hash,
+                Bytes32(got_hash),
+            ),
+            Err(err) => panic!(
+                "Failed to hash preimage from resolver (expecting hash {}): {}",
+                $hash, err,
+            ),
+        }
+        data
+    }};
+}
+
 #[no_mangle]
 pub unsafe extern "C" fn arbitrator_set_preimage_resolver(
     mach: *mut Machine,
     resolver: unsafe extern "C" fn(u64, u8, *const u8) -> ResolvedPreimage,
 ) {
     (*mach).set_preimage_resolver(Arc::new(
         move |context: u64, ty: PreimageType, hash: Bytes32| -> Option<CBytes> {
-            let res = resolver(context, ty.into(), hash.as_ptr());
-            if res.len < 0 {
-                return None;
-            }
-            let data = CBytes::from_raw_parts(res.ptr, res.len as usize);
-            // Check if preimage rehashes to the provided hash
-            let cache_key = (hash, ty);
-            let cache = CACHE_PREIMAGE_REHASH_CHECK.read().unwrap();
-            if !cache.contains(&cache_key) {
-                drop(cache);
-                match crate::utils::hash_preimage(&data, ty) {
-                    Ok(have_hash) if have_hash.as_slice() == *hash => {}
-                    Ok(got_hash) => panic!(
-                        "Resolved incorrect data for hash {} (rehashed to {})",
-                        hash,
-                        Bytes32(got_hash),
-                    ),
-                    Err(err) => panic!(
-                        "Failed to hash preimage from resolver (expecting hash {}): {}",
-                        hash, err,
-                    ),
+            if let PreimageType::EthVersionedHash = ty {
+                let mut cache = BLOBHASH_PREIMAGE_CACHE.lock().unwrap();
+                if cache.contains(&hash) {
+                    return cache.get(&hash).cloned();
                 }
-                let mut cache = CACHE_PREIMAGE_REHASH_CHECK.write().unwrap();
-                cache.put(cache_key, true);
-            } else {
-                drop(cache);
-                if let Ok(mut cache) = CACHE_PREIMAGE_REHASH_CHECK.try_write() {
-                    let _ = cache.pop(&cache_key);
-                    cache.put(cache_key, true);
-                };
+                let data = handle_preimage_resolution!(context, ty, hash, resolver);
+                cache.put(hash, data.clone());
+                return Some(data);
             }
-            Some(data)
+            Some(handle_preimage_resolution!(context, ty, hash, resolver))
         },
     ) as PreimageResolver);
 }