worktree: Add, logic improvements

1 files changed, 15 insertions, 21 deletions

diff --git a/worktree_status.go b/worktree_status.go
index 46922da..e6d1745 100644
--- a/worktree_status.go
+++ b/worktree_status.go
@@ -216,42 +216,36 @@ func (w *Worktree) addOrUpdateFileToIndex(filename string, h plumbing.Hash) erro
 		return err
 	}
 
-	_, err = idx.Entry(filename)
-	if err == index.ErrEntryNotFound {
-		err = w.doAddFileToIndex(idx, filename)
-	}
-
-	if err != nil {
+	e, err := idx.Entry(filename)
+	if err != nil && err != index.ErrEntryNotFound {
 		return err
 	}
 
-	err = w.doUpdateFileToIndex(idx, filename, h)
-	if err != nil {
-		return err
+	if err == index.ErrEntryNotFound {
+		if err := w.doAddFileToIndex(idx, filename, h); err != nil {
+			return err
+		}
+	} else {
+		if err := w.doUpdateFileToIndex(e, filename, h); err != nil {
+			return err
+		}
 	}
-
 	return w.r.Storer.SetIndex(idx)
 }
 
-func (w *Worktree) doAddFileToIndex(idx *index.Index, filename string) error {
-	idx.Entries = append(idx.Entries, &index.Entry{
-		Name: filename,
-	})
+func (w *Worktree) doAddFileToIndex(idx *index.Index, filename string, h plumbing.Hash) error {
+	e := &index.Entry{Name: filename}
+	idx.Entries = append(idx.Entries, e)
 
-	return nil
+	return w.doUpdateFileToIndex(e, filename, h)
 }
 
-func (w *Worktree) doUpdateFileToIndex(idx *index.Index, filename string, h plumbing.Hash) error {
+func (w *Worktree) doUpdateFileToIndex(e *index.Entry, filename string, h plumbing.Hash) error {
 	info, err := w.fs.Stat(filename)
 	if err != nil {
 		return err
 	}
 
-	e, err := idx.Entry(filename)
-	if err != nil {
-		return err
-	}
-
 	e.Hash = h
 	e.ModifiedAt = info.ModTime()
 	e.Mode, err = filemode.NewFromOSFileMode(info.Mode())