misc fixes github/erikj/fixes_batch erikj/fixes_batch

1 files changed, 17 insertions, 20 deletions

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