Fix Notion nested archive ordering

crates/locality-core/src/journal.rs

@@ -239,6 +239,40 @@ pub enum JournalApplyEffect {
     },
 }
 
+impl JournalApplyEffect {
+    pub fn operation_index(&self) -> usize {
+        match self {
+            Self::UpdatedBlock {
+                operation_index, ..
+            }
+            | Self::CreatedBlock {
+                operation_index, ..
+            }
+            | Self::MovedBlock {
+                operation_index, ..
+            }
+            | Self::ArchivedBlock {
+                operation_index, ..
+            }
+            | Self::ArchivedEntity {
+                operation_index, ..
+            }
+            | Self::UpdatedEntityBody {
+                operation_index, ..
+            }
+            | Self::UpdatedProperties {
+                operation_index, ..
+            }
+            | Self::MovedEntity {
+                operation_index, ..
+            }
+            | Self::CreatedEntity {
+                operation_index, ..
+            } => *operation_index,
+        }
+    }
+}
+
 #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub struct JournalPreimage {
     pub entity_id: RemoteId,

crates/locality-core/src/undo.rs

@@ -118,7 +118,8 @@ pub fn plan_journal_undo(entry: &JournalEntry) -> UndoPlan {
     let mut operations = Vec::new();
     let mut unsupported = Vec::new();
 
-    for (operation_index, operation) in entry.plan.operations.iter().enumerate().rev() {
+    for operation_index in reverse_apply_operation_indices(entry) {
+        let operation = &entry.plan.operations[operation_index];
         match operation {
             PushOperation::UpdateBlock { block_id, .. } => {
                 match find_preimage_block(entry, block_id) {
@@ -393,6 +394,61 @@ pub fn plan_journal_undo(entry: &JournalEntry) -> UndoPlan {
     }
 }
 
+fn reverse_apply_operation_indices(entry: &JournalEntry) -> Vec<usize> {
+    let operation_count = entry.plan.operations.len();
+    let fallback = || (0..operation_count).rev().collect::<Vec<_>>();
+    let non_archive_indices = entry
+        .plan
+        .operations
+        .iter()
+        .enumerate()
+        .filter_map(|(index, operation)| {
+            (!matches!(
+                operation,
+                PushOperation::ArchiveBlock { .. } | PushOperation::ArchiveEntity { .. }
+            ))
+            .then_some(index)
+        })
+        .collect::<Vec<_>>();
+    let archive_indices = entry
+        .plan
+        .operations
+        .iter()
+        .enumerate()
+        .filter_map(|(index, operation)| {
+            matches!(
+                operation,
+                PushOperation::ArchiveBlock { .. } | PushOperation::ArchiveEntity { .. }
+            )
+            .then_some(index)
+        })
+        .collect::<Vec<_>>();
+    if archive_indices.is_empty() {
+        return fallback();
+    }
+
+    let mut apply_order = Vec::with_capacity(operation_count);
+
+    for effect in &entry.apply_effects {
+        let operation_index = effect.operation_index();
+        if operation_index < operation_count && !apply_order.contains(&operation_index) {
+            apply_order.push(operation_index);
+        }
+    }
+
+    if apply_order.len() != operation_count
+        || !apply_order.starts_with(&non_archive_indices)
+        || apply_order[non_archive_indices.len()..]
+            .iter()
+            .any(|index| !archive_indices.contains(index))
+    {
+        return fallback();
+    }
+
+    apply_order.reverse();
+    apply_order
+}
+
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct UndoApplyRequest<'a> {
     pub target_push_id: &'a PushId,

crates/locality-core/tests/undo.rs

@@ -537,6 +537,215 @@ fn archive_entity_reverses_to_restore_archived_entity() {
     );
 }
 
+#[test]
+fn undo_reverses_complete_journaled_apply_order_instead_of_original_plan_order() {
+    let operations = vec![
+        PushOperation::ArchiveEntity {
+            entity_id: RemoteId::new("page-1"),
+        },
+        PushOperation::ArchiveBlock {
+            block_id: RemoteId::new("paragraph-1"),
+        },
+    ];
+    let mut entry = journal_entry_with_shadow(operations, shadow_with_frontmatter());
+    entry.apply_effects = vec![
+        JournalApplyEffect::ArchivedBlock {
+            operation_id: PushOperationId::for_operation(
+                &entry.push_id,
+                1,
+                &entry.plan.operations[1],
+            ),
+            operation_index: 1,
+            block_id: RemoteId::new("paragraph-1"),
+        },
+        JournalApplyEffect::ArchivedEntity {
+            operation_id: PushOperationId::for_operation(
+                &entry.push_id,
+                0,
+                &entry.plan.operations[0],
+            ),
+            operation_index: 0,
+            entity_id: RemoteId::new("page-1"),
+        },
+    ];
+
+    let plan = plan_journal_undo(&entry);
+
+    assert_eq!(plan.status, UndoPlanStatus::Complete);
+    assert_eq!(
+        plan.operations,
+        vec![
+            UndoOperation::RestoreArchivedEntity {
+                entity_id: RemoteId::new("page-1"),
+                expected: EntityUndoState {
+                    parent_id: RemoteId::new("old-parent"),
+                    title: "Roadmap".to_string(),
+                    properties: BTreeMap::from([
+                        ("Points".to_string(), PropertyValue::Number("2".to_string())),
+                        (
+                            "Status".to_string(),
+                            PropertyValue::String("Todo".to_string()),
+                        ),
+                    ]),
+                    body: "# Roadmap\n\nOld paragraph.".to_string(),
+                    archived: true,
+                },
+            },
+            UndoOperation::RestoreArchivedBlock {
+                block_id: RemoteId::new("paragraph-1"),
+                parent_id: RemoteId::new("page-1"),
+                after: Some(RemoteId::new("heading-1")),
+                content: "Old paragraph.".to_string(),
+                native_kind: None,
+            },
+        ]
+    );
+}
+
+#[test]
+fn nested_archive_effects_reverse_to_ancestor_first_restore_order() {
+    let operations = vec![
+        PushOperation::UpdateBlock {
+            block_id: RemoteId::new("survivor-1"),
+            content: "Updated survivor.".to_string(),
+        },
+        PushOperation::ArchiveBlock {
+            block_id: RemoteId::new("parent-1"),
+        },
+        PushOperation::ArchiveBlock {
+            block_id: RemoteId::new("child-1"),
+        },
+        PushOperation::ArchiveBlock {
+            block_id: RemoteId::new("grandchild-1"),
+        },
+    ];
+    let shadow = ShadowDocument::from_synced_body(
+        RemoteId::new("page-1"),
+        "Original survivor.\n\nParent.\n\nChild.\n\nGrandchild.",
+        9,
+        [
+            RemoteId::new("survivor-1"),
+            RemoteId::new("parent-1"),
+            RemoteId::new("child-1"),
+            RemoteId::new("grandchild-1"),
+        ],
+    )
+    .expect("nested archive shadow");
+    let mut entry = journal_entry_with_shadow(operations, shadow);
+    entry.apply_effects = vec![
+        JournalApplyEffect::UpdatedBlock {
+            operation_id: PushOperationId::for_operation(
+                &entry.push_id,
+                0,
+                &entry.plan.operations[0],
+            ),
+            operation_index: 0,
+            block_id: RemoteId::new("survivor-1"),
+        },
+        JournalApplyEffect::ArchivedBlock {
+            operation_id: PushOperationId::for_operation(
+                &entry.push_id,
+                3,
+                &entry.plan.operations[3],
+            ),
+            operation_index: 3,
+            block_id: RemoteId::new("grandchild-1"),
+        },
+        JournalApplyEffect::ArchivedBlock {
+            operation_id: PushOperationId::for_operation(
+                &entry.push_id,
+                2,
+                &entry.plan.operations[2],
+            ),
+            operation_index: 2,
+            block_id: RemoteId::new("child-1"),
+        },
+        JournalApplyEffect::ArchivedBlock {
+            operation_id: PushOperationId::for_operation(
+                &entry.push_id,
+                1,
+                &entry.plan.operations[1],
+            ),
+            operation_index: 1,
+            block_id: RemoteId::new("parent-1"),
+        },
+    ];
+
+    let plan = plan_journal_undo(&entry);
+
+    assert_eq!(plan.status, UndoPlanStatus::Complete);
+    assert_eq!(
+        plan.operations
+            .iter()
+            .filter_map(|operation| match operation {
+                UndoOperation::RestoreArchivedBlock { block_id, .. } => {
+                    Some(block_id.as_str())
+                }
+                _ => None,
+            })
+            .collect::<Vec<_>>(),
+        ["parent-1", "child-1", "grandchild-1"]
+    );
+    assert!(matches!(
+        plan.operations.last(),
+        Some(UndoOperation::RestoreBlockContent { block_id, content })
+            if block_id == &RemoteId::new("survivor-1")
+                && content == "Original survivor."
+    ));
+}
+
+#[test]
+fn undo_does_not_treat_non_archive_effect_order_as_remote_execution_order() {
+    let operations = vec![
+        PushOperation::UpdateBlock {
+            block_id: RemoteId::new("a"),
+            content: "Updated A".to_string(),
+        },
+        PushOperation::UpdateBlock {
+            block_id: RemoteId::new("b"),
+            content: "Updated B".to_string(),
+        },
+    ];
+    let mut entry = journal_entry_with_shadow(operations, multi_block_shadow());
+    entry.apply_effects = vec![
+        JournalApplyEffect::UpdatedBlock {
+            operation_id: PushOperationId::for_operation(
+                &entry.push_id,
+                1,
+                &entry.plan.operations[1],
+            ),
+            operation_index: 1,
+            block_id: RemoteId::new("b"),
+        },
+        JournalApplyEffect::UpdatedBlock {
+            operation_id: PushOperationId::for_operation(
+                &entry.push_id,
+                0,
+                &entry.plan.operations[0],
+            ),
+            operation_index: 0,
+            block_id: RemoteId::new("a"),
+        },
+    ];
+
+    let plan = plan_journal_undo(&entry);
+
+    assert_eq!(plan.status, UndoPlanStatus::Complete);
+    assert_eq!(
+        plan.operations,
+        vec![
+            UndoOperation::RestoreBlockContent {
+                block_id: RemoteId::new("b"),
+                content: "B".to_string(),
+            },
+            UndoOperation::RestoreBlockContent {
+                block_id: RemoteId::new("a"),
+                content: "A".to_string(),
+            },
+        ]
+    );
+}
+
 #[test]
 fn archive_entity_without_entity_preimage_is_blocked() {
     let mut entry = journal_entry(vec![PushOperation::ArchiveEntity {