add optimistic rendering for staging and unstaging files

8 files changed, 228 insertions, 58 deletions

diff --git a/pkg/commands/loaders/files.go b/pkg/commands/loaders/files.go
index 825cad595..db37da935 100644
--- a/pkg/commands/loaders/files.go
+++ b/pkg/commands/loaders/files.go
@@ -7,7 +7,6 @@ import (
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
 	"github.com/jesseduffield/lazygit/pkg/commands/oscommands"
 	"github.com/jesseduffield/lazygit/pkg/common"
-	"github.com/samber/lo"
 )
 
 type FileLoaderConfig interface {
@@ -54,28 +53,15 @@ func (self *FileLoader) GetStatusFiles(opts GetStatusFileOptions) []*models.File
 			self.Log.Warningf("warning when calling git status: %s", status.StatusString)
 			continue
 		}
-		change := status.Change
-		stagedChange := change[0:1]
-		unstagedChange := change[1:2]
-		untracked := lo.Contains([]string{"??", "A ", "AM"}, change)
-		hasNoStagedChanges := lo.Contains([]string{" ", "U", "?"}, stagedChange)
-		hasInlineMergeConflicts := lo.Contains([]string{"UU", "AA"}, change)
-		hasMergeConflicts := hasInlineMergeConflicts || lo.Contains([]string{"DD", "AU", "UA", "UD", "DU"}, change)
 
 		file := &models.File{
-			Name:                    status.Name,
-			PreviousName:            status.PreviousName,
-			DisplayString:           status.StatusString,
-			HasStagedChanges:        !hasNoStagedChanges,
-			HasUnstagedChanges:      unstagedChange != " ",
-			Tracked:                 !untracked,
-			Deleted:                 unstagedChange == "D" || stagedChange == "D",
-			Added:                   unstagedChange == "A" || untracked,
-			HasMergeConflicts:       hasMergeConflicts,
-			HasInlineMergeConflicts: hasInlineMergeConflicts,
-			Type:                    self.getFileType(status.Name),
-			ShortStatus:             change,
+			Name:          status.Name,
+			PreviousName:  status.PreviousName,
+			DisplayString: status.StatusString,
+			Type:          self.getFileType(status.Name),
 		}
+
+		models.SetStatusFields(file, status.Change)
 		files = append(files, file)
 	}
 
diff --git a/pkg/commands/models/file.go b/pkg/commands/models/file.go
index bfb40ee53..88e5a14e1 100644
--- a/pkg/commands/models/file.go
+++ b/pkg/commands/models/file.go
@@ -2,6 +2,7 @@ package models
 
 import (
 	"github.com/jesseduffield/lazygit/pkg/utils"
+	"github.com/samber/lo"
 )
 
 // File : A file from git status
@@ -90,3 +91,48 @@ func (f *File) GetPath() string {
 func (f *File) GetPreviousPath() string {
 	return f.PreviousName
 }
+
+type StatusFields struct {
+	HasStagedChanges        bool
+	HasUnstagedChanges      bool
+	Tracked                 bool
+	Deleted                 bool
+	Added                   bool
+	HasMergeConflicts       bool
+	HasInlineMergeConflicts bool
+	ShortStatus             string
+}
+
+func SetStatusFields(file *File, shortStatus string) {
+	derived := deriveStatusFields(shortStatus)
+
+	file.HasStagedChanges = derived.HasStagedChanges
+	file.HasUnstagedChanges = derived.HasUnstagedChanges
+	file.Tracked = derived.Tracked
+	file.Deleted = derived.Deleted
+	file.Added = derived.Added
+	file.HasMergeConflicts = derived.HasMergeConflicts
+	file.HasInlineMergeConflicts = derived.HasInlineMergeConflicts
+	file.ShortStatus = derived.ShortStatus
+}
+
+// shortStatus is something like '??' or 'A '
+func deriveStatusFields(shortStatus string) StatusFields {
+	stagedChange := shortStatus[0:1]
+	unstagedChange := shortStatus[1:2]
+	untracked := lo.Contains([]string{"??", "A ", "AM"}, shortStatus)
+	hasNoStagedChanges := lo.Contains([]string{" ", "U", "?"}, stagedChange)
+	hasInlineMergeConflicts := lo.Contains([]string{"UU", "AA"}, shortStatus)
+	hasMergeConflicts := hasInlineMergeConflicts || lo.Contains([]string{"DD", "AU", "UA", "UD", "DU"}, shortStatus)
+
+	return StatusFields{
+		HasStagedChanges:        !hasNoStagedChanges,
+		HasUnstagedChanges:      unstagedChange != " ",
+		Tracked:                 !untracked,
+		Deleted:                 unstagedChange == "D" || stagedChange == "D",
+		Added:                   unstagedChange == "A" || untracked,
+		HasMergeConflicts:       hasMergeConflicts,
+		HasInlineMergeConflicts: hasInlineMergeConflicts,
+		ShortStatus:             shortStatus,
+	}
+}
diff --git a/pkg/gui/controllers.go b/pkg/gui/controllers.go
index 2c851739d..dcb3377cf 100644
--- a/pkg/gui/controllers.go
+++ b/pkg/gui/controllers.go
@@ -62,6 +62,7 @@ func (gui *Gui) resetControllers() {
 		model,
 		gui.State.Contexts,
 		gui.State.Modes,
+		&gui.Mutexes,
 	)
 
 	syncController := controllers.NewSyncController(
diff --git a/pkg/gui/controllers/common.go b/pkg/gui/controllers/common.go
index 55ba4b176..12a3788fd 100644
--- a/pkg/gui/controllers/common.go
+++ b/pkg/gui/controllers/common.go
@@ -16,6 +16,7 @@ type controllerCommon struct {
 	model    *types.Model
 	contexts *context.ContextTree
 	modes    *types.Modes
+	mutexes  *types.Mutexes
 }
 
 func NewControllerCommon(
@@ -26,6 +27,7 @@ func NewControllerCommon(
 	model *types.Model,
 	contexts *context.ContextTree,
 	modes *types.Modes,
+	mutexes *types.Mutexes,
 ) *controllerCommon {
 	return &controllerCommon{
 		c:        c,
@@ -35,5 +37,6 @@ func NewControllerCommon(
 		model:    model,
 		contexts: contexts,
 		modes:    modes,
+		mutexes:  mutexes,
 	}
 }
diff --git a/pkg/gui/controllers/files_controller.go b/pkg/gui/controllers/files_controller.go
index 4aed5cf0a..220891a4d 100644
--- a/pkg/gui/controllers/files_controller.go
+++ b/pkg/gui/controllers/files_controller.go
@@ -108,7 +108,7 @@ func (self *FilesController) GetKeybindings(opts types.KeybindingsOpts) []*types
 		},
 		{
 			Key:         opts.GetKey(opts.Config.Files.ToggleStagedAll),
-			Handler:     self.stageAll,
+			Handler:     self.toggleStagedAll,
 			Description: self.c.Tr.LcToggleStagedAll,
 		},
 		{
@@ -167,7 +167,86 @@ func (self *FilesController) GetOnClick() func() error {
 	return self.checkSelectedFileNode(self.press)
 }
 
-func (self *FilesController) press(node *filetree.FileNode) error {
+// if we are dealing with a status for which there is no key in this map,
+// then we won't optimistically render: we'll just let `git status` tell
+// us what the new status is.
+// There are no doubt more entries that could be added to these two maps.
+var stageStatusMap = map[string]string{
+	"??": "A ",
+	" M": "M ",
+	"MM": "M ",
+	" D": "D ",
+	" A": "A ",
+	"AM": "A ",
+	"MD": "D ",
+}
+
+var unstageStatusMap = map[string]string{
+	"A ": "??",
+	"M ": " M",
+	"D ": " D",
+}
+
+func (self *FilesController) optimisticStage(file *models.File) bool {
+	newShortStatus, ok := stageStatusMap[file.ShortStatus]
+	if !ok {
+		return false
+	}
+
+	models.SetStatusFields(file, newShortStatus)
+	return true
+}
+
+func (self *FilesController) optimisticUnstage(file *models.File) bool {
+	newShortStatus, ok := unstageStatusMap[file.ShortStatus]
+	if !ok {
+		return false
+	}
+
+	models.SetStatusFields(file, newShortStatus)
+	return true
+}
+
+// Running a git add command followed by a git status command can take some time (e.g. 200ms).
+// Given how often users stage/unstage files in Lazygit, we're adding some
+// optimistic rendering to make things feel faster. When we go to stage
+// a file, we'll first update that file's status in-memory, then re-render
+// the files panel. Then we'll immediately do a proper git status call
+// so that if the optimistic rendering got something wrong, it's quickly
+// corrected.
+func (self *FilesController) optimisticChange(node *filetree.FileNode, optimisticChangeFn func(*models.File) bool) error {
+	rerender := false
+	err := node.ForEachFile(func(f *models.File) error {
+		// can't act on the file itself: we need to update the original model file
+		for _, modelFile := range self.model.Files {
+			if modelFile.Name == f.Name {
+				if optimisticChangeFn(modelFile) {
+					rerender = true
+				}
+				break
+			}
+		}
+
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+	if rerender {
+		if err := self.c.PostRefreshUpdate(self.contexts.Files); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (self *FilesController) pressWithLock(node *filetree.FileNode) error {
+	// Obtaining this lock because optimistic rendering requires us to mutate
+	// the files in our model.
+	self.mutexes.RefreshingFilesMutex.Lock()
+	defer self.mutexes.RefreshingFilesMutex.Unlock()
+
 	if node.IsLeaf() {
 		file := node.File
 
@@ -177,11 +256,21 @@ func (self *FilesController) press(node *filetree.FileNode) error {
 
 		if file.HasUnstagedChanges {
 			self.c.LogAction(self.c.Tr.Actions.StageFile)
+
+			if err := self.optimisticChange(node, self.optimisticStage); err != nil {
+				return err
+			}
+
 			if err := self.git.WorkingTree.StageFile(file.Name); err != nil {
 				return self.c.Error(err)
 			}
 		} else {
 			self.c.LogAction(self.c.Tr.Actions.UnstageFile)
+
+			if err := self.optimisticChange(node, self.optimisticUnstage); err != nil {
+				return err
+			}
+
 			if err := self.git.WorkingTree.UnStageFile(file.Names(), file.Tracked); err != nil {
 				return self.c.Error(err)
 			}
@@ -195,19 +284,37 @@ func (self *FilesController) press(node *filetree.FileNode) error {
 
 		if node.GetHasUnstagedChanges() {
 			self.c.LogAction(self.c.Tr.Actions.StageFile)
+
+			if err := self.optimisticChange(node, self.optimisticStage); err != nil {
+				return err
+			}
+
 			if err := self.git.WorkingTree.StageFile(node.Path); err != nil {
 				return self.c.Error(err)
 			}
 		} else {
-			// pretty sure it doesn't matter that we're always passing true here
 			self.c.LogAction(self.c.Tr.Actions.UnstageFile)
+
+			if err := self.optimisticChange(node, self.optimisticUnstage); err != nil {
+				return err
+			}
+
+			// pretty sure it doesn't matter that we're always passing true here
 			if err := self.git.WorkingTree.UnStageFile([]string{node.Path}, true); err != nil {
 				return self.c.Error(err)
 			}
 		}
 	}
 
-	if err := self.c.Refresh(types.RefreshOptions{Scope: []types.RefreshableView{types.FILES}}); err != nil {
+	return nil
+}
+
+func (self *FilesController) press(node *filetree.FileNode) error {
+	if err := self.pressWithLock(node); err != nil {
+		return err
+	}
+
+	if err := self.c.Refresh(types.RefreshOptions{Scope: []types.RefreshableView{types.FILES}, Mode: types.ASYNC}); err != nil {
 		return err
 	}
 
@@ -273,33 +380,53 @@ func (self *FilesController) EnterFile(opts types.OnFocusOpts) error {
 	return self.c.PushContext(self.contexts.Staging, opts)
 }
 
-func (self *FilesController) allFilesStaged() bool {
-	for _, file := range self.model.Files {
-		if file.HasUnstagedChanges {
-			return false
-		}
+func (self *FilesController) toggleStagedAll() error {
+	if err := self.toggleStagedAllWithLock(); err != nil {
+		return err
 	}
-	return true
-}
 
-func (self *FilesController) stageAll() error {
-	var err error
-	if self.allFilesStaged() {
-		self.c.LogAction(self.c.Tr.Actions.UnstageAllFiles)
-		err = self.git.WorkingTree.UnstageAll()
-	} else {
-		self.c.LogAction(self.c.Tr.Actions.StageAllFiles)
-		err = self.git.WorkingTree.StageAll()
+	if err := self.c.Refresh(types.RefreshOptions{Scope: []types.RefreshableView{types.FILES}, Mode: types.ASYNC}); err != nil {
+		return err
 	}
-	if err != nil {
-		_ = self.c.Error(err)
+
+	return self.context().HandleFocus()
+}
+
+func (self *FilesController) toggleStagedAllWithLock() error {
+	self.mutexes.RefreshingFilesMutex.Lock()
+	defer self.mutexes.RefreshingFilesMutex.Unlock()
+
+	root := self.context().FileTreeViewModel.GetRoot()
+
+	// if any files within have inline merge conflicts we can't stage or unstage,
+	// or it'll end up with those >>>>>> lines actually staged
+	if root.GetHasInlineMergeConflicts() {
+		return self.c.ErrorMsg(self.c.Tr.ErrStageDirWithInlineMergeConflicts)
 	}
 
-	if err := self.c.Refresh(types.RefreshOptions{Scope: []types.RefreshableView{types.FILES}}); err != nil {
-		return err
+	if root.GetHasUnstagedChanges() {
+		self.c.LogAction(self.c.Tr.Actions.StageAllFiles)
+
+		if err := self.optimisticChange(root, self.optimisticStage); err != nil {
+			return err
+		}
+
+		if err := self.git.WorkingTree.StageAll(); err != nil {
+			return self.c.Error(err)
+		}
+	} else {
+		self.c.LogAction(self.c.Tr.Actions.UnstageAllFiles)
+
+		if err := self.optimisticChange(root, self.optimisticUnstage); err != nil {
+			return err
+		}
+
+		if err := self.git.WorkingTree.UnstageAll(); err != nil {
+			return self.c.Error(err)
+		}
 	}
 
-	return self.contexts.Files.HandleFocus()
+	return nil
 }
 
 func (self *FilesController) unstageFiles(node *filetree.FileNode) error {
diff --git a/pkg/gui/filetree/file_tree.go b/pkg/gui/filetree/file_tree.go
index d4bb8e596..17dafaa5d 100644
--- a/pkg/gui/filetree/file_tree.go
+++ b/pkg/gui/filetree/file_tree.go
@@ -41,6 +41,7 @@ type IFileTree interface {
 	GetAllItems() []*FileNode
 	GetAllFiles() []*models.File
 	GetFilter() FileTreeDisplayFilter
+	GetRoot() *FileNode
 }
 
 type FileTree struct {
@@ -159,6 +160,10 @@ func (self *FileTree) Tree() INode {
 	return self.tree
 }
 
+func (self *FileTree) GetRoot() *FileNode {
+	return self.tree
+}
+
 func (self *FileTree) CollapsedPaths() *CollapsedPaths {
 	return self.collapsedPaths
 }
diff --git a/pkg/gui/gui.go b/pkg/gui/gui.go
index 531a43035..04b53e074 100644
--- a/pkg/gui/gui.go
+++ b/pkg/gui/gui.go
@@ -97,7 +97,7 @@ type Gui struct {
 	// recent repo with the recent repos popup showing
 	showRecentRepos bool
 
-	Mutexes guiMutexes
+	Mutexes types.Mutexes
 
 	// findSuggestions will take a string that the user has typed into a prompt
 	// and return a slice of suggestions which match that string.
@@ -237,18 +237,6 @@ const (
 	COMPLETE
 )
 
-// if you add a new mutex here be sure to instantiate it. We're using pointers to
-// mutexes so that we can pass the mutexes to controllers.
-type guiMutexes struct {
-	RefreshingFilesMutex  *sync.Mutex
-	RefreshingStatusMutex *sync.Mutex
-	SyncMutex             *sync.Mutex
-	LocalCommitsMutex     *sync.Mutex
-	LineByLinePanelMutex  *sync.Mutex
-	SubprocessMutex       *sync.Mutex
-	PopupMutex            *sync.Mutex
-}
-
 func (gui *Gui) onNewRepo(startArgs types.StartArgs, reuseState bool) error {
 	var err error
 	gui.git, err = commands.NewGitCommand(
@@ -418,7 +406,7 @@ func NewGui(
 		// but now we do it via state. So we need to still support the config for the
 		// sake of backwards compatibility. We're making use of short circuiting here
 		ShowExtrasWindow: cmn.UserConfig.Gui.ShowCommandLog && !config.GetAppState().HideCommandLog,
-		Mutexes: guiMutexes{
+		Mutexes: types.Mutexes{
 			RefreshingFilesMutex:  &sync.Mutex{},
 			RefreshingStatusMutex: &sync.Mutex{},
 			SyncMutex:             &sync.Mutex{},
diff --git a/pkg/gui/types/common.go b/pkg/gui/types/common.go
index 72bba0b6c..21808705a 100644
--- a/pkg/gui/types/common.go
+++ b/pkg/gui/types/common.go
@@ -1,6 +1,8 @@
 package types
 
 import (
+	"sync"
+
 	"github.com/jesseduffield/lazygit/pkg/commands/git_commands"
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
 	"github.com/jesseduffield/lazygit/pkg/commands/oscommands"
@@ -156,3 +158,15 @@ type Model struct {
 	// for displaying suggestions while typing in a file name
 	FilesTrie *patricia.Trie
 }
+
+// if you add a new mutex here be sure to instantiate it. We're using pointers to
+// mutexes so that we can pass the mutexes to controllers.
+type Mutexes struct {
+	RefreshingFilesMutex  *sync.Mutex
+	RefreshingStatusMutex *sync.Mutex
+	SyncMutex             *sync.Mutex
+	LocalCommitsMutex     *sync.Mutex
+	LineByLinePanelMutex  *sync.Mutex
+	SubprocessMutex       *sync.Mutex
+	PopupMutex            *sync.Mutex
+}