Discard with support for whole hunks.

Whole chunks match the worktree change completely.

Author: Byron

crates/but-cli/src/command/mod.rs

@@ -151,6 +151,7 @@ pub(crate) fn discard_change(
     debug_print(but_workspace::discard_workspace_changes(
         &repo,
         Some(spec.into()),
+        UI_CONTEXT_LINES,
     )?)
 }
 

crates/but-core/src/diff/mod.rs

@@ -80,6 +80,19 @@ impl ModeFlags {
     }
 }
 
+impl ModeFlags {
+    /// Returns `true` if this instance indicates a type-change.
+    /// The only reason this isn't the case is if the executable bit changed.
+    pub fn is_typechange(&self) -> bool {
+        match self {
+            ModeFlags::ExecutableBitAdded | ModeFlags::ExecutableBitRemoved => false,
+            ModeFlags::TypeChangeFileToLink
+            | ModeFlags::TypeChangeLinkToFile
+            | ModeFlags::TypeChange => true,
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     mod flags {

crates/but-core/src/diff/worktree.rs

@@ -805,7 +805,7 @@ impl TreeChange {
         let mut diff_filter = crate::unified_diff::filter_from_state(
             repo,
             self.status.state(),
-            gix::diff::blob::pipeline::Mode::ToGitUnlessBinaryToTextIsPresent,
+            UnifiedDiff::CONVERSION_MODE,
         )?;
         self.unified_diff_with_filter(repo, context_lines, &mut diff_filter)
     }

crates/but-workspace/src/commit_engine/mod.rs

@@ -13,7 +13,7 @@ use gix::prelude::ObjectIdExt as _;
 use gix::refs::transaction::PreviousValue;
 use serde::{Deserialize, Serialize};
 
-mod tree;
+pub(crate) mod tree;
 use tree::{CreateTreeOutcome, create_tree};
 
 pub(crate) mod index;
@@ -94,7 +94,7 @@ pub struct DiffSpec {
 }
 
 /// The header of a hunk that represents a change to a file.
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct HunkHeader {
     /// The 1-based line number at which the previous version of the file started.

crates/but-workspace/src/commit_engine/tree.rs

@@ -2,13 +2,14 @@ use crate::commit_engine::{
     Destination, DiffSpec, HunkHeader, MoveSourceCommit, RejectionReason, apply_hunks,
 };
 use anyhow::bail;
-use bstr::ByteSlice;
+use bstr::{BStr, ByteSlice};
 use but_core::{RepositoryExt, UnifiedDiff};
 use gix::filter::plumbing::pipeline::convert::ToGitOutcome;
 use gix::merge::tree::TreatAsUnresolved;
 use gix::object::tree::EntryKind;
 use gix::prelude::ObjectIdExt;
 use std::io::Read;
+use std::path::Path;
 
 /// Additional information about the outcome of a [`create_tree()`] call.
 #[derive(Debug)]
@@ -222,6 +223,7 @@ fn apply_worktree_changes<'repo>(
                 gix::diff::blob::pipeline::Mode::ToGitUnlessBinaryToTextIsPresent,
                 "BUG: if this changes, the uses of worktree filters need a review"
             );
+            // TODO(perf): avoid computing the unified diff here, we only need hunks with, usually with zero context.
             let UnifiedDiff::Patch { hunks } =
                 worktree_change.unified_diff_with_filter(repo, context_lines, &mut diff_filter)?
             else {
@@ -281,22 +283,13 @@ fn apply_worktree_changes<'repo>(
                 continue;
             };
 
-            current_worktree.clear();
-            let to_git = pipeline.convert_to_git(
-                std::fs::File::open(path)?,
-                &gix::path::from_bstr(base_rela_path),
+            worktree_file_to_git_in_buf(
+                &mut current_worktree,
+                base_rela_path,
+                &path,
+                &mut pipeline,
                 &index,
             )?;
-            match to_git {
-                ToGitOutcome::Unchanged(mut file) => {
-                    file.read_to_end(&mut current_worktree)?;
-                }
-                ToGitOutcome::Process(mut stream) => {
-                    stream.read_to_end(&mut current_worktree)?;
-                }
-                ToGitOutcome::Buffer(buf) => current_worktree.extend_from_slice(buf),
-            };
-
             let worktree_hunks: Vec<HunkHeader> = hunks.into_iter().map(Into::into).collect();
             for selected_hunk in &change_request.hunk_headers {
                 if !worktree_hunks.contains(selected_hunk)
@@ -327,3 +320,28 @@ fn apply_worktree_changes<'repo>(
     let maybe_new_tree = (actual_base_tree != altered_base_tree_id).then_some(altered_base_tree_id);
     Ok((maybe_new_tree, actual_base_tree))
 }
+
+pub(crate) fn worktree_file_to_git_in_buf(
+    buf: &mut Vec<u8>,
+    rela_path: &BStr,
+    path: &Path,
+    pipeline: &mut gix::filter::Pipeline<'_>,
+    index: &gix::index::State,
+) -> anyhow::Result<()> {
+    buf.clear();
+    let to_git = pipeline.convert_to_git(
+        std::fs::File::open(path)?,
+        &gix::path::from_bstr(rela_path),
+        index,
+    )?;
+    match to_git {
+        ToGitOutcome::Unchanged(mut file) => {
+            file.read_to_end(buf)?;
+        }
+        ToGitOutcome::Process(mut stream) => {
+            stream.read_to_end(buf)?;
+        }
+        ToGitOutcome::Buffer(buf2) => buf.extend_from_slice(buf2),
+    };
+    Ok(())
+}

crates/but-workspace/src/discard/function.rs

@@ -1,14 +1,18 @@
-use crate::discard::{DiscardSpec, file};
-use anyhow::Context;
+use crate::discard::{DiscardSpec, file, hunk};
+use anyhow::{Context, bail};
 use bstr::ByteSlice;
 use but_core::{ChangeState, TreeStatus};
 
 /// Discard the given `changes` in the worktree of `repo`. If a change could not be matched with an actual worktree change, for
-/// instance due to a race, that's not an error, instead it will be returned in the result Vec.
+/// instance due to a race, that's not an error, instead it will be returned in the result Vec, along with all hunks that couldn't
+/// be matched.
 /// The returned Vec is typically empty, meaning that all `changes` could be discarded.
 ///
+/// `context_lines` is the amount of context lines we should assume when obtaining hunks of worktree changes to match against
+/// the ones we have specified in the hunks contained within `changes`.
+///
 /// Discarding a change is really more of an 'undo' of a change as it will restore the previous state to the desired extent - Git
-/// doesn't have a notion of this.
+/// doesn't have a notion of this on a whole-file basis.
 ///
 /// Each of the `changes` will be matched against actual worktree changes to make this operation as safe as possible, after all, it
 /// discards changes without recovery.
@@ -19,6 +23,7 @@ use but_core::{ChangeState, TreeStatus};
 pub fn discard_workspace_changes(
     repo: &gix::Repository,
     changes: impl IntoIterator<Item = DiscardSpec>,
+    context_lines: u32,
 ) -> anyhow::Result<Vec<DiscardSpec>> {
     let wt_changes = but_core::diff::worktree_changes(repo)?;
     let mut dropped = Vec::new();
@@ -30,7 +35,7 @@ pub fn discard_workspace_changes(
     let mut path_check = gix::status::plumbing::SymlinkCheck::new(
         repo.workdir().context("non-bare repository")?.into(),
     );
-    for spec in changes {
+    for mut spec in changes {
         let Some(wt_change) = wt_changes.changes.iter().find(|c| {
             c.path == spec.path
                 && c.previous_path() == spec.previous_path.as_ref().map(|p| p.as_bstr())
@@ -115,7 +120,43 @@ pub fn discard_workspace_changes(
                 }
             }
         } else {
-            todo!("hunk-based undo")
+            match wt_change.status {
+                TreeStatus::Addition { .. } | TreeStatus::Deletion { .. } => {
+                    bail!(
+                        "Deletions or additions aren't well-defined for hunk-based operations - use the whole-file mode instead: '{}'",
+                        wt_change.path
+                    )
+                }
+                TreeStatus::Modification {
+                    previous_state,
+                    flags,
+                    ..
+                }
+                | TreeStatus::Rename {
+                    previous_state,
+                    flags,
+                    ..
+                } => {
+                    if flags.is_some_and(|f| f.is_typechange()) {
+                        bail!(
+                            "Type-changed items can't be discard by hunks - use the whole-file mode isntead"
+                        )
+                    }
+                    hunk::restore_state_to_worktree(
+                        wt_change,
+                        previous_state,
+                        &mut spec.hunk_headers,
+                        &mut path_check,
+                        &mut pipeline,
+                        &index,
+                        context_lines,
+                    )?;
+                    if !spec.hunk_headers.is_empty() {
+                        dropped.push(spec);
+                        continue;
+                    }
+                }
+            }
         }
     }
 

crates/but-workspace/src/discard/hunk.rs

@@ -0,0 +1,93 @@
+use crate::commit_engine::tree::worktree_file_to_git_in_buf;
+use crate::commit_engine::{HunkHeader, apply_hunks};
+use anyhow::bail;
+use bstr::ByteSlice;
+use but_core::{ChangeState, TreeChange, UnifiedDiff};
+use gix::filter::plumbing::driver::apply::{Delay, MaybeDelayed};
+use gix::filter::plumbing::pipeline::convert::ToWorktreeOutcome;
+use gix::prelude::ObjectIdExt;
+
+/// Discard `hunks_to_discard` in the resource at `wt_change`, whose previous version is `previous_state` and is expected to
+/// be tracked and readable from the object database. We will always read what's currently on disk as the current version
+/// and overwrite it.
+///
+/// The general idea is to rebuild the file, but without the `hunks_to_discard`, write it into the worktree replacing
+/// the previous file. `hunks_to_discard` are hunks based on a diff of what's in Git.
+///
+/// Note that an index update isn't necessary, as for the purpose of the application it might as well not exist due to the
+/// way our worktree-changes function ignores the index entirely.
+///
+/// Note that `hunks_to_discard` that weren't present in the actual set of worktree changes will remain, so when everything
+/// worked `hunks_to_discard` will be empty.
+pub fn restore_state_to_worktree(
+    wt_change: &TreeChange,
+    previous_state: ChangeState,
+    hunks_to_discard: &mut Vec<HunkHeader>,
+    path_check: &mut gix::status::plumbing::SymlinkCheck,
+    pipeline: &mut gix::filter::Pipeline<'_>,
+    index: &gix::index::State,
+    context_lines: u32,
+) -> anyhow::Result<()> {
+    let repo = pipeline.repo;
+    let state_in_worktree = ChangeState {
+        id: repo.object_hash().null(),
+        kind: previous_state.kind,
+    };
+    let mut diff_filter = but_core::unified_diff::filter_from_state(
+        repo,
+        Some(state_in_worktree),
+        UnifiedDiff::CONVERSION_MODE,
+    )?;
+    let UnifiedDiff::Patch { hunks } =
+        wt_change.unified_diff_with_filter(repo, context_lines, &mut diff_filter)?
+    else {
+        bail!("Couldn't obtain diff for worktree changes.")
+    };
+
+    let prev_len = hunks_to_discard.len();
+    let hunks_to_keep: Vec<HunkHeader> = hunks
+        .into_iter()
+        .map(Into::into)
+        .filter(|hunk| {
+            match hunks_to_discard
+                .iter()
+                .enumerate()
+                .find_map(|(idx, hunk_to_discard)| (hunk_to_discard == hunk).then_some(idx))
+            {
+                None => true,
+                Some(idx_to_remove) => {
+                    hunks_to_discard.remove(idx_to_remove);
+                    false
+                }
+            }
+        })
+        .collect();
+    if prev_len == hunks_to_discard.len() {
+        // We want to keep everything, so nothing has to be done.
+        return Ok(());
+    }
+
+    let old = previous_state.id.attach(repo).object()?.detach().data;
+    let mut new = repo.empty_reusable_buffer();
+    let rela_path = wt_change.path.as_bstr();
+    let worktree_path = path_check.verified_path_allow_nonexisting(rela_path)?;
+    worktree_file_to_git_in_buf(&mut new, rela_path, &worktree_path, pipeline, index)?;
+
+    let base_with_patches = apply_hunks(old.as_bstr(), new.as_bstr(), &hunks_to_keep)?;
+
+    let to_worktree = pipeline.convert_to_worktree(&base_with_patches, rela_path, Delay::Forbid)?;
+    match to_worktree {
+        ToWorktreeOutcome::Unchanged(buf) => {
+            std::fs::write(&worktree_path, buf)?;
+        }
+        ToWorktreeOutcome::Buffer(buf) => {
+            std::fs::write(&worktree_path, buf)?;
+        }
+        ToWorktreeOutcome::Process(MaybeDelayed::Immediate(mut stream)) => {
+            let mut file = std::fs::File::create(&worktree_path)?;
+            std::io::copy(&mut stream, &mut file)?;
+        }
+        ToWorktreeOutcome::Process(MaybeDelayed::Delayed(_)) => unreachable!("disabled"),
+    }
+    Ok(())
+}

crates/but-workspace/src/discard/mod.rs

@@ -2,7 +2,7 @@
 
 use crate::commit_engine::DiffSpec;
 use gix::object::tree::EntryKind;
-use std::ops::Deref;
+use std::ops::{Deref, DerefMut};
 use std::path::{Path, PathBuf};
 
 /// A specification of what should be discarded, either changes to the whole file, or a portion of it.
@@ -24,6 +24,12 @@ impl Deref for DiscardSpec {
     }
 }
 
+impl DerefMut for DiscardSpec {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
 impl From<DiscardSpec> for DiffSpec {
     fn from(value: DiscardSpec) -> Self {
         value.0
@@ -38,6 +44,7 @@ pub mod ui {
 }
 
 mod file;
+mod hunk;
 
 #[cfg(unix)]
 fn locked_resource_at(

crates/but-workspace/tests/fixtures/generated-archives/.gitignore

@@ -21,3 +21,5 @@
 /all-file-types-renamed-and-overwriting-existing.tar
 /move-directory-into-sibling-file.tar
 /merge-with-two-branches-conflict.tar
+/deletion-addition-untracked.tar
+/mixed-hunk-modifications.tar

crates/but-workspace/tests/fixtures/scenario/deletion-addition-untracked.sh

@@ -0,0 +1,16 @@
+#!/usr/bin/env bash
+
+### Description
+# Provide a deleted file, an added one, and an untracked one.
+set -eu -o pipefail
+
+git init
+echo content >to-be-deleted
+echo content2 >to-be-deleted-in-index
+git add . && git commit -m "init"
+
+echo new >added-to-index && git add added-to-index
+echo untracked >untracked
+rm to-be-deleted
+git rm to-be-deleted-in-index
+