fix(resolve): Prevent submodule deletion

crates/gitbutler-edit-mode/src/lib.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
 use anyhow::{Context as _, Result, bail};
 use bstr::{BString, ByteSlice};
@@ -19,7 +19,13 @@ use gitbutler_operating_modes::{
 };
 use gitbutler_repo::{RepositoryExt as _, SignaturePurpose, signature};
 use gitbutler_stack::VirtualBranchesHandle;
-use gitbutler_workspace::branch_trees::{WorkspaceState, update_uncommitted_changes_with_tree};
+use gitbutler_workspace::{
+    branch_trees::{WorkspaceState, update_uncommitted_changes_with_tree},
+    submodules::{
+        configured_submodule_paths, is_submodule_related_path,
+        remove_untracked_excluding_submodule_paths,
+    },
+};
 use serde::Serialize;
 
 pub mod commands;
@@ -149,29 +155,95 @@ fn get_uncommited_changes(ctx: &Context) -> Result<git2::Oid> {
     Ok(uncommited_changes)
 }
 
+fn collect_checkout_paths(
+    repo: &git2::Repository,
+    baseline: &git2::Tree,
+    target: &git2::Tree,
+    submodule_paths: &[String],
+) -> Result<Vec<String>> {
+    let mut paths = HashSet::new();
+    let diff = repo.diff_tree_to_tree(Some(baseline), Some(target), None)?;
+
+    for delta in diff.deltas() {
+        if let Some(path) = delta.old_file().path() {
+            let path = path.to_string_lossy().to_string();
+            if !is_submodule_related_path(&path, submodule_paths) {
+                paths.insert(path);
+            }
+        }
+        if let Some(path) = delta.new_file().path() {
+            let path = path.to_string_lossy().to_string();
+            if !is_submodule_related_path(&path, submodule_paths) {
+                paths.insert(path);
+            }
+        }
+    }
+
+    let mut output = paths.into_iter().collect::<Vec<_>>();
+    output.sort();
+    Ok(output)
+}
+
 fn checkout_edit_branch(ctx: &Context, commit: git2::Commit) -> Result<()> {
     let repo = &*ctx.git2_repo.get()?;
+    let submodule_paths = configured_submodule_paths(repo);
+    let preserve_submodule_worktrees = !submodule_paths.is_empty();
+    let result = (|| -> Result<()> {
+        let current_head_tree = repo.head()?.peel_to_tree()?;
+
+        // Checkout commits's parent
+        let commit_parent = find_or_create_base_commit(repo, &commit)?;
+        let commit_parent_tree = commit_parent.tree()?;
+        let checkout_head_paths = collect_checkout_paths(
+            repo,
+            &current_head_tree,
+            &commit_parent_tree,
+            &submodule_paths,
+        )?;
 
-    // Checkout commits's parent
-    let commit_parent = find_or_create_base_commit(repo, &commit)?;
-    repo.reference(EDIT_BRANCH_REF, commit_parent.id(), true, "")?;
-    repo.set_head(EDIT_BRANCH_REF)?;
-    repo.checkout_head(Some(CheckoutBuilder::new().force().remove_untracked(true)))?;
+        repo.reference(EDIT_BRANCH_REF, commit_parent.id(), true, "")?;
+        repo.set_head(EDIT_BRANCH_REF)?;
+        let mut checkout_head = CheckoutBuilder::new();
+        checkout_head.force();
+        if preserve_submodule_worktrees {
+            for path in &checkout_head_paths {
+                checkout_head.path(path);
+            }
+        } else {
+            checkout_head.remove_untracked(true);
+        }
+        if !preserve_submodule_worktrees || !checkout_head_paths.is_empty() {
+            repo.checkout_head(Some(&mut checkout_head))?;
+            if preserve_submodule_worktrees {
+                remove_untracked_excluding_submodule_paths(repo)?;
+            }
+        }
 
-    // Checkout the commit as unstaged changes
-    let mut index = get_commit_index(ctx, &commit)?;
-
-    repo.checkout_index(
-        Some(&mut index),
-        Some(
-            CheckoutBuilder::new()
-                .force()
-                .remove_untracked(true)
-                .conflict_style_diff3(true),
-        ),
-    )?;
+        // Checkout the commit as unstaged changes
+        let mut index = get_commit_index(ctx, &commit)?;
+        let commit_tree = commit.tree()?;
+        let checkout_index_paths =
+            collect_checkout_paths(repo, &commit_parent_tree, &commit_tree, &submodule_paths)?;
+        let mut checkout_index = CheckoutBuilder::new();
+        checkout_index.force().conflict_style_diff3(true);
+        if preserve_submodule_worktrees {
+            for path in &checkout_index_paths {
+                checkout_index.path(path);
+            }
+        } else {
+            checkout_index.remove_untracked(true);
+        }
 
-    Ok(())
+        if !preserve_submodule_worktrees || !checkout_index_paths.is_empty() {
+            repo.checkout_index(Some(&mut index), Some(&mut checkout_index))?;
+            if preserve_submodule_worktrees {
+                remove_untracked_excluding_submodule_paths(repo)?;
+            }
+        }
+
+        Ok(())
+    })();
+    result
 }
 
 pub(crate) fn enter_edit_mode(
@@ -201,25 +273,47 @@ pub(crate) fn abort_and_return_to_workspace(
     force: bool,
     perm: &mut RepoExclusive,
 ) -> Result<()> {
-    if !force && !changes_from_initial(ctx, perm.read_permission())?.is_empty() {
+    let repo = &*ctx.git2_repo.get()?;
+    let submodule_paths = configured_submodule_paths(repo);
+    let preserve_submodule_worktrees = !submodule_paths.is_empty();
+    let changes = changes_from_initial(ctx, perm.read_permission())?;
+    if !force && !changes.is_empty() {
+        let has_submodule_or_gitlink_changes = changes
+            .iter()
+            .any(|change| is_submodule_related_path(&change.path.to_str_lossy(), &submodule_paths));
+
+        if has_submodule_or_gitlink_changes {
+            bail!(
+                "The working tree differs from the original commit, including submodule or gitlink changes. A forced abort is necessary and will revert changes made during edit mode.\nIf you are seeing this message, please report it as a bug. The UI should have prevented this line getting hit."
+            );
+        }
+
         bail!(
             "The working tree differs from the original commit. A forced abort is necessary.\nIf you are seeing this message, please report it as a bug. The UI should have prevented this line getting hit."
         );
     }
 
-    let repo = &*ctx.git2_repo.get()?;
-
     // Checkout gitbutler workspace branch
     repo.set_head(WORKSPACE_BRANCH_REF)
         .context("Failed to set head reference")?;
 
     let uncommited_changes = get_uncommited_changes(ctx)?;
     let uncommited_changes = repo.find_tree(uncommited_changes)?;
 
-    repo.checkout_tree(
-        uncommited_changes.as_object(),
-        Some(CheckoutBuilder::new().force().remove_untracked(true)),
-    )?;
+    let mut checkout_tree = CheckoutBuilder::new();
+    checkout_tree.force();
+    if !preserve_submodule_worktrees {
+        checkout_tree.remove_untracked(true);
+    }
+    repo.checkout_tree(uncommited_changes.as_object(), Some(&mut checkout_tree))?;
+    if preserve_submodule_worktrees {
+        remove_untracked_excluding_submodule_paths(repo)?;
+    }
+
+    // Keep index in sync with HEAD after the forceful checkout.
+    let mut index = repo.index()?;
+    index.read_tree(&repo.head()?.peel_to_tree()?)?;
+    index.write()?;
 
     Ok(())
 }

crates/gitbutler-edit-mode/tests/edit_mode.rs

@@ -6,6 +6,7 @@ use gitbutler_edit_mode::commands::{
 };
 use gitbutler_operating_modes::{EDIT_BRANCH_REF, WORKSPACE_BRANCH_REF};
 use gitbutler_stack::VirtualBranchesHandle;
+use gitbutler_testsupport::run_git_at_dir;
 use tempfile::TempDir;
 
 fn command_ctx(folder: &str) -> Result<(Context, TempDir)> {
@@ -25,7 +26,6 @@ fn command_ctx(folder: &str) -> Result<(Context, TempDir)> {
 fn conficted_entries_get_written_when_leaving_edit_mode() -> Result<()> {
     let (mut ctx, _tempdir) = command_ctx("conficted_entries_get_written_when_leaving_edit_mode")?;
     let repo = ctx.git2_repo.get()?;
-
     let foobar = repo.head()?.peel_to_commit()?.parent(0)?.id();
 
     let vb_state = VirtualBranchesHandle::new(ctx.project_data_dir());
@@ -87,6 +87,8 @@ fn abort_requires_force_when_changes_were_made() -> Result<()> {
     drop(repo);
 
     std::fs::write(worktree_dir.join("file"), "edited during edit mode\n")?;
+    let untracked_path = worktree_dir.join("new-untracked-during-edit-mode.txt");
+    std::fs::write(&untracked_path, "temporary file\n")?;
 
     let result = abort_and_return_to_workspace(&mut ctx, false);
     assert!(result.is_err());
@@ -102,6 +104,155 @@ fn abort_requires_force_when_changes_were_made() -> Result<()> {
         ctx.git2_repo.get()?.head()?.name(),
         Some(WORKSPACE_BRANCH_REF)
     );
+    assert!(
+        !untracked_path.exists(),
+        "forced abort should clean untracked files in non-submodule repos"
+    );
+
+    Ok(())
+}
+
+#[test]
+fn save_and_return_to_workspace_preserves_submodule_worktree() -> Result<()> {
+    let (mut ctx, _tempdir) =
+        command_ctx("save_and_return_to_workspace_preserves_submodule_worktree")?;
+    let (foobar, worktree_dir) = {
+        let repo = ctx.git2_repo.get()?;
+        let foobar = repo.head()?.peel_to_commit()?.parent(0)?.id();
+        (foobar, repo.path().parent().unwrap().to_path_buf())
+    };
+    let vb_state = VirtualBranchesHandle::new(ctx.project_data_dir());
+    let stacks = vb_state.list_stacks_in_workspace()?;
+    let stack = stacks.first().unwrap();
+    let submodule_probe = worktree_dir.join("submodules/test-module/probe.txt");
+    assert!(
+        submodule_probe.exists(),
+        "fixture should start with populated submodule worktree"
+    );
+
+    enter_edit_mode(&mut ctx, foobar, stack.id)?;
+    assert!(
+        submodule_probe.exists(),
+        "submodule file should remain after entering edit mode"
+    );
+    save_and_return_to_workspace(&mut ctx)?;
+
+    assert!(
+        submodule_probe.exists(),
+        "submodule file should remain after leaving edit mode"
+    );
+
+    Ok(())
+}
+
+#[test]
+fn abort_preserves_preexisting_dirty_and_diverged_submodule_state() -> Result<()> {
+    let (mut ctx, _tempdir) =
+        command_ctx("save_and_return_to_workspace_preserves_submodule_worktree")?;
+    let (foobar, worktree_dir) = {
+        let repo = ctx.git2_repo.get()?;
+        let foobar = repo.head()?.peel_to_commit()?.parent(0)?.id();
+        (foobar, repo.path().parent().unwrap().to_path_buf())
+    };
+    let submodule_dir = worktree_dir.join("submodules/test-module");
+
+    run_git_at_dir(&submodule_dir, &["config", "user.name", "Submodule Author"])?;
+    run_git_at_dir(
+        &submodule_dir,
+        &["config", "user.email", "submodule@example.com"],
+    )?;
+    std::fs::write(submodule_dir.join("diverged.txt"), "diverged commit\n")?;
+    run_git_at_dir(&submodule_dir, &["add", "diverged.txt"])?;
+    run_git_at_dir(&submodule_dir, &["commit", "-m", "local diverged commit"])?;
+
+    std::fs::write(submodule_dir.join("dirty.txt"), "dirty worktree change\n")?;
+
+    let baseline_submodule_head = run_git_at_dir(&submodule_dir, &["rev-parse", "HEAD"])?;
+    let baseline_submodule_status = run_git_at_dir(&submodule_dir, &["status", "--porcelain"])?;
+    let baseline_superproject_submodule_status = run_git_at_dir(
+        &worktree_dir,
+        &["status", "--porcelain", "submodules/test-module"],
+    )?;
+
+    let vb_state = VirtualBranchesHandle::new(ctx.project_data_dir());
+    let stacks = vb_state.list_stacks_in_workspace()?;
+    let stack = stacks.first().unwrap();
+
+    enter_edit_mode(&mut ctx, foobar, stack.id)?;
+    abort_and_return_to_workspace(&mut ctx, true)?;
+
+    let final_submodule_head = run_git_at_dir(&submodule_dir, &["rev-parse", "HEAD"])?;
+    let final_submodule_status = run_git_at_dir(&submodule_dir, &["status", "--porcelain"])?;
+    let final_superproject_submodule_status = run_git_at_dir(
+        &worktree_dir,
+        &["status", "--porcelain", "submodules/test-module"],
+    )?;
+
+    assert_eq!(
+        final_submodule_head, baseline_submodule_head,
+        "abort should preserve pre-existing submodule commit divergence"
+    );
+    assert_eq!(
+        final_submodule_status, baseline_submodule_status,
+        "abort should preserve pre-existing dirty submodule working tree state"
+    );
+    assert_eq!(
+        final_superproject_submodule_status, baseline_superproject_submodule_status,
+        "abort should restore the same superproject-visible submodule state"
+    );
+
+    Ok(())
+}
+
+#[test]
+fn abort_requires_force_warns_about_submodule_and_gitlink_reversion() -> Result<()> {
+    let (mut ctx, _tempdir) =
+        command_ctx("save_and_return_to_workspace_preserves_submodule_worktree")?;
+    let (foobar, worktree_dir) = {
+        let repo = ctx.git2_repo.get()?;
+        let foobar = repo.head()?.peel_to_commit()?.parent(0)?.id();
+        (foobar, repo.path().parent().unwrap().to_path_buf())
+    };
+    let submodule_dir = worktree_dir.join("submodules/test-module");
+
+    let vb_state = VirtualBranchesHandle::new(ctx.project_data_dir());
+    let stacks = vb_state.list_stacks_in_workspace()?;
+    let stack = stacks.first().unwrap();
+
+    enter_edit_mode(&mut ctx, foobar, stack.id)?;
+
+    run_git_at_dir(&submodule_dir, &["config", "user.name", "Submodule Author"])?;
+    run_git_at_dir(
+        &submodule_dir,
+        &["config", "user.email", "submodule@example.com"],
+    )?;
+    std::fs::write(
+        submodule_dir.join("during-edit-mode-gitlink-change.txt"),
+        "changed during edit mode\n",
+    )?;
+    run_git_at_dir(
+        &submodule_dir,
+        &["add", "during-edit-mode-gitlink-change.txt"],
+    )?;
+    run_git_at_dir(
+        &submodule_dir,
+        &["commit", "-m", "gitlink change during edit mode"],
+    )?;
+
+    // Stage the submodule entry so superproject tree diff includes the gitlink change.
+    run_git_at_dir(&worktree_dir, &["add", "submodules/test-module"])?;
+
+    let result = abort_and_return_to_workspace(&mut ctx, false);
+    assert!(result.is_err());
+    let err = result.err().unwrap().to_string();
+    assert!(
+        err.contains("including submodule or gitlink changes"),
+        "expected submodule/gitlink warning in abort error, got: {err}"
+    );
+    assert!(
+        err.contains("will revert changes made during edit mode"),
+        "expected explicit reversion warning in abort error, got: {err}"
+    );
 
     Ok(())
 }