Calculate the state delta in Rust

Port `_get_new_state_after_events`, `_calculate_state_delta`,
`_is_server_still_joined` and the two steps that call them with the
extremity calculations, `_calculate_new_forward_extremities_and_state_delta`
and `_process_state_dag_forward_extremities_and_state_delta`, to
`rust/src/persist/state_delta.rs`. Reads go through a `StateDeltaReads`
trait on top of `ExtremityReads`, and state resolution through a
`StateResolver`, so the unit tests can fake both.

`NativeStateResolver` checks and bumps the input state groups as
`check_state_groups_and_bump_deletion` does, then resolves with the Rust
state resolution on the native pool, through `resolve_on_pool`, which is
split out of `StateResolutionHandler::start_converted`. The persister
shares the handler's conflict cache. There is no `_state_cache` of whole
resolutions. Rooms on state resolution v1 or without a chain cover index
fail with `StateResFallback`, for the caller to run the Python step.

`NewEvent` gains the fields the state delta reads from an event and its
context. `PersistReaderForTest` gains the five calculations, and
`test_persist_reader.py` compares each with the Python method on events
built by the event creation handler, including a resolved fork and an
MSC4242 room.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UHKYPSc37kWqGACn9aKxKL

rust/src/handlers/mod.rs

@@ -81,6 +81,11 @@ impl RustHandlers {
         let (state_resolution, sync, persister) =
             if connection_pool.is_instance(&rust_connection_pool_class)? {
                 let connection_pool = connection_pool.unbind();
+                let state_resolution = StateResolutionHandler::new(
+                    connection_pool.clone_ref(py),
+                    runtime.clone(),
+                    Arc::clone(&store),
+                );
                 let persister = if config.experimental.rust_event_persistence != "off" {
                     Persister::for_homeserver(
                         py,
@@ -88,6 +93,7 @@ impl RustHandlers {
                         connection_pool.clone_ref(py),
                         runtime.clone(),
                         Arc::clone(&store),
+                        state_resolution.conflict_cache(),
                     )?
                     .map(|persister| Py::new(py, persister))
                     .transpose()?
@@ -95,14 +101,7 @@ impl RustHandlers {
                     None
                 };
                 (
-                    Some(Py::new(
-                        py,
-                        StateResolutionHandler::new(
-                            connection_pool.clone_ref(py),
-                            runtime.clone(),
-                            Arc::clone(&store),
-                        ),
-                    )?),
+                    Some(Py::new(py, state_resolution)?),
                     Some(Py::new(
                         py,
                         SyncHandler::new(connection_pool, runtime.clone(), store, &config),

rust/src/persist/context.rs

@@ -42,7 +42,8 @@ pub struct PersistContext {
     pub state_delta_due_to_event: Option<StateMap>,
     /// Deltas between pairs of state groups `(prev, current)` known when the
     /// context was built, which can save a state resolution.
-    pub state_group_deltas: HashMap<(i64, i64), StateMap>,
+    /// Shared, since every event of a chunk carries a copy.
+    pub state_group_deltas: Arc<HashMap<(i64, i64), StateMap>>,
     /// Whether the state at the event is partial.
     pub partial_state: bool,
     /// Whether an application service sent the event. Only used to label
@@ -79,7 +80,7 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PersistContext {
                 .getattr(intern!(py, "state_group_before_event"))?
                 .extract()?,
             state_delta_due_to_event,
-            state_group_deltas,
+            state_group_deltas: Arc::new(state_group_deltas),
             partial_state: context.getattr(intern!(py, "partial_state"))?.extract()?,
             app_service: !context.getattr(intern!(py, "app_service"))?.is_none(),
         })

rust/src/persist/extremities.rs

@@ -31,49 +31,94 @@ use std::{
 use anyhow::{anyhow, bail};
 use async_trait::async_trait;
 use log::{error, info};
+use serde_json::Value;
 
 use super::{context::PersistContext, metrics::PersistCounts, reader::PersistReader};
-use crate::{events::Event, storage::events::RedactBehaviour};
+use crate::{events::Event, state::StateMap, storage::events::RedactBehaviour};
 
 /// `EventTypes.Dummy`.
 const DUMMY: &str = "org.matrix.dummy_event";
+/// `EventTypes.Create`.
+const CREATE: &str = "m.room.create";
 
-/// What the extremity calculations read from an event being persisted.
+/// What the pre-transaction computation reads from an event being
+/// persisted and its context.
 #[derive(Clone, Debug, Default)]
 pub struct NewEvent {
     pub event_id: Arc<str>,
+    pub event_type: Arc<str>,
+    /// The state key, if it is a state event. In an MSC4242 room, that is
+    /// what puts it in the state DAG.
+    pub state_key: Option<Arc<str>>,
     pub sender: Arc<str>,
     pub depth: i64,
     pub prev_event_ids: Vec<Arc<str>>,
     /// The MSC4242 `prev_state_events`, if the event's format has them.
     pub prev_state_events: Option<Vec<Arc<str>>>,
-    /// Whether the event has a state key, which in an MSC4242 room is
-    /// what puts it in the state DAG.
-    pub is_state: bool,
+    /// `content.membership`, if it is a string.
+    pub membership: Option<Arc<str>>,
+    /// For a create event, `content.room_version` as
+    /// `_get_new_state_after_events` reads it: `"1"` if it is missing, and
+    /// `None` if it is falsy in Python's terms, which sends the reader to
+    /// the database. A value that is not a string fails to parse as a
+    /// room version, as Python's `KNOWN_ROOM_VERSIONS` lookup fails for it.
+    pub create_room_version: Option<Arc<str>>,
     pub outlier: bool,
     pub soft_failed: bool,
     pub rejected: bool,
+    /// The context's state group, the state after the event. `None` for an
+    /// outlier.
+    pub state_group: Option<i64>,
+    /// The context's `state_group_deltas`.
+    pub state_group_deltas: Arc<HashMap<(i64, i64), StateMap>>,
 }
 
 impl NewEvent {
     pub fn new(event: &Event, context: &PersistContext) -> anyhow::Result<Self> {
         let metadata = event.internal_metadata();
+        let content = event.content();
+        let membership = match content.get_field("membership") {
+            Some(Value::String(membership)) => Some(Arc::from(membership.as_str())),
+            _ => None,
+        };
+        let create_room_version = if event.type_() == CREATE && event.state_key() == Some("") {
+            match content.get_field("room_version") {
+                None => Some(Arc::from("1")),
+                Some(Value::String(version)) if version.is_empty() => None,
+                Some(Value::String(version)) => Some(Arc::from(version.as_str())),
+                Some(value) if is_falsy(value) => None,
+                // Prefixed so that a number such as `10` is not read as
+                // the room version `"10"`.
+                Some(value) => Some(Arc::from(format!("not a string: {value}"))),
+            }
+        } else {
+            None
+        };
         Ok(NewEvent {
             event_id: Arc::from(event.event_id()),
+            event_type: Arc::from(event.type_()),
+            state_key: event.state_key().map(Arc::from),
             sender: Arc::from(event.sender()),
             depth: event.depth(),
             prev_event_ids: event.prev_event_ids().into_iter().map(Arc::from).collect(),
             prev_state_events: event
                 .prev_state_event_ids()
                 .map(|ids| ids.iter().map(|id| Arc::from(id.as_str())).collect()),
-            is_state: event.state_key().is_some(),
+            membership,
+            create_room_version,
             outlier: metadata.is_outlier()?,
             soft_failed: metadata.soft_failed()?,
             rejected: context.is_rejected(),
+            state_group: context.state_group,
+            state_group_deltas: Arc::clone(&context.state_group_deltas),
         })
     }
 
-    /// The `prev_state_events`, which an event in the state DAG must have.
+    pub fn is_state(&self) -> bool {
+        self.state_key.is_some()
+    }
+
+    /// The prev state events, which an event in the state DAG must have.
     fn prev_state(&self) -> anyhow::Result<&[Arc<str>]> {
         self.prev_state_events
             .as_deref()
@@ -81,6 +126,18 @@ impl NewEvent {
     }
 }
 
+/// Whether Python would read a JSON value as false.
+fn is_falsy(value: &Value) -> bool {
+    match value {
+        Value::Null => true,
+        Value::Bool(b) => !b,
+        Value::Number(n) => n.as_f64() == Some(0.0),
+        Value::String(s) => s.is_empty(),
+        Value::Array(a) => a.is_empty(),
+        Value::Object(o) => o.is_empty(),
+    }
+}
+
 /// What `_prune_extremities` reads from an extremity it might drop, and
 /// from the prev events of a local dummy event.
 #[derive(Clone, Debug, Default)]
@@ -245,7 +302,7 @@ pub async fn calculate_new_state_dag_extremities(
     events: &[NewEvent],
     counts: &mut PersistCounts,
 ) -> anyhow::Result<HashSet<Arc<str>>> {
-    let state_events: Vec<&NewEvent> = events.iter().filter(|event| event.is_state).collect();
+    let state_events: Vec<&NewEvent> = events.iter().filter(|event| event.is_state()).collect();
     if state_events.is_empty() {
         return Ok(existing.clone());
     }
@@ -434,7 +491,7 @@ fn domain_from_id(id: &str) -> anyhow::Result<&str> {
 }
 
 /// `HomeServer.is_mine_id`: false for a malformed ID.
-fn is_mine_id(id: &str, server_name: &str) -> bool {
+pub(super) fn is_mine_id(id: &str, server_name: &str) -> bool {
     id.split_once(':')
         .is_some_and(|(_, domain)| domain == server_name)
 }
@@ -556,7 +613,7 @@ mod tests {
 
     fn state_event(id: &str, prev_state: &[&str]) -> NewEvent {
         NewEvent {
-            is_state: true,
+            state_key: Some(Arc::from("")),
             prev_state_events: Some(prev_state.iter().map(|p| Arc::from(*p)).collect()),
             ..new_event(id, &[])
         }

rust/src/persist/mod.rs

@@ -25,6 +25,7 @@
 //! - [`metrics`]: counts for the Python metrics, flushed when scraped.
 //! - [`python`]: the [`python::Persister`] and test helpers.
 //! - [`reader`]: what persisting reads before its write transaction.
+//! - [`state_delta`]: the room's new current state after a chunk.
 //!
 //! Every GIL hold in this tree goes through [`crate::gil::with_gil`] with a
 //! site named `persist_*`, so that tests can pin the count per site.
@@ -37,6 +38,7 @@ pub mod extremities;
 pub mod metrics;
 pub mod python;
 pub mod reader;
+pub mod state_delta;
 
 /// Called when registering modules with python.
 pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {

rust/src/persist/python.rs

@@ -29,6 +29,7 @@ use crate::{
     events::Event,
     homeserver::HomeServer,
     runtime::RustRuntime,
+    state::ConflictCache,
     storage::{
         native_pool::NativePool,
         store::Store,
@@ -64,6 +65,8 @@ pub(super) struct PersisterInner {
     pub server_name: String,
     pub instance_name: String,
     pub metrics: PersistMetrics,
+    /// The Rust state resolution handler's conflict cache.
+    pub conflict_cache: Arc<ConflictCache>,
 }
 
 impl Persister {
@@ -76,6 +79,7 @@ impl Persister {
         connection_pool: Py<PyAny>,
         runtime: RustRuntime,
         store: Arc<Store>,
+        conflict_cache: Arc<ConflictCache>,
     ) -> PyResult<Option<Persister>> {
         let instance_name = homeserver.instance_name(py)?;
         if !homeserver.event_writers(py)?.contains(&instance_name) {
@@ -112,6 +116,7 @@ impl Persister {
                 server_name: homeserver.hostname(py)?,
                 instance_name,
                 metrics: PersistMetrics::default(),
+                conflict_cache,
             }),
         }))
     }
@@ -178,7 +183,7 @@ pub fn persist_context_for_test<'py>(
             .transpose()?,
     )?;
     let deltas = PyDict::new(py);
-    for (key, delta) in &context.state_group_deltas {
+    for (key, delta) in context.state_group_deltas.iter() {
         deltas.set_item(key, state_map_to_py(py, delta)?)?;
     }
     dict.set_item("state_group_deltas", deltas)?;

rust/src/persist/reader.rs

@@ -50,13 +50,17 @@ use super::{
     extremities::{self, NewEvent},
     metrics::PersistCounts,
     python::persist_error_to_py,
+    state_delta::{self, DeltaState, NativeStateResolver, StateAndExtremities},
 };
 use crate::database::postgres::pool::{ConnectionPool, PyConnectionPool};
 use crate::events::Event;
+use crate::state::{ConflictCache, EventMap, StateKey, StateMap, StateResEvent};
+use crate::storage::caches::StoreCaches;
 use crate::storage::events::Querier;
 use crate::storage::sync_store::{
     PersistedEventPosition, SyncReader, SyncStore, SELECT_MANY_BATCH_SIZE,
 };
+use crate::types::state_filter::state_map_to_py;
 
 /// How many events `_get_events_which_are_prevs` and
 /// `_get_prevs_before_rejected` handle per transaction.
@@ -692,6 +696,8 @@ pub struct PyPersistReaderForTest {
     store: Arc<SyncStore>,
     pool: ConnectionPool,
     handle: Handle,
+    caches: Arc<StoreCaches>,
+    conflict_cache: Arc<ConflictCache>,
 }
 
 impl PyPersistReaderForTest {
@@ -707,6 +713,62 @@ impl PyPersistReaderForTest {
         })
         .map_err(persist_error_to_py)
     }
+
+    /// Run one of the two orchestrating steps.
+    fn run_orchestration(
+        &self,
+        py: Python<'_>,
+        room_id: String,
+        server_name: String,
+        events_and_contexts: Vec<(Py<Event>, PersistContext)>,
+        now_ms: i64,
+        state_dag: bool,
+    ) -> PyResult<(StateAndExtremities, PersistCounts)> {
+        let (events, event_map) = chunk(py, &events_and_contexts)?;
+        let resolver = self.resolver();
+        self.run(py, |reader| {
+            async move {
+                let mut counts = PersistCounts::default();
+                let result = if state_dag {
+                    state_delta::process_state_dag_forward_extremities_and_state_delta(
+                        &reader,
+                        &resolver,
+                        &room_id,
+                        &server_name,
+                        &events,
+                        &event_map,
+                        now_ms,
+                        &mut counts,
+                    )
+                    .await?
+                } else {
+                    state_delta::calculate_new_forward_extremities_and_state_delta(
+                        &reader,
+                        &resolver,
+                        &room_id,
+                        &server_name,
+                        &events,
+                        &event_map,
+                        now_ms,
+                        &mut counts,
+                    )
+                    .await?
+                };
+                Ok((result, counts))
+            }
+            .boxed()
+        })
+    }
+
+    /// A state resolver over the same pool.
+    fn resolver(&self) -> NativeStateResolver {
+        NativeStateResolver {
+            pool: self.pool.clone(),
+            handle: self.handle.clone(),
+            caches: Arc::clone(&self.caches),
+            conflict_cache: Arc::clone(&self.conflict_cache),
+        }
+    }
 }
 
 #[pymethods]
@@ -718,6 +780,8 @@ impl PyPersistReaderForTest {
             store: Arc::clone(&persister.inner.sync),
             pool: pool.pool.clone(),
             handle: pool.handle.clone(),
+            caches: Arc::clone(&persister.inner.store.caches),
+            conflict_cache: Arc::clone(&persister.inner.conflict_cache),
         }
     }
 
@@ -1076,6 +1140,218 @@ impl PyPersistReaderForTest {
         })?;
         Ok((strings(result), counts.to_py(py)?))
     }
+
+    // The state delta over the reader and the native state resolver, each
+    // returning its result and what it counted.
+
+    /// `(current_state, delta_ids, new_latest_event_ids, counts)`.
+    #[allow(clippy::too_many_arguments, clippy::type_complexity)]
+    fn get_new_state_after_events<'py>(
+        &self,
+        py: Python<'py>,
+        room_id: String,
+        server_name: String,
+        events_and_contexts: Vec<(Py<Event>, PersistContext)>,
+        old_latest_event_ids: Vec<String>,
+        new_latest_event_ids: Vec<String>,
+        should_prune: bool,
+        now_ms: i64,
+    ) -> PyResult<(
+        Option<Bound<'py, PyDict>>,
+        Option<Bound<'py, PyDict>>,
+        HashSet<String>,
+        Bound<'py, PyDict>,
+    )> {
+        let (events, event_map) = chunk(py, &events_and_contexts)?;
+        let old: HashSet<Arc<str>> = arcs(old_latest_event_ids).into_iter().collect();
+        let new: HashSet<Arc<str>> = arcs(new_latest_event_ids).into_iter().collect();
+        let resolver = self.resolver();
+        let (result, counts) = self.run(py, |reader| {
+            async move {
+                let mut counts = PersistCounts::default();
+                let result = state_delta::get_new_state_after_events(
+                    &reader,
+                    &resolver,
+                    &room_id,
+                    &server_name,
+                    &events,
+                    &event_map,
+                    &old,
+                    new,
+                    should_prune,
+                    now_ms,
+                    &mut counts,
+                )
+                .await?;
+                Ok((result, counts))
+            }
+            .boxed()
+        })?;
+        Ok((
+            result
+                .current_state
+                .map(|state| state_map_to_py(py, &state))
+                .transpose()?,
+            result
+                .delta_ids
+                .map(|delta| state_map_to_py(py, &delta))
+                .transpose()?,
+            strings(result.new_latest_event_ids),
+            counts.to_py(py)?,
+        ))
+    }
+
+    /// `(to_delete, to_insert)`.
+    #[allow(clippy::type_complexity)]
+    fn calculate_state_delta<'py>(
+        &self,
+        py: Python<'py>,
+        room_id: String,
+        current_state: HashMap<(String, String), String>,
+    ) -> PyResult<(Vec<(String, String)>, Bound<'py, PyDict>)> {
+        let current_state = state_map(current_state);
+        let delta = self.run(py, |reader| {
+            async move {
+                    state_delta::calculate_state_delta(&reader, &room_id, &current_state).await
+                }
+                .boxed()
+        })?;
+        let (to_delete, to_insert, _) = delta_to_py(py, &delta)?;
+        Ok((to_delete, to_insert))
+    }
+
+    fn is_server_still_joined(
+        &self,
+        py: Python<'_>,
+        room_id: String,
+        server_name: String,
+        events_and_contexts: Vec<(Py<Event>, PersistContext)>,
+        to_delete: Vec<(String, String)>,
+        to_insert: HashMap<(String, String), String>,
+    ) -> PyResult<bool> {
+        let (events, _) = chunk(py, &events_and_contexts)?;
+        let delta = DeltaState {
+            to_delete: to_delete
+                .iter()
+                .map(|(event_type, state_key)| StateKey::new(event_type, state_key))
+                .collect(),
+            to_insert: state_map(to_insert),
+            no_longer_in_room: false,
+        };
+        self.run(py, |reader| {
+            async move {
+                state_delta::is_server_still_joined(
+                    &reader,
+                    &room_id,
+                    &server_name,
+                    &events,
+                    &delta,
+                )
+                .await
+            }
+            .boxed()
+        })
+    }
+
+    /// `(new_forward_extremities, delta, counts)`, the delta as
+    /// `(to_delete, to_insert, no_longer_in_room)`.
+    #[allow(clippy::type_complexity)]
+    fn calculate_new_forward_extremities_and_state_delta<'py>(
+        &self,
+        py: Python<'py>,
+        room_id: String,
+        server_name: String,
+        events_and_contexts: Vec<(Py<Event>, PersistContext)>,
+        now_ms: i64,
+    ) -> PyResult<(
+        Option<HashSet<String>>,
+        Option<PyDelta<'py>>,
+        Bound<'py, PyDict>,
+    )> {
+        let (result, counts) =
+            self.run_orchestration(py, room_id, server_name, events_and_contexts, now_ms, false)?;
+        Ok((
+            result.new_forward_extremities.map(strings),
+            result
+                .state_delta
+                .map(|delta| delta_to_py(py, &delta))
+                .transpose()?,
+            counts.to_py(py)?,
+        ))
+    }
+
+    /// `(new_forward_extremities, delta, new_state_dag_extremities, counts)`,
+    /// the delta as `(to_delete, to_insert, no_longer_in_room)`.
+    #[allow(clippy::type_complexity)]
+    fn process_state_dag_forward_extremities_and_state_delta<'py>(
+        &self,
+        py: Python<'py>,
+        room_id: String,
+        server_name: String,
+        events_and_contexts: Vec<(Py<Event>, PersistContext)>,
+        now_ms: i64,
+    ) -> PyResult<(
+        Option<HashSet<String>>,
+        Option<PyDelta<'py>>,
+        Option<HashSet<String>>,
+        Bound<'py, PyDict>,
+    )> {
+        let (result, counts) =
+            self.run_orchestration(py, room_id, server_name, events_and_contexts, now_ms, true)?;
+        Ok((
+            result.new_forward_extremities.map(strings),
+            result
+                .state_delta
+                .map(|delta| delta_to_py(py, &delta))
+                .transpose()?,
+            result.new_state_dag_extremities.map(strings),
+            counts.to_py(py)?,
+        ))
+    }
+}
+
+/// A delta as Python's test reads it: `(to_delete, to_insert,
+/// no_longer_in_room)`.
+type PyDelta<'py> = (Vec<(String, String)>, Bound<'py, PyDict>, bool);
+
+fn delta_to_py<'py>(py: Python<'py>, delta: &DeltaState) -> PyResult<PyDelta<'py>> {
+    Ok((
+        delta
+            .to_delete
+            .iter()
+            .map(|key| (key.0.to_string(), key.1.to_string()))
+            .collect(),
+        state_map_to_py(py, &delta.to_insert)?,
+        delta.no_longer_in_room,
+    ))
+}
+
+fn state_map(map: HashMap<(String, String), String>) -> StateMap {
+    map.into_iter()
+        .map(|((event_type, state_key), event_id)| {
+            (StateKey::new(&event_type, &state_key), Arc::from(event_id))
+        })
+        .collect()
+}
+
+/// The chunk's events as the computation reads them, and as state
+/// resolution's map of events in flight.
+fn chunk(
+    py: Python<'_>,
+    events_and_contexts: &[(Py<Event>, PersistContext)],
+) -> PyResult<(Vec<NewEvent>, EventMap)> {
+    let events = new_events(py, events_and_contexts)?;
+    let event_map = events_and_contexts
+        .iter()
+        .map(|(event, _)| {
+            let event = event.bind(py).get();
+            (
+                Arc::from(event.event_id()),
+                StateResEvent::new(event.shallow_clone()),
+            )
+        })
+        .collect();
+    Ok((events, event_map))
 }
 
 fn new_events(

rust/src/state/mod.rs

@@ -47,7 +47,7 @@ use pyo3::{
 };
 
 use crate::caches::{python::PyCacheKey, Cache};
-use crate::database::postgres::pool::PyConnectionPool;
+use crate::database::postgres::pool::{ConnectionPool, PyConnectionPool};
 use crate::deferred::create_deferred;
 use crate::event_auth::AuthFailure;
 use crate::events::Event;
@@ -351,9 +351,9 @@ impl StateResStats {
 }
 
 /// What a resolution's deferred resolves to: the state and the stats.
-struct ResolveOutput {
-    state: StateMap,
-    stats: StateResStats,
+pub(crate) struct ResolveOutput {
+    pub(crate) state: StateMap,
+    pub(crate) stats: StateResStats,
 }
 
 impl<'py> IntoPyObject<'py> for ResolveOutput {
@@ -381,15 +381,9 @@ enum StateGroupsResult {
     State(StateMap),
 }
 
-/// Port of `_make_state_cache_entry`: how `new_state` relates to the input
-/// groups, in the order they were given, as the Python breaks ties by its
-/// dict's order.
-fn make_state_cache_entry(
-    new_state: StateMap,
-    groups: &[(i64, Arc<StateMap>)],
-) -> StateGroupsResult {
-    // First look for exact matches, by the set of event IDs as the Python
-    // does.
+/// The first of the input groups whose state is `new_state`, compared by
+/// the set of event IDs as `_make_state_cache_entry` does.
+pub(crate) fn matching_group(new_state: &StateMap, groups: &[(i64, Arc<StateMap>)]) -> Option<i64> {
     let new_event_ids: HashSet<&Arc<str>> = new_state.values().collect();
     for (group, state) in groups {
         if new_event_ids.len() != state.len() {
@@ -397,9 +391,22 @@ fn make_state_cache_entry(
         }
         let old_event_ids: HashSet<&Arc<str>> = state.values().collect();
         if new_event_ids == old_event_ids {
-            return StateGroupsResult::Group(*group);
+            return Some(*group);
         }
     }
+    None
+}
+
+/// Port of `_make_state_cache_entry`: how `new_state` relates to the input
+/// groups, in the order they were given, as the Python breaks ties by its
+/// dict's order.
+fn make_state_cache_entry(
+    new_state: StateMap,
+    groups: &[(i64, Arc<StateMap>)],
+) -> StateGroupsResult {
+    if let Some(group) = matching_group(&new_state, groups) {
+        return StateGroupsResult::Group(group);
+    }
 
     // Failing that, the closest match. A delta cannot remove keys, so a
     // group with a key the new state lacks is no candidate.
@@ -524,6 +531,11 @@ impl StateResolutionHandler {
         Arc::clone(&self.store.caches)
     }
 
+    /// The conflict cache, which the Rust event persister shares.
+    pub(crate) fn conflict_cache(&self) -> Arc<ConflictCache> {
+        Arc::clone(&self.cache)
+    }
+
     /// Convert the state sets under the GIL and start a resolution: either
     /// the answer, if nothing conflicts, or the future that computes it.
     fn start(
@@ -653,67 +665,95 @@ impl StateResolutionHandler {
         let caches = self.caches();
 
         Ok(Started::Pending(Box::pin(async move {
-            let start = Instant::now();
-            let store = PostgresStateResStore::begin(&pool, handle, Some(caches))
-                .await
-                .map_err(resolve_error_to_py)?;
-            let mut connection_time = start.elapsed();
-            let prepared = v2::prepare(
-                &store,
+            resolve_on_pool(
+                &pool,
+                handle,
+                caches,
+                &cache,
                 &room_id,
                 room_version,
                 &state_sets,
                 (unconflicted, conflicted),
                 events,
-                Some(&cache),
+                stats,
             )
             .await
-            .map_err(resolve_error_to_py)?;
-
-            // The rest does no I/O, so give the connection back first.
-            let finish_start = Instant::now();
-            let db = store.finish().await.map_err(resolve_error_to_py)?;
-            connection_time += finish_start.elapsed();
-            let resolution = match prepared {
-                v2::Prepared::Done(resolution) => resolution,
-                v2::Prepared::Pending(pending) => pending
-                    .run(Some(&cache))
-                    .await
-                    .map_err(resolve_error_to_py)?,
-            };
-
-            let core = &resolution.stats;
-            let stats = StateResStats {
-                total_time: start.elapsed().as_secs_f64(),
-                convert_in_time: stats.convert_in_time,
-                state_map_cache_hits: stats.state_map_cache_hits,
-                db_time: db.db_time.as_secs_f64(),
-                parse_time: db.parse_time.as_secs_f64(),
-                cpu_time: (core.sort_time + core.auth_time + core.mainline_time).as_secs_f64(),
-                sort_time: core.sort_time.as_secs_f64(),
-                auth_time: core.auth_time.as_secs_f64(),
-                mainline_time: core.mainline_time.as_secs_f64(),
-                auth_chain_time: core.auth_chain_time.as_secs_f64(),
-                load_events_time: core.load_events_time.as_secs_f64(),
-                connection_time: connection_time.as_secs_f64(),
-                db_queries: db.queries,
-                db_rows: db.rows,
-                db_events: db.events_fetched,
-                event_cache_hits: db.events_cached,
-                chain_position_cache_hits: db.chain_positions_cached,
-                chain_link_cache_hits: db.chain_links_cached,
-                convert_out_time: 0.0,
-                conflicted_events: core.conflicted_events,
-                cache_hit: core.cache_hit,
-            };
-            Ok(ResolveOutput {
-                state: resolution.merged(),
-                stats,
-            })
+            .map_err(resolve_error_to_py)
         })))
     }
 }
 
+/// Resolve converted state sets with conflicts over the native pool: the
+/// reads in one read-only transaction on a pooled connection, which goes
+/// back to the pool before the algorithm runs. `stats` holds what the
+/// caller recorded before. For callers in Rust, and behind
+/// [`StateResolutionHandler`]'s deferreds.
+#[allow(clippy::too_many_arguments)]
+pub(crate) async fn resolve_on_pool(
+    pool: &ConnectionPool,
+    handle: tokio::runtime::Handle,
+    caches: Arc<StoreCaches>,
+    cache: &ConflictCache,
+    room_id: &str,
+    room_version: &'static RoomVersion,
+    state_sets: &[Arc<StateMap>],
+    separated: (StateMap, v2::Conflicted),
+    events: EventMap,
+    stats: StateResStats,
+) -> anyhow::Result<ResolveOutput> {
+    let start = Instant::now();
+    let store = PostgresStateResStore::begin(pool, handle, Some(caches)).await?;
+    let mut connection_time = start.elapsed();
+    let prepared = v2::prepare(
+        &store,
+        room_id,
+        room_version,
+        state_sets,
+        separated,
+        events,
+        Some(cache),
+    )
+    .await?;
+
+    // The rest does no I/O, so give the connection back first.
+    let finish_start = Instant::now();
+    let db = store.finish().await?;
+    connection_time += finish_start.elapsed();
+    let resolution = match prepared {
+        v2::Prepared::Done(resolution) => resolution,
+        v2::Prepared::Pending(pending) => pending.run(Some(cache)).await?,
+    };
+
+    let core = &resolution.stats;
+    let stats = StateResStats {
+        total_time: start.elapsed().as_secs_f64(),
+        convert_in_time: stats.convert_in_time,
+        state_map_cache_hits: stats.state_map_cache_hits,
+        db_time: db.db_time.as_secs_f64(),
+        parse_time: db.parse_time.as_secs_f64(),
+        cpu_time: (core.sort_time + core.auth_time + core.mainline_time).as_secs_f64(),
+        sort_time: core.sort_time.as_secs_f64(),
+        auth_time: core.auth_time.as_secs_f64(),
+        mainline_time: core.mainline_time.as_secs_f64(),
+        auth_chain_time: core.auth_chain_time.as_secs_f64(),
+        load_events_time: core.load_events_time.as_secs_f64(),
+        connection_time: connection_time.as_secs_f64(),
+        db_queries: db.queries,
+        db_rows: db.rows,
+        db_events: db.events_fetched,
+        event_cache_hits: db.events_cached,
+        chain_position_cache_hits: db.chain_positions_cached,
+        chain_link_cache_hits: db.chain_links_cached,
+        convert_out_time: 0.0,
+        conflicted_events: core.conflicted_events,
+        cache_hit: core.cache_hit,
+    };
+    Ok(ResolveOutput {
+        state: resolution.merged(),
+        stats,
+    })
+}
+
 /// State sets converted from Python, with the interner that built them (so
 /// the event map shares its strings) and what the conversion recorded.
 struct Converted {

synapse/synapse_rust/persist.pyi

@@ -169,6 +169,61 @@ class PersistReaderForTest:
         events_and_contexts: list[tuple[EventBase, EventContext]],
         now_ms: int,
     ) -> tuple[set[str], dict[str, Any]]: ...
+    def get_new_state_after_events(
+        self,
+        room_id: str,
+        server_name: str,
+        events_and_contexts: list[tuple[EventBase, EventContext]],
+        old_latest_event_ids: list[str],
+        new_latest_event_ids: list[str],
+        should_prune: bool,
+        now_ms: int,
+    ) -> tuple[
+        dict[tuple[str, str], str] | None,
+        dict[tuple[str, str], str] | None,
+        set[str],
+        dict[str, Any],
+    ]:
+        """`(current_state, delta_ids, new_latest_event_ids, counts)`."""
+    def calculate_state_delta(
+        self, room_id: str, current_state: dict[tuple[str, str], str]
+    ) -> tuple[list[tuple[str, str]], dict[tuple[str, str], str]]:
+        """`(to_delete, to_insert)`."""
+    def is_server_still_joined(
+        self,
+        room_id: str,
+        server_name: str,
+        events_and_contexts: list[tuple[EventBase, EventContext]],
+        to_delete: list[tuple[str, str]],
+        to_insert: dict[tuple[str, str], str],
+    ) -> bool: ...
+    def calculate_new_forward_extremities_and_state_delta(
+        self,
+        room_id: str,
+        server_name: str,
+        events_and_contexts: list[tuple[EventBase, EventContext]],
+        now_ms: int,
+    ) -> tuple[
+        set[str] | None,
+        tuple[list[tuple[str, str]], dict[tuple[str, str], str], bool] | None,
+        dict[str, Any],
+    ]:
+        """`(new_forward_extremities, delta, counts)`, the delta as
+        `(to_delete, to_insert, no_longer_in_room)`."""
+    def process_state_dag_forward_extremities_and_state_delta(
+        self,
+        room_id: str,
+        server_name: str,
+        events_and_contexts: list[tuple[EventBase, EventContext]],
+        now_ms: int,
+    ) -> tuple[
+        set[str] | None,
+        tuple[list[tuple[str, str]], dict[tuple[str, str], str], bool] | None,
+        set[str] | None,
+        dict[str, Any],
+    ]:
+        """`(new_forward_extremities, delta, new_state_dag_extremities,
+        counts)`."""
 
 def persist_gil_stats() -> dict[str, tuple[int, float]]:
     """`{site: (calls, seconds held)}` for every time event persistence took

tests/storage/test_persist_reader.py

@@ -25,7 +25,7 @@ from unittest import SkipTest
 
 from twisted.internet.testing import MemoryReactor
 
-from synapse.api.constants import EventTypes
+from synapse.api.constants import EventTypes, Membership
 from synapse.api.errors import NotFoundError, SynapseError
 from synapse.api.room_versions import RoomVersions
 from synapse.events import EventBase
@@ -35,8 +35,10 @@ from synapse.rest.client import login, room
 from synapse.server import HomeServer
 from synapse.storage.database import LoggingTransaction
 from synapse.storage.databases.main.event_federation import EventFederationStore
+from synapse.storage.databases.main.events import DeltaState
 from synapse.storage.rust_pool import RustConnectionPool
 from synapse.synapse_rust.persist import PersistReaderForTest
+from synapse.types import StateMap, create_requester
 from synapse.types.handlers import SLIDING_SYNC_DEFAULT_BUMP_EVENT_TYPES
 from synapse.util.clock import Clock
 
@@ -955,3 +957,440 @@ class ExtremitiesTestCase(PersistReaderTestBase):
             ),
             {keep, being_persisted[0].event_id},
         )
+
+
+_Delta = tuple[set[tuple[str, str]], dict[tuple[str, str], str], bool]
+
+
+def _python_delta(delta: DeltaState | None) -> _Delta | None:
+    if delta is None:
+        return None
+    return (set(delta.to_delete), dict(delta.to_insert), delta.no_longer_in_room)
+
+
+def _rust_delta(
+    delta: tuple[list[tuple[str, str]], dict[tuple[str, str], str], bool] | None,
+) -> _Delta | None:
+    if delta is None:
+        return None
+    to_delete, to_insert, no_longer_in_room = delta
+    return (set(to_delete), to_insert, no_longer_in_room)
+
+
+class StateDeltaTestCase(PersistReaderTestBase):
+    """The state delta in `rust/src/persist/state_delta.rs` against the
+    controller's Python methods, over the same database and the same new
+    events, which are built with real contexts but not persisted."""
+
+    def default_config(self) -> dict[str, Any]:
+        config = super().default_config()
+        config["experimental_features"]["msc4242_enabled"] = True
+        return config
+
+    def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
+        super().prepare(reactor, clock, hs)
+        persistence = hs.get_storage_controllers().persistence
+        assert persistence is not None
+        self.persistence = persistence
+
+    def _create(
+        self,
+        event_type: str,
+        content: dict[str, Any],
+        *,
+        state_key: str | None = None,
+        sender: str | None = None,
+        prev_event_ids: list[str] | None = None,
+        room_id: str | None = None,
+    ) -> tuple[EventBase, EventContext]:
+        """An event as the event creation handler builds it, with its state
+        group stored but the event not persisted."""
+        event_dict: dict[str, Any] = {
+            "type": event_type,
+            "room_id": room_id or self.room_id,
+            "sender": sender or self.alice,
+            "content": content,
+        }
+        if state_key is not None:
+            event_dict["state_key"] = state_key
+        event, unpersisted = self.get_success(
+            self.hs.get_event_creation_handler().create_event(
+                create_requester(sender or self.alice),
+                event_dict,
+                prev_event_ids=prev_event_ids,
+            )
+        )
+        context = self.get_success(unpersisted.persist(event))
+        return event, context
+
+    def _topic(
+        self, topic: str, prev_event_ids: list[str] | None = None
+    ) -> tuple[EventBase, EventContext]:
+        return self._create(
+            EventTypes.Topic,
+            {"topic": topic},
+            state_key="",
+            prev_event_ids=prev_event_ids,
+        )
+
+    def _leave(
+        self, user: str, prev_event_ids: list[str] | None = None
+    ) -> tuple[EventBase, EventContext]:
+        return self._create(
+            EventTypes.Member,
+            {"membership": Membership.LEAVE},
+            state_key=user,
+            sender=user,
+            prev_event_ids=prev_event_ids,
+        )
+
+    def _persist(self, event_and_context: tuple[EventBase, EventContext]) -> str:
+        event, context = event_and_context
+        self.get_success(self.persistence.persist_event(event, context))
+        return event.event_id
+
+    def _latest(self) -> set[str]:
+        return set(
+            self.get_success(self.store.get_latest_event_ids_in_room(self.room_id))
+        )
+
+    def _new_state(
+        self,
+        events: list[tuple[EventBase, EventContext]],
+        old: set[str],
+        new: set[str],
+    ) -> tuple[StateMap[str] | None, StateMap[str] | None, set[str], dict[str, Any]]:
+        current_state, delta_ids, new_latest = self.get_success(
+            self.persistence._get_new_state_after_events(
+                self.room_id, events, old, set(new)
+            )
+        )
+        rust_state, rust_delta, rust_latest, counts = (
+            self.rust.get_new_state_after_events(
+                self.room_id,
+                self.hs.hostname,
+                events,
+                list(old),
+                list(new),
+                True,
+                self.clock.time_msec(),
+            )
+        )
+        self.assertEqual(
+            rust_state, None if current_state is None else dict(current_state)
+        )
+        self.assertEqual(rust_delta, None if delta_ids is None else dict(delta_ids))
+        self.assertEqual(rust_latest, new_latest)
+        return rust_state, rust_delta, rust_latest, counts
+
+    def test_get_new_state_after_events(self) -> None:
+        latest = self._latest()
+
+        # A message: the state groups do not change.
+        message = self._create(EventTypes.Message, {"msgtype": "m.text", "body": "x"})
+        current_state, delta, _, _ = self._new_state(
+            [message], latest, {message[0].event_id}
+        )
+        self.assertIsNone(current_state)
+        self.assertIsNone(delta)
+
+        # A state event: its context knows the delta.
+        topic = self._topic("one")
+        current_state, delta, _, _ = self._new_state(
+            [topic], latest, {topic[0].event_id}
+        )
+        self.assertIsNone(current_state)
+        self.assertEqual(delta, {(EventTypes.Topic, ""): topic[0].event_id})
+
+        # Without the delta, the new group's state.
+        topic[1].state_group_deltas = {}
+        current_state, delta, _, counts = self._new_state(
+            [topic], latest, {topic[0].event_id}
+        )
+        assert current_state is not None
+        self.assertIsNone(delta)
+        self.assertEqual(current_state[(EventTypes.Topic, "")], topic[0].event_id)
+        self.assertEqual(counts["state_resolutions"], 0)
+
+    def test_resolution_of_a_fork(self) -> None:
+        (parent,) = self._latest()
+        persisted = self._persist(self._topic("persisted", [parent]))
+        fork = self._topic("fork", [parent])
+        current_state, delta, latest, counts = self._new_state(
+            [fork], {persisted}, {persisted, fork[0].event_id}
+        )
+        assert current_state is not None
+        self.assertIsNone(delta)
+        self.assertIn(
+            current_state[(EventTypes.Topic, "")], {persisted, fork[0].event_id}
+        )
+        self.assertEqual(latest, {persisted, fork[0].event_id})
+        self.assertEqual(counts["state_resolutions"], 1)
+
+    def _pending_deletions(self, groups: list[int]) -> dict[int, tuple[int, int]]:
+        rows = self.get_success(
+            self.hs.get_datastores().state_deletion.db_pool.simple_select_many_batch(
+                table="state_groups_pending_deletion",
+                column="state_group",
+                iterable=groups,
+                retcols=("state_group", "sequence_number", "insertion_ts"),
+            )
+        )
+        return {group: (sequence, ts) for group, sequence, ts in rows}
+
+    def test_resolution_bumps_pending_deletions(self) -> None:
+        (parent,) = self._latest()
+        persisted_topic = self._topic("persisted", [parent])
+        persisted = self._persist(persisted_topic)
+        fork = self._topic("fork", [parent])
+        persisted_group = persisted_topic[1].state_group
+        fork_group = fork[1].state_group
+        assert persisted_group is not None and fork_group is not None
+        groups = [persisted_group, fork_group]
+
+        self.get_success(
+            self.hs.get_datastores().state_deletion.mark_state_groups_as_pending_deletion(
+                groups
+            )
+        )
+        before = self._pending_deletions(groups)
+        self.assertEqual(before.keys(), set(groups))
+
+        # Only Rust runs, so the bump is its own.
+        self.reactor.advance(60)
+        self.rust.get_new_state_after_events(
+            self.room_id,
+            self.hs.hostname,
+            [fork],
+            [persisted],
+            [persisted, fork[0].event_id],
+            True,
+            self.clock.time_msec(),
+        )
+        after = self._pending_deletions(groups)
+        for group in groups:
+            self.assertGreater(after[group][0], before[group][0])
+            self.assertEqual(after[group][1], self.clock.time_msec())
+
+    def test_resolution_of_a_deleted_group(self) -> None:
+        (parent,) = self._latest()
+        persisted = self._persist(self._topic("persisted", [parent]))
+        fork = self._topic("fork", [parent])
+        fork[1]._state_group = 1_000_000_000
+        new = {persisted, fork[0].event_id}
+
+        failure = self.get_failure(
+            self.persistence._get_new_state_after_events(
+                self.room_id, [fork], {persisted}, set(new)
+            ),
+            Exception,
+        )
+        self.assertIn("deleted", str(failure.value))
+        with self.assertRaisesRegex(Exception, "State groups have been deleted"):
+            self.rust.get_new_state_after_events(
+                self.room_id,
+                self.hs.hostname,
+                [fork],
+                [persisted],
+                list(new),
+                True,
+                self.clock.time_msec(),
+            )
+
+    def test_calculate_state_delta(self) -> None:
+        current = dict(
+            self.get_success(self.store.get_partial_current_state_ids(self.room_id))
+        )
+        new_state = dict(current)
+        del new_state[(EventTypes.JoinRules, "")]
+        new_state[(EventTypes.Name, "")] = "$changed"
+        new_state[(EventTypes.Topic, "")] = "$added"
+
+        expected = self.get_success(
+            self.persistence._calculate_state_delta(self.room_id, new_state)
+        )
+        to_delete, to_insert = self.rust.calculate_state_delta(self.room_id, new_state)
+        self.assertEqual(set(to_delete), set(expected.to_delete))
+        self.assertEqual(to_insert, expected.to_insert)
+        self.assertEqual(set(to_delete), {(EventTypes.JoinRules, "")})
+        self.assertEqual(
+            to_insert,
+            {(EventTypes.Name, ""): "$changed", (EventTypes.Topic, ""): "$added"},
+        )
+
+    def _still_joined(
+        self,
+        events: list[tuple[EventBase, EventContext]],
+        to_delete: list[tuple[str, str]],
+        to_insert: dict[tuple[str, str], str],
+    ) -> bool:
+        expected = self.get_success(
+            self.persistence._is_server_still_joined(
+                self.room_id, events, DeltaState(to_delete, to_insert)
+            )
+        )
+        result = self.rust.is_server_still_joined(
+            self.room_id, self.hs.hostname, events, to_delete, to_insert
+        )
+        self.assertEqual(result, expected)
+        return result
+
+    def test_is_server_still_joined(self) -> None:
+        alice_key = (EventTypes.Member, self.alice)
+        bob_key = (EventTypes.Member, self.bob)
+        alice_leave = self._leave(self.alice)
+        bob_leave = self._leave(self.bob)
+
+        # No local membership changes.
+        self.assertTrue(
+            self._still_joined(
+                [], [(EventTypes.Member, "@x:remote")], {(EventTypes.Topic, ""): "$t"}
+            )
+        )
+        # Alice leaves, Bob is still joined.
+        self.assertTrue(
+            self._still_joined([alice_leave], [], {alice_key: alice_leave[0].event_id})
+        )
+        # Both leave, or one leaves and the other's membership goes.
+        self.assertFalse(
+            self._still_joined(
+                [alice_leave, bob_leave],
+                [],
+                {
+                    alice_key: alice_leave[0].event_id,
+                    bob_key: bob_leave[0].event_id,
+                },
+            )
+        )
+        self.assertFalse(
+            self._still_joined(
+                [alice_leave], [bob_key], {alice_key: alice_leave[0].event_id}
+            )
+        )
+        # A join in the database, which is Alice's current membership.
+        alice_join = self.get_success(
+            self.store.get_partial_current_state_ids(self.room_id)
+        )[alice_key]
+        self.assertTrue(
+            self._still_joined([alice_leave], [bob_key], {alice_key: alice_join})
+        )
+
+    def test_still_joined_by_a_join_in_the_chunk(self) -> None:
+        alice_key = (EventTypes.Member, self.alice)
+        bob_key = (EventTypes.Member, self.bob)
+        # Alice changes her display name, a join. Bob's membership goes, so
+        # only the join in the chunk keeps the server in the room.
+        rejoin = self._create(
+            EventTypes.Member,
+            {"membership": Membership.JOIN, "displayname": "Alice"},
+            state_key=self.alice,
+        )
+        self.assertTrue(
+            self._still_joined([rejoin], [bob_key], {alice_key: rejoin[0].event_id})
+        )
+        # Without the event in the chunk, the server reads the database,
+        # where the join is not stored.
+        self.assertFalse(
+            self._still_joined([], [bob_key], {alice_key: rejoin[0].event_id})
+        )
+
+    def _orchestrate(
+        self, events: list[tuple[EventBase, EventContext]]
+    ) -> tuple[set[str] | None, _Delta | None, dict[str, Any]]:
+        forward, delta = self.get_success(
+            self.persistence._calculate_new_forward_extremities_and_state_delta(
+                self.room_id, events
+            )
+        )
+        rust_forward, rust_delta, counts = (
+            self.rust.calculate_new_forward_extremities_and_state_delta(
+                self.room_id, self.hs.hostname, events, self.clock.time_msec()
+            )
+        )
+        self.assertEqual(rust_forward, forward)
+        self.assertEqual(_rust_delta(rust_delta), _python_delta(delta))
+        return rust_forward, _rust_delta(rust_delta), counts
+
+    def test_calculate_new_forward_extremities_and_state_delta(self) -> None:
+        # A message only moves the extremity.
+        message = self._create(EventTypes.Message, {"msgtype": "m.text", "body": "x"})
+        forward, delta, counts = self._orchestrate([message])
+        self.assertEqual(forward, {message[0].event_id})
+        self.assertIsNone(delta)
+        self.assertEqual(counts["state_delta"], 0)
+
+        # A state event, whose delta the context knows.
+        topic = self._topic("one")
+        forward, delta, counts = self._orchestrate([topic])
+        self.assertEqual(
+            delta, (set(), {(EventTypes.Topic, ""): topic[0].event_id}, False)
+        )
+        self.assertEqual(counts["state_delta"], 1)
+        self.assertEqual(counts["state_delta_single_event"], 1)
+        self.assertEqual(counts["state_delta_reuse_delta"], 1)
+
+        # The last local user leaves, so the server leaves the room.
+        self._persist(topic)
+        self._persist(self._leave(self.alice))
+        bob_leave = self._leave(self.bob)
+        forward, delta, counts = self._orchestrate([bob_leave])
+        self.assertEqual(
+            delta,
+            (set(), {(EventTypes.Member, self.bob): bob_leave[0].event_id}, True),
+        )
+
+    def test_orchestration_resolves_a_fork(self) -> None:
+        (parent,) = self._latest()
+        persisted = self._persist(self._topic("persisted", [parent]))
+        fork = self._topic("fork", [parent])
+        forward, delta, counts = self._orchestrate([fork])
+        self.assertEqual(forward, {persisted, fork[0].event_id})
+        assert delta is not None
+        self.assertEqual(counts["state_resolutions"], 1)
+        self.assertEqual(counts["state_delta"], 1)
+        self.assertEqual(counts["state_delta_single_event"], 0)
+
+    def test_state_dag_orchestration(self) -> None:
+        room_id = self.helper.create_room_as(
+            self.alice,
+            room_version=RoomVersions.MSC4242v12.identifier,
+            tok=self.alice_tok,
+        )
+
+        def compare(
+            events: list[tuple[EventBase, EventContext]],
+        ) -> tuple[set[str] | None, _Delta | None, set[str] | None, dict[str, Any]]:
+            forward, delta, state_dag = self.get_success(
+                self.persistence._process_state_dag_forward_extremities_and_state_delta(
+                    room_id,
+                    events,  # type: ignore[arg-type]
+                )
+            )
+            rust_forward, rust_delta, rust_state_dag, counts = (
+                self.rust.process_state_dag_forward_extremities_and_state_delta(
+                    room_id, self.hs.hostname, events, self.clock.time_msec()
+                )
+            )
+            self.assertEqual(rust_forward, forward)
+            self.assertEqual(_rust_delta(rust_delta), _python_delta(delta))
+            self.assertEqual(rust_state_dag, state_dag)
+            return rust_forward, _rust_delta(rust_delta), rust_state_dag, counts
+
+        # A message moves only the room DAG.
+        message = self._create(
+            EventTypes.Message, {"msgtype": "m.text", "body": "x"}, room_id=room_id
+        )
+        forward, delta, state_dag, _ = compare([message])
+        self.assertEqual(forward, {message[0].event_id})
+        self.assertIsNone(delta)
+
+        # A state event moves both, and the state.
+        name = self._create(
+            EventTypes.Name, {"name": "dag"}, state_key="", room_id=room_id
+        )
+        forward, delta, state_dag, counts = compare([name])
+        self.assertEqual(state_dag, {name[0].event_id})
+        self.assertEqual(
+            delta, (set(), {(EventTypes.Name, ""): name[0].event_id}, False)
+        )
+        self.assertEqual(counts["state_dag_forward_extremities"], {1: 1})