perf(vm): key the class cache on declaration height instead of load height

vm/rust/src/state_reader/ffi.rs

@@ -22,6 +22,7 @@ extern "C" {
     pub fn JunoStateGetCompiledClass(
         reader_handle: usize,
         class_hash: *const c_uchar,
+        declared_at: *mut u64,
     ) -> *const c_char;
     pub fn JunoStateGetCompiledClassHash(
         reader_handle: usize,

vm/rust/src/state_reader/state_reader.rs

@@ -26,7 +26,7 @@ use crate::{
 
 struct CachedRunnableCompiledClass {
     pub definition: RunnableCompiledClass,
-    pub cached_on_height: u64,
+    pub declared_at: u64,
 }
 
 static CLASS_CACHE: Lazy<Mutex<SizedCache<ClassHash, CachedRunnableCompiledClass>>> =
@@ -135,25 +135,19 @@ impl StateReader for JunoStateReader {
     /// Returns the contract class of the given class hash.
     fn get_compiled_class(&self, class_hash: ClassHash) -> StateResult<RunnableCompiledClass> {
         if let Some(cached_class) = CLASS_CACHE.lock().unwrap().cache_get(&class_hash) {
-            // skip the cache if it comes from a height higher than ours. Class might be undefined on the height
-            // that we are reading from right now.
-            //
-            // About the changes in the second attempt at making class cache behave as expected;
-            //
-            // The initial assumption here was that `self.height` uniquely identifies and strictly orders the underlying state
-            // instances. The first assumption doesn't necessarily hold, because we can pass different state instaces with the
-            // same height. This most commonly happens with call/estimate/simulate and trace flows. Trace flow calls the VM
-            // for block number N with the state at the beginning of the block, while call/estimate/simulate flows call the VM
-            // with the same block number but with the state at the end of that block. That is why, we cannot use classes from cache
-            // if they are cached on the same height that we are executing on. Because they might be cached using a state instance that
-            // is in the future compared to the state that we are currently executing on, even tho they have the same height.
-            if self.height.is_after(cached_class.cached_on_height) {
+            // A class declared at height D exists in every state from the end of block D on. Trace
+            // and call flows can run block N on different states (start vs end of N), so a height
+            // equal to D is not a hit.
+            if self.height.is_after(cached_class.declared_at) {
                 return Ok(cached_class.definition.clone());
             }
         }
 
         let class_hash_bytes = felt_to_byte_array(&class_hash.0);
-        let ptr = unsafe { JunoStateGetCompiledClass(self.handle, class_hash_bytes.as_ptr()) };
+        let mut declared_at: u64 = 0;
+        let ptr = unsafe {
+            JunoStateGetCompiledClass(self.handle, class_hash_bytes.as_ptr(), &mut declared_at)
+        };
         if ptr.is_null() {
             Err(StateError::UndeclaredClassHash(class_hash))
         } else {
@@ -166,14 +160,15 @@ impl StateReader for JunoStateReader {
                 Ok(class) => {
                     let runnable_compiled_class =
                         RunnableCompiledClass::try_from(class.contract_class).unwrap();
-                    if let BlockHeight::Height(height) = self.height {
+                    // Pending states report declared_at 0 for classes declared in the pending block.
+                    if matches!(self.height, BlockHeight::Height(_)) {
                         CLASS_CACHE.lock().unwrap().cache_set(
                             class_hash,
                             CachedRunnableCompiledClass {
                                 // This clone is cheap, it is just a reference copy in the underlying
                                 // RunnableCompiledClass implementation
                                 definition: runnable_compiled_class.clone(),
-                                cached_on_height: height,
+                                declared_at,
                             },
                         );
                     }

vm/state.go

@@ -71,7 +71,11 @@ func JunoStateGetClassHashAt(readerHandle C.uintptr_t, contractAddress, buffer u
 }
 
 //export JunoStateGetCompiledClass
-func JunoStateGetCompiledClass(readerHandle C.uintptr_t, classHash unsafe.Pointer) unsafe.Pointer {
+func JunoStateGetCompiledClass(
+	readerHandle C.uintptr_t,
+	classHash unsafe.Pointer,
+	declaredAt *C.uint64_t,
+) unsafe.Pointer {
 	context := unwrapContext(readerHandle)
 
 	classHashFelt := makeFeltFromPtr(classHash)
@@ -89,6 +93,7 @@ func JunoStateGetCompiledClass(readerHandle C.uintptr_t, classHash unsafe.Pointe
 		return nil
 	}
 
+	*declaredAt = C.uint64_t(val.At)
 	return unsafe.Pointer(cstring(compiledClass))
 }
 

vm/vm_test.go

@@ -440,3 +440,83 @@ func NewState(
 	stateDB := state.NewStateDB(testDB, triedb)
 	return state.New(stateRoot, stateDB, batch)
 }
+
+// classFetchCounter counts how often the VM asks Go for a class definition.
+type classFetchCounter struct {
+	core.StateReader
+	fetches int
+}
+
+func (c *classFetchCounter) Class(classHash *felt.Felt) (*core.DeclaredClassDefinition, error) {
+	c.fetches++
+	return c.StateReader.Class(classHash)
+}
+
+func TestClassCacheKeyedOnDeclarationHeight(t *testing.T) {
+	testDB := memory.New()
+	batch := testDB.NewBatch()
+	client := feeder.NewTestClient(t, &networks.Mainnet)
+	gw := adaptfeeder.New(client)
+
+	contractAddr := felt.NewUnsafeFromString[felt.Felt]("0xDEADBEEF")
+	// https://voyager.online/class/0x03297a93c52357144b7da71296d7e8231c3e0959f0a1d37222204f2f7712010e
+	classHash := felt.NewUnsafeFromString[felt.Felt](
+		"0x3297a93c52357144b7da71296d7e8231c3e0959f0a1d37222204f2f7712010e",
+	)
+	simpleClass, err := gw.Class(t.Context(), classHash)
+	require.NoError(t, err)
+
+	testState, err := NewState(t, &felt.Zero, testDB, batch)
+	require.NoError(t, err)
+	require.NoError(t, testState.Update(&core.Header{Number: 0}, &core.StateUpdate{
+		OldRoot: &felt.Zero,
+		NewRoot: felt.NewUnsafeFromString[felt.Felt](
+			"0x3d452fbb3c3a32fe85b1a3fbbcdec316d5fc940cefc028ee808ad25a15991c8",
+		),
+		StateDiff: &core.StateDiff{
+			DeployedContracts: map[felt.Felt]*felt.Felt{
+				*contractAddr: classHash,
+			},
+		},
+	}, map[felt.Felt]core.ClassDefinition{
+		*classHash: simpleClass,
+	}, false))
+	require.NoError(t, batch.Write())
+
+	entryPoint := felt.NewUnsafeFromString[felt.Felt](
+		"0x39e11d48192e4333233c7eb19d10ad67c362bb28580c604d67884c85da39695",
+	)
+	chainInfo := ChainInfo{
+		ChainID:           networks.Mainnet.L2ChainID,
+		FeeTokenAddresses: networks.DefaultFeeTokenAddresses,
+	}
+	counter := &classFetchCounter{StateReader: testState}
+	callAt := func(height uint64) {
+		_, err := New(&chainInfo, false, nil).Call(
+			&CallInfo{
+				ContractAddress: contractAddr,
+				ClassHash:       classHash,
+				Selector:        entryPoint,
+			},
+			// A non-nil hash marks the block as not pending, which enables the class cache.
+			&BlockInfo{Header: &core.Header{Number: height, Hash: &felt.One}},
+			counter,
+			DefaultMaxSteps,
+			DefaultMaxGas,
+			false,
+			false,
+		)
+		require.NoError(t, err)
+	}
+
+	callAt(10)
+	fetchesAfterFirstCall := counter.fetches
+
+	callAt(5)
+	assert.Equal(t, fetchesAfterFirstCall, counter.fetches,
+		"a class declared at block 0 must be served from the cache at any later height")
+
+	callAt(0)
+	assert.Equal(t, fetchesAfterFirstCall+1, counter.fetches,
+		"the declaration height itself must not be served from the cache")
+}