1 files changed, 21 insertions, 30 deletions

diff --git a/worktree_status.go b/worktree_status.go
index 6becada..e6d1745 100644
--- a/worktree_status.go
+++ b/worktree_status.go
@@ -216,52 +216,43 @@ 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
 	}
 
-	for i, e := range idx.Entries {
-		if e.Name != filename {
-			continue
-		}
-
-		e.Hash = h
-		e.ModifiedAt = info.ModTime()
-		e.Mode, err = filemode.NewFromOSFileMode(info.Mode())
-		if err != nil {
-			return err
-		}
-
-		fillSystemInfo(&e, info.Sys())
-		idx.Entries[i] = e
+	e.Hash = h
+	e.ModifiedAt = info.ModTime()
+	e.Mode, err = filemode.NewFromOSFileMode(info.Mode())
+	if err != nil {
+		return err
 	}
 
+	fillSystemInfo(e, info.Sys())
 	return nil
 }