handle reorgs

arbitrator/stylus/src/cache.rs

@@ -1,22 +1,21 @@
 // Copyright 2022-2024, Offchain Labs, Inc.
 // For license information, see https://github.com/OffchainLabs/nitro/blob/master/LICENSE
 
-use std::{
-    collections::{hash_map::Entry, HashMap},
-    num::NonZeroUsize,
-};
-
 use arbutil::Bytes32;
 use eyre::Result;
 use lazy_static::lazy_static;
 use lru::LruCache;
 use parking_lot::Mutex;
 use prover::programs::config::CompileConfig;
 use rand::Rng;
+use std::{
+    collections::{hash_map::Entry, HashMap},
+    num::NonZeroUsize,
+};
 use wasmer::{Engine, Module, Store};
 
 lazy_static! {
-    static ref INIT_CACHE: Mutex<InitCache> = Mutex::new(InitCache::new(64));
+    static ref INIT_CACHE: Mutex<InitCache> = Mutex::new(InitCache::new(256));
 }
 
 macro_rules! cache {
@@ -110,7 +109,7 @@ impl InitCache {
         None
     }
 
-    /// Inserts an item into the long term cache, stealing from the LRU one if able.
+    /// Inserts an item into the long term cache, stealing from the LRU cache if able.
     pub fn insert(
         module_hash: Bytes32,
         module: &[u8],
@@ -157,4 +156,13 @@ impl InitCache {
         let key = CacheKey::new(module_hash, version, debug);
         cache!().arbos.remove(&key);
     }
+
+    /// Modifies the cache for reorg, dropping the long-term cache.
+    pub fn reorg(_block: u64) {
+        let mut cache = cache!();
+        let cache = &mut *cache;
+        for (key, item) in cache.arbos.drain() {
+            cache.lru.put(key, item); // not all will fit, just a heuristic
+        }
+    }
 }

arbitrator/stylus/src/lib.rs

@@ -227,15 +227,17 @@ pub unsafe extern "C" fn stylus_cache_module(
 }
 
 /// Evicts an activated user program from the init cache.
-///
-/// # Safety
-///
-/// `module` must represent a valid module produced from `stylus_activate`.
 #[no_mangle]
-pub unsafe extern "C" fn stylus_evict_module(module_hash: Bytes32, version: u16, debug: bool) {
+pub extern "C" fn stylus_evict_module(module_hash: Bytes32, version: u16, debug: bool) {
     InitCache::evict(module_hash, version, debug);
 }
 
+/// Reorgs the init cache. This will likely never happen.
+#[no_mangle]
+pub extern "C" fn stylus_reorg_vm(block: u64) {
+    InitCache::reorg(block);
+}
+
 /// Frees the vector. Does nothing when the vector is null.
 ///
 /// # Safety

execution/gethexec/executionengine.go

@@ -1,5 +1,14 @@
+// Copyright 2022-2024, Offchain Labs, Inc.
+// For license information, see https://github.com/OffchainLabs/nitro/blob/master/LICENSE
+
 package gethexec
 
+/*
+#cgo CFLAGS: -g -Wall -I../../target/include/
+#cgo LDFLAGS: ${SRCDIR}/../../target/lib/libstylus.a -ldl -lm
+#include "arbitrator.h"
+*/
+import "C"
 import (
 	"context"
 	"encoding/binary"
@@ -128,6 +137,9 @@ func (s *ExecutionEngine) Reorg(count arbutil.MessageIndex, newMessages []arbost
 		return nil
 	}
 
+	// reorg Rust-side VM state
+	C.stylus_reorg_vm(C.uint64_t(blockNum))
+
 	err := s.bc.ReorgToOldBlock(targetBlock)
 	if err != nil {
 		return err