fix: support millisecond quote stamps in simulations

crates/simulator/src/state_override_stream.rs

@@ -41,6 +41,13 @@ use {
 /// stamp. The remaining bytes are the maker's price and must survive
 /// restamping untouched.
 const STAMP_LEN: usize = 4;
+const MILLIS_STAMP_RANGE: std::ops::Range<usize> = 25..31;
+
+#[derive(Clone, Default)]
+struct MillisecondStamps {
+    stamp: Option<u32>,
+    words: BTreeMap<(Address, B256, bool), u32>,
+}
 
 /// How many recent block gaps are kept to infer the chain's block spacing. A
 /// handful is enough to see past a slot nobody proposed, and few enough that a
@@ -57,6 +64,7 @@ struct Snapshot {
     /// Newest stamp any venue quoted for. Only words carrying it belong to a
     /// lane a maker is quoting for `block_number`.
     stamp: Option<u32>,
+    millisecond_stamps: MillisecondStamps,
     received_at: Option<Instant>,
 }
 
@@ -91,7 +99,7 @@ impl SimulationOverrides {
         let metrics = Metrics::get();
         // Holding this borrow blocks the stream task from publishing, so it is
         // released before the copy is restamped.
-        let (overrides, stamp) = {
+        let (overrides, stamp, millisecond_stamps) = {
             let snapshot = self.0.snapshots.borrow();
             let Some(received_at) = snapshot.received_at else {
                 metrics.record_override_result(OverrideResult::Empty);
@@ -109,10 +117,14 @@ impl SimulationOverrides {
                 metrics.record_override_result(OverrideResult::Empty);
                 return None;
             }
-            (snapshot.overrides.clone(), snapshot.stamp)
+            (
+                snapshot.overrides.clone(),
+                snapshot.stamp,
+                snapshot.millisecond_stamps.clone(),
+            )
         };
         metrics.record_override_result(OverrideResult::Fresh);
-        let overrides = restamp(overrides, stamp, timestamp);
+        let overrides = restamp(overrides, stamp, &millisecond_stamps, timestamp);
         Some(overrides)
     }
 }
@@ -125,7 +137,34 @@ impl SimulationOverrides {
 /// would on chain; rewriting it too would forge liveness for a price nobody is
 /// quoting. Only the stamp bytes are touched, never the price bytes next to
 /// them.
-fn restamp(mut overrides: StateOverride, stamp: Option<u32>, timestamp: u64) -> StateOverride {
+fn restamp(
+    mut overrides: StateOverride,
+    stamp: Option<u32>,
+    millisecond_stamps: &MillisecondStamps,
+    timestamp: u64,
+) -> StateOverride {
+    if let Some(millisecond_stamp) = millisecond_stamps.stamp
+        && timestamp <= u64::from(millisecond_stamp)
+    {
+        for ((address, slot, is_state_diff), quoted_at) in &millisecond_stamps.words {
+            if *quoted_at != millisecond_stamp {
+                continue;
+            }
+            let Some(account) = overrides.get_mut(address) else {
+                continue;
+            };
+            let words = if *is_state_diff {
+                account.state_diff.as_mut()
+            } else {
+                account.state.as_mut()
+            };
+            if let Some(word) = words.and_then(|words| words.get_mut(slot))
+                && !stamp.is_some_and(|stamp| word[..STAMP_LEN] == stamp.to_be_bytes())
+            {
+                word[MILLIS_STAMP_RANGE].copy_from_slice(&(timestamp * 1000).to_be_bytes()[2..]);
+            }
+        }
+    }
     let Some(stamp) = stamp else {
         return overrides;
     };
@@ -142,6 +181,11 @@ fn restamp(mut overrides: StateOverride, stamp: Option<u32>, timestamp: u64) ->
     overrides
 }
 
+fn matches_millisecond_stamp(word: &B256, stamp: u32) -> bool {
+    word[..STAMP_LEN] != stamp.to_be_bytes()
+        && word[MILLIS_STAMP_RANGE] == (u64::from(stamp) * 1000).to_be_bytes()[2..]
+}
+
 #[derive(Debug, Deserialize)]
 #[serde(rename_all = "camelCase")]
 struct Frame {
@@ -264,6 +308,7 @@ pub fn spawn_pamm_stream(cfg: &Config, blocks: CurrentBlockWatcher) -> Simulatio
         overrides: StateOverride::default(),
         block_number: 0,
         stamp: None,
+        millisecond_stamps: MillisecondStamps::default(),
         received_at: None,
     });
 
@@ -351,6 +396,7 @@ struct Quotes {
     /// Block timestamp this venue's freshly quoted lanes are stamped with,
     /// when its newest frame quoted any.
     stamp: Option<u32>,
+    millisecond_stamp: Option<u32>,
 }
 
 /// The newest frame of every venue seen so far.
@@ -373,7 +419,17 @@ impl Venues {
                 let stamp = stamp.to_be_bytes();
                 words(&overrides).any(|word| word[..STAMP_LEN] == stamp)
             });
-            self.0.insert(venue, Quotes { overrides, stamp });
+            let millisecond_stamp = quoted_at.filter(|stamp| {
+                words(&overrides).any(|word| matches_millisecond_stamp(word, *stamp))
+            });
+            self.0.insert(
+                venue,
+                Quotes {
+                    overrides,
+                    stamp,
+                    millisecond_stamp,
+                },
+            );
         }
     }
 
@@ -383,19 +439,35 @@ impl Venues {
     /// Every venue keeps its lanes in the same shared registry account and a
     /// frame only ever carries its own, so storage is merged word by word:
     /// inserting the account wholesale would drop the other venues' lanes.
-    fn fold(&self) -> (StateOverride, Option<u32>) {
+    fn fold(&self) -> (StateOverride, Option<u32>, MillisecondStamps) {
         let mut overrides = StateOverride::default();
         let mut stamp = None;
+        let mut millisecond_stamps = MillisecondStamps::default();
         for quotes in self.0.values() {
             stamp = stamp.max(quotes.stamp);
+            millisecond_stamps.stamp = millisecond_stamps.stamp.max(quotes.millisecond_stamp);
             for (account, account_override) in &quotes.overrides {
+                for (is_state_diff, words) in [
+                    (false, account_override.state.as_ref()),
+                    (true, account_override.state_diff.as_ref()),
+                ] {
+                    for (slot, word) in words.into_iter().flatten() {
+                        let key = (*account, *slot, is_state_diff);
+                        millisecond_stamps.words.remove(&key);
+                        if let Some(stamp) = quotes.millisecond_stamp
+                            && matches_millisecond_stamp(word, stamp)
+                        {
+                            millisecond_stamps.words.insert(key, stamp);
+                        }
+                    }
+                }
                 merge_account(
                     overrides.entry(*account).or_default(),
                     account_override.clone(),
                 );
             }
         }
-        (overrides, stamp)
+        (overrides, stamp, millisecond_stamps)
     }
 }
 
@@ -436,12 +508,13 @@ fn merge_words(target: &mut Option<B256Map<B256>>, update: Option<B256Map<B256>>
 }
 
 fn publish(venues: &Venues, block_number: u64, sender: &watch::Sender<Snapshot>) {
-    let (overrides, stamp) = venues.fold();
+    let (overrides, stamp, millisecond_stamps) = venues.fold();
     Metrics::get().venue_count.set(overrides.len() as i64);
     let snapshot = Snapshot {
         overrides,
         block_number,
         stamp,
+        millisecond_stamps,
         received_at: Some(Instant::now()),
     };
     if let Err(err) = sender.send(snapshot) {
@@ -705,11 +778,12 @@ mod tests {
         for frame in frames {
             venues.update(frame, Some(QUOTED_AT));
         }
-        let (overrides, stamp) = venues.fold();
+        let (overrides, stamp, millisecond_stamps) = venues.fold();
         let (_sender, receiver) = watch::channel(Snapshot {
             overrides,
             block_number,
             stamp,
+            millisecond_stamps,
             received_at: Some(Instant::now()),
         });
         handle(receiver, max_age)
@@ -725,6 +799,7 @@ mod tests {
             overrides,
             block_number,
             stamp: None,
+            millisecond_stamps: MillisecondStamps::default(),
             received_at: Some(received_at),
         }
     }
@@ -759,6 +834,7 @@ mod tests {
             overrides: StateOverride::default(),
             block_number: 100,
             stamp: None,
+            millisecond_stamps: MillisecondStamps::default(),
             received_at: Some(Instant::now()),
         });
         assert!(
@@ -834,12 +910,12 @@ mod tests {
     fn restamping_leaves_every_other_byte_untouched() {
         let mut venues = Venues::default();
         venues.update(serde_json::from_str(FERMI_FRAME).unwrap(), Some(QUOTED_AT));
-        let (overrides, stamp) = venues.fold();
+        let (overrides, stamp, millisecond_stamps) = venues.fold();
         assert_eq!(stamp, Some(QUOTED_AT));
 
         let simulated_at = QUOTED_AT - SPACING;
         let original = overrides.clone();
-        let restamped = restamp(overrides, stamp, simulated_at.into());
+        let restamped = restamp(overrides, stamp, &millisecond_stamps, simulated_at.into());
 
         // Only the stamp bytes of the registry words moved; the maker's price
         // bytes and every other account are byte-identical.
@@ -1023,7 +1099,7 @@ mod tests {
 
         let mut venues = Venues::default();
         venues.update(fermi, Some(QUOTED_AT));
-        let (overrides, stamp) = venues.fold();
+        let (overrides, stamp, _) = venues.fold();
         assert_eq!(stamp, Some(QUOTED_AT));
         assert_eq!(overrides.len(), 5);
 
@@ -1049,7 +1125,7 @@ mod tests {
 
         let mut venues = Venues::default();
         venues.update(serde_json::from_str(FERMI_FRAME).unwrap(), Some(projected));
-        let (overrides, stamp) = venues.fold();
+        let (overrides, stamp, _) = venues.fold();
         assert_eq!(stamp, Some(QUOTED_AT));
         let stamp = projected.to_be_bytes();
         assert!(