1 files changed, 22 insertions, 16 deletions

diff --git a/app/javascript/mastodon/reducers/contexts.js b/app/javascript/mastodon/reducers/contexts.js
index fe8308d0c..c1ecf7e8a 100644
--- a/app/javascript/mastodon/reducers/contexts.js
+++ b/app/javascript/mastodon/reducers/contexts.js
@@ -21,24 +21,30 @@ const normalizeContext = (state, id, ancestors, descendants) => {
   });
 };
 
-const deleteFromContexts = (state, id) => {
-  state.getIn(['descendants', id], ImmutableList()).forEach(descendantId => {
-    state = state.updateIn(['ancestors', descendantId], ImmutableList(), list => list.filterNot(itemId => itemId === id));
-  });
+const deleteFromContexts = (immutableState, ids) => immutableState.withMutations(state => {
+  state.update('ancestors', immutableAncestors => immutableAncestors.withMutations(ancestors => {
+    state.update('descendants', immutableDescendants => immutableDescendants.withMutations(descendants => {
+      ids.forEach(id => {
+        descendants.get(id, ImmutableList()).forEach(descendantId => {
+          ancestors.update(descendantId, ImmutableList(), list => list.filterNot(itemId => itemId === id));
+        });
 
-  state.getIn(['ancestors', id], ImmutableList()).forEach(ancestorId => {
-    state = state.updateIn(['descendants', ancestorId], ImmutableList(), list => list.filterNot(itemId => itemId === id));
-  });
+        ancestors.get(id, ImmutableList()).forEach(ancestorId => {
+          descendants.update(ancestorId, ImmutableList(), list => list.filterNot(itemId => itemId === id));
+        });
 
-  state = state.deleteIn(['descendants', id]).deleteIn(['ancestors', id]);
+        descendants.delete(id);
+        ancestors.delete(id);
+      });
+    }));
+  }));
+});
 
-  return state;
-};
+const filterContexts = (state, relationship, statuses) => {
+  const ownedStatusIds = statuses.filter(status => status.get('account') === relationship.id)
+                                 .map(status => status.get('id'));
 
-const filterContexts = (state, relationship) => {
-  return state.map(
-    statuses => statuses.filter(
-      status => status.get('account') !== relationship.id));
+  return deleteFromContexts(state, ownedStatusIds);
 };
 
 const updateContext = (state, status, references) => {
@@ -61,11 +67,11 @@ export default function contexts(state = initialState, action) {
   switch(action.type) {
   case ACCOUNT_BLOCK_SUCCESS:
   case ACCOUNT_MUTE_SUCCESS:
-    return filterContexts(state, action.relationship);
+    return filterContexts(state, action.relationship, action.statuses);
   case CONTEXT_FETCH_SUCCESS:
     return normalizeContext(state, action.id, action.ancestors, action.descendants);
   case TIMELINE_DELETE:
-    return deleteFromContexts(state, action.id);
+    return deleteFromContexts(state, [action.id]);
   case TIMELINE_CONTEXT_UPDATE:
     return updateContext(state, action.status, action.references);
   default: