From d5fe56d3fbc39eb97470c85241ee88601fda5ec0 Mon Sep 17 00:00:00 2001
From: Richard Jennings <richardjennings@gmail.com>
Date: Mon, 28 Oct 2024 09:18:23 +0000
Subject: [PATCH] git restore --staged

---
 index.go            | 11 +++++++++++
 integration_test.go | 34 +++++++++++++++++++++++++++++++---
 restore.go          | 42 +++++++++++++++++++++++++++++++++++-------
 3 files changed, 77 insertions(+), 10 deletions(-)

diff --git a/index.go b/index.go
index a73fb5d..b7de9f8 100644
--- a/index.go
+++ b/index.go
@@ -169,6 +169,17 @@ func (idx *Index) addFromWorkTree(f *FileStatus) error {
 	return nil
 }
 
+func (idx *Index) upsertItem(item *indexItem) error {
+	for k, v := range idx.items {
+		if bytes.Equal(v.Name, item.Name) {
+			idx.items[k] = item
+			return nil
+		}
+	}
+	idx.items = append(idx.items, item)
+	return nil
+}
+
 func (idx *Index) updateItem(i *indexItem) error {
 	found := false
 	for k, v := range idx.items {
diff --git a/integration_test.go b/integration_test.go
index 575cd1e..67e171f 100644
--- a/integration_test.go
+++ b/integration_test.go
@@ -146,8 +146,30 @@ func Test_Integration(t *testing.T) {
 		assertStatus(t, map[string]IndexStatus{"c": AddedInIndex}, map[string]WDStatus{"c": IndexAndWorkingTreeMatch})
 	}
 
-	// check for semantics when switching branch
-	//
+	// restore a file
+	{
+		// commit change to 'c'
+		_ = assertCreateCommit(t, &Commit{
+			Author:        fmt.Sprintf("%s <%s>", "tester", "tester@test.com"),
+			AuthoredTime:  time.Now(),
+			Committer:     fmt.Sprintf("%s <%s>", "tester", "tester@test.com"),
+			CommittedTime: time.Now(),
+			Message:       []byte("this is yet another commit message"),
+		})
+		// change 'c'
+		e(os.WriteFile(filepath.Join(dir, "c"), []byte("cc"), 0644), t)
+		// check status is correct
+		assertStatus(t, map[string]IndexStatus{"c": NotUpdated}, map[string]WDStatus{"c": WorktreeChangedSinceIndex})
+		// add to index
+		assertAddFiles(t, []string{"c"})
+		// check status
+		assertStatus(t, map[string]IndexStatus{"c": UpdatedInIndex}, map[string]WDStatus{"c": IndexAndWorkingTreeMatch})
+		// restore --staged c
+		assertRestore(t, "c", true)
+		// check status
+		assertStatus(t, map[string]IndexStatus{"c": NotUpdated}, map[string]WDStatus{"c": WorktreeChangedSinceIndex})
+
+	}
 
 }
 
@@ -158,6 +180,12 @@ func e(err error, t *testing.T) {
 	}
 }
 
+func assertRestore(t *testing.T, path string, staged bool) {
+	if err := Restore(path, staged); err != nil {
+		t.Fatal(err)
+	}
+}
+
 func assertLookupCommit(t *testing.T, sha Sha, f func(*testing.T, *Commit)) {
 	t.Helper()
 	commit, err := ReadCommit(sha)
@@ -256,7 +284,7 @@ func assertStatus(t *testing.T, i map[string]IndexStatus, w map[string]WDStatus)
 			t.Fatalf("expected file '%s'", k)
 		}
 		if f.wdStatus != v {
-			t.Errorf("expected file '%s' to have working directory status '%d' got '%d'", k, v, f.wdStatus)
+			t.Errorf("expected file '%s' to have working directory status '%s' got '%s'", k, v, f.wdStatus)
 		}
 	}
 }
diff --git a/restore.go b/restore.go
index cc74b4b..21f1210 100644
--- a/restore.go
+++ b/restore.go
@@ -8,7 +8,42 @@ import (
 	"path/filepath"
 )
 
+func RestoreStaged(path string) error {
+	// update the index with the sha from the last commit
+	idx, err := ReadIndex()
+	if err != nil {
+		return err
+	}
+	status, err := CurrentStatus()
+	if err != nil {
+		return err
+	}
+	f, ok := status.idx[path]
+	if !ok {
+		return fmt.Errorf("file %s not found in index", path)
+	}
+	item, err := newItem(f.wd.Finfo, f.commit.Sha, f.path)
+	if err != nil {
+		return err
+	}
+	// Work Tree files do not have a hash associated until they are added to the
+	// index. Therefore, we use filesystem modification time to compare work tree
+	// files against the index. We do not know the previous actual modification
+	// time from the commit - as git does not track timestamps like that outside
+	// of the local index. So - just subtract 1 from the working directory
+	// timestamp, so that the index is detected as different form the working
+	// tree.
+	item.MTimeN--
+	if err := idx.upsertItem(item); err != nil {
+		return err
+	}
+	return idx.Write()
+}
+
 func Restore(path string, staged bool) error {
+	if staged {
+		return RestoreStaged(path)
+	}
 	idx, err := ReadIndex()
 	if err != nil {
 		return err
@@ -21,13 +56,6 @@ func Restore(path string, staged bool) error {
 	if err != nil {
 		return err
 	}
-	if staged {
-		// remove file from index
-		if err := idx.Rm(path); err != nil {
-			return err
-		}
-		return idx.Write()
-	}
 
 	fileStatus, ok := currentStatus.Contains(path)
 	// if the path not found or is untracked working directory fileStatus then error