diff --git a/synapse/events/snapshot.py b/synapse/events/snapshot.py
index a91a5d1e3c..1bb3d8f476 100644
--- a/synapse/events/snapshot.py
+++ b/synapse/events/snapshot.py
@@ -135,6 +135,8 @@ class EventContext(UnpersistedEventContextBase):
delta_ids: Optional[StateMap[str]] = None
app_service: Optional[ApplicationService] = None
+ _state_map_before_event: Optional[StateMap[str]] = None
+
partial_state: bool = False
@staticmethod
@@ -293,6 +295,11 @@ class EventContext(UnpersistedEventContextBase):
Maps a (type, state_key) to the event ID of the state event matching
this tuple.
"""
+ if self._state_map_before_event is not None:
+ if state_filter is not None:
+ return state_filter.filter_state(self._state_map_before_event)
+ return self._state_map_before_event
+
assert self.state_group_before_event is not None
return await self._storage.state.get_state_ids_for_group(
self.state_group_before_event, state_filter
@@ -374,26 +381,16 @@ class UnpersistedEventContext(UnpersistedEventContextBase):
events_and_persisted_context = []
for event, unpersisted_context in amended_events_and_context:
- if event.is_state():
- context = EventContext(
- storage=unpersisted_context._storage,
- state_group=unpersisted_context.state_group_after_event,
- state_group_before_event=unpersisted_context.state_group_before_event,
- state_delta_due_to_event=unpersisted_context.state_delta_due_to_event,
- partial_state=unpersisted_context.partial_state,
- prev_group=unpersisted_context.state_group_before_event,
- delta_ids=unpersisted_context.state_delta_due_to_event,
- )
- else:
- context = EventContext(
- storage=unpersisted_context._storage,
- state_group=unpersisted_context.state_group_after_event,
- state_group_before_event=unpersisted_context.state_group_before_event,
- state_delta_due_to_event=unpersisted_context.state_delta_due_to_event,
- partial_state=unpersisted_context.partial_state,
- prev_group=unpersisted_context.prev_group_for_state_group_before_event,
- delta_ids=unpersisted_context.delta_ids_to_state_group_before_event,
- )
+ context = EventContext(
+ storage=unpersisted_context._storage,
+ state_group=unpersisted_context.state_group_after_event,
+ state_group_before_event=unpersisted_context.state_group_before_event,
+ state_delta_due_to_event=unpersisted_context.state_delta_due_to_event,
+ partial_state=unpersisted_context.partial_state,
+ prev_group=unpersisted_context.prev_group_for_state_group_before_event,
+ delta_ids=unpersisted_context.delta_ids_to_state_group_before_event,
+ state_map_before_event=unpersisted_context.state_map_before_event,
+ )
events_and_persisted_context.append((event, context))
return events_and_persisted_context
|
