migrate files context to new structure

13 files changed, 215 insertions, 152 deletions

diff --git a/pkg/gui/context/context.go b/pkg/gui/context/context.go
index 13d63b9e9..1bcefeadb 100644
--- a/pkg/gui/context/context.go
+++ b/pkg/gui/context/context.go
@@ -56,7 +56,7 @@ var AllContextKeys = []types.ContextKey{
 
 type ContextTree struct {
 	Status         types.Context
-	Files          types.IListContext
+	Files          *WorkingTreeContext
 	Submodules     types.IListContext
 	Menu           types.IListContext
 	Branches       types.IListContext
diff --git a/pkg/gui/context/tags_context.go b/pkg/gui/context/tags_context.go
index 5c2592cf9..b79e4d31f 100644
--- a/pkg/gui/context/tags_context.go
+++ b/pkg/gui/context/tags_context.go
@@ -7,7 +7,7 @@ import (
 )
 
 type TagsContext struct {
-	*TagsList
+	*TagsViewModel
 	*BaseContext
 	*ListContextTrait
 }
@@ -35,7 +35,7 @@ func NewTagsContext(
 	self := &TagsContext{}
 	takeFocus := func() error { return c.PushContext(self) }
 
-	list := NewTagsList(getModel)
+	list := NewTagsViewModel(getModel)
 	viewTrait := NewViewTrait(getView)
 	listContextTrait := &ListContextTrait{
 		base:      baseContext,
@@ -56,21 +56,21 @@ func NewTagsContext(
 
 	self.BaseContext = baseContext
 	self.ListContextTrait = listContextTrait
-	self.TagsList = list
+	self.TagsViewModel = list
 
 	return self
 }
 
-type TagsList struct {
+type TagsViewModel struct {
 	*ListTrait
 	getModel func() []*models.Tag
 }
 
-func (self *TagsList) GetItemsLength() int {
+func (self *TagsViewModel) GetItemsLength() int {
 	return len(self.getModel())
 }
 
-func (self *TagsList) GetSelectedTag() *models.Tag {
+func (self *TagsViewModel) GetSelectedTag() *models.Tag {
 	if self.GetItemsLength() == 0 {
 		return nil
 	}
@@ -78,13 +78,13 @@ func (self *TagsList) GetSelectedTag() *models.Tag {
 	return self.getModel()[self.GetSelectedLineIdx()]
 }
 
-func (self *TagsList) GetSelectedItem() (types.ListItem, bool) {
-	tag := self.GetSelectedTag()
-	return tag, tag != nil
+func (self *TagsViewModel) GetSelectedItem() (types.ListItem, bool) {
+	item := self.GetSelectedTag()
+	return item, item != nil
 }
 
-func NewTagsList(getModel func() []*models.Tag) *TagsList {
-	self := &TagsList{
+func NewTagsViewModel(getModel func() []*models.Tag) *TagsViewModel {
+	self := &TagsViewModel{
 		getModel: getModel,
 	}
 
diff --git a/pkg/gui/context/working_tree_context.go b/pkg/gui/context/working_tree_context.go
new file mode 100644
index 000000000..4df16fd11
--- /dev/null
+++ b/pkg/gui/context/working_tree_context.go
@@ -0,0 +1,101 @@
+package context
+
+import (
+	"github.com/jesseduffield/gocui"
+	"github.com/jesseduffield/lazygit/pkg/commands/models"
+	"github.com/jesseduffield/lazygit/pkg/gui/filetree"
+	"github.com/jesseduffield/lazygit/pkg/gui/types"
+	"github.com/sirupsen/logrus"
+)
+
+type WorkingTreeContext struct {
+	*WorkingTreeViewModal
+	*BaseContext
+	*ListContextTrait
+}
+
+var _ types.IListContext = (*WorkingTreeContext)(nil)
+
+func NewWorkingTreeContext(
+	getModel func() []*models.File,
+	getView func() *gocui.View,
+	getDisplayStrings func(startIdx int, length int) [][]string,
+
+	onFocus func(...types.OnFocusOpts) error,
+	onRenderToMain func(...types.OnFocusOpts) error,
+	onFocusLost func() error,
+
+	c *types.ControllerCommon,
+) *WorkingTreeContext {
+	baseContext := NewBaseContext(NewBaseContextOpts{
+		ViewName:   "files",
+		WindowName: "files",
+		Key:        FILES_CONTEXT_KEY,
+		Kind:       types.SIDE_CONTEXT,
+	})
+
+	self := &WorkingTreeContext{}
+	takeFocus := func() error { return c.PushContext(self) }
+
+	list := NewWorkingTreeViewModal(getModel, c.Log, c.UserConfig.Gui.ShowFileTree)
+	viewTrait := NewViewTrait(getView)
+	listContextTrait := &ListContextTrait{
+		base:      baseContext,
+		listTrait: list.ListTrait,
+		viewTrait: viewTrait,
+
+		GetDisplayStrings: getDisplayStrings,
+		OnFocus:           onFocus,
+		OnRenderToMain:    onRenderToMain,
+		OnFocusLost:       onFocusLost,
+		takeFocus:         takeFocus,
+
+		// TODO: handle this in a trait
+		RenderSelection: false,
+
+		c: c,
+	}
+
+	self.BaseContext = baseContext
+	self.ListContextTrait = listContextTrait
+	self.WorkingTreeViewModal = list
+
+	return self
+}
+
+type WorkingTreeViewModal struct {
+	*ListTrait
+	*filetree.FileTreeViewModel
+	getModel func() []*models.File
+}
+
+func (self *WorkingTreeViewModal) GetItemsLength() int {
+	return self.FileTreeViewModel.GetItemsLength()
+}
+
+func (self *WorkingTreeViewModal) GetSelectedFileNode() *filetree.FileNode {
+	if self.GetItemsLength() == 0 {
+		return nil
+	}
+
+	return self.FileTreeViewModel.GetItemAtIndex(self.selectedIdx)
+}
+
+func (self *WorkingTreeViewModal) GetSelectedItem() (types.ListItem, bool) {
+	item := self.GetSelectedFileNode()
+	return item, item != nil
+}
+
+func NewWorkingTreeViewModal(getModel func() []*models.File, log *logrus.Entry, showTree bool) *WorkingTreeViewModal {
+	self := &WorkingTreeViewModal{
+		getModel:          getModel,
+		FileTreeViewModel: filetree.NewFileTreeViewModel(getModel, log, showTree),
+	}
+
+	self.ListTrait = &ListTrait{
+		selectedIdx: 0,
+		HasLength:   self,
+	}
+
+	return self
+}
diff --git a/pkg/gui/controllers/files_controller.go b/pkg/gui/controllers/files_controller.go
index b8e853d48..a27197948 100644
--- a/pkg/gui/controllers/files_controller.go
+++ b/pkg/gui/controllers/files_controller.go
@@ -22,13 +22,13 @@ type FilesController struct {
 	// struct embedding, but Go does not allow hiding public fields in an embedded struct
 	// to the client
 	c          *types.ControllerCommon
-	getContext func() types.IListContext
+	getContext func() *context.WorkingTreeContext
+	getFiles   func() []*models.File
 	git        *commands.GitCommand
 	os         *oscommands.OSCommand
 
 	getSelectedFileNode    func() *filetree.FileNode
 	getContexts            func() context.ContextTree
-	getViewModel           func() *filetree.FileTreeViewModel
 	enterSubmodule         func(submodule *models.SubmoduleConfig) error
 	getSubmodules          func() []*models.SubmoduleConfig
 	setCommitMessage       func(message string)
@@ -48,12 +48,12 @@ var _ types.IController = &FilesController{}
 
 func NewFilesController(
 	c *types.ControllerCommon,
-	getContext func() types.IListContext,
+	getContext func() *context.WorkingTreeContext,
+	getFiles func() []*models.File,
 	git *commands.GitCommand,
 	os *oscommands.OSCommand,
 	getSelectedFileNode func() *filetree.FileNode,
 	allContexts func() context.ContextTree,
-	getViewModel func() *filetree.FileTreeViewModel,
 	enterSubmodule func(submodule *models.SubmoduleConfig) error,
 	getSubmodules func() []*models.SubmoduleConfig,
 	setCommitMessage func(message string),
@@ -70,11 +70,11 @@ func NewFilesController(
 	return &FilesController{
 		c:                      c,
 		getContext:             getContext,
+		getFiles:               getFiles,
 		git:                    git,
 		os:                     os,
 		getSelectedFileNode:    getSelectedFileNode,
 		getContexts:            allContexts,
-		getViewModel:           getViewModel,
 		enterSubmodule:         enterSubmodule,
 		getSubmodules:          getSubmodules,
 		setCommitMessage:       setCommitMessage,
@@ -185,6 +185,7 @@ func (self *FilesController) Keybindings(getKey func(key string) interface{}, co
 			Description: self.c.Tr.LcViewResetToUpstreamOptions,
 			OpensMenu:   true,
 		},
+		// here
 		{
 			Key:         getKey(config.Files.ToggleTreeView),
 			Handler:     self.toggleTreeView,
@@ -303,7 +304,7 @@ func (self *FilesController) EnterFile(opts types.OnFocusOpts) error {
 }
 
 func (self *FilesController) allFilesStaged() bool {
-	for _, file := range self.getViewModel().GetAllFiles() {
+	for _, file := range self.getFiles() {
 		if file.HasUnstagedChanges {
 			return false
 		}
@@ -433,7 +434,7 @@ func (self *FilesController) HandleCommitPress() error {
 		return self.c.Error(err)
 	}
 
-	if self.getViewModel().GetItemsLength() == 0 {
+	if len(self.getFiles()) == 0 {
 		return self.c.ErrorMsg(self.c.Tr.NoFilesStagedTitle)
 	}
 
@@ -484,7 +485,7 @@ func (self *FilesController) promptToStageAllAndRetry(retry func() error) error
 }
 
 func (self *FilesController) handleAmendCommitPress() error {
-	if self.getViewModel().GetItemsLength() == 0 {
+	if len(self.getFiles()) == 0 {
 		return self.c.ErrorMsg(self.c.Tr.NoFilesStagedTitle)
 	}
 
@@ -510,7 +511,7 @@ func (self *FilesController) handleAmendCommitPress() error {
 // HandleCommitEditorPress - handle when the user wants to commit changes via
 // their editor rather than via the popup panel
 func (self *FilesController) HandleCommitEditorPress() error {
-	if self.getViewModel().GetItemsLength() == 0 {
+	if len(self.getFiles()) == 0 {
 		return self.c.ErrorMsg(self.c.Tr.NoFilesStagedTitle)
 	}
 
@@ -551,7 +552,7 @@ func (self *FilesController) handleStatusFilterPressed() error {
 }
 
 func (self *FilesController) setStatusFiltering(filter filetree.FileTreeDisplayFilter) error {
-	self.getViewModel().SetFilter(filter)
+	self.getContext().FileTreeViewModel.SetFilter(filter)
 	return self.c.PostRefreshUpdate(self.getContext())
 }
 
@@ -642,7 +643,7 @@ func (self *FilesController) handleToggleDirCollapsed() error {
 		return nil
 	}
 
-	self.getViewModel().ToggleCollapsed(node.GetPath())
+	self.getContext().FileTreeViewModel.ToggleCollapsed(node.GetPath())
 
 	if err := self.c.PostRefreshUpdate(self.getContexts().Files); err != nil {
 		self.c.Log.Error(err)
@@ -655,12 +656,12 @@ func (self *FilesController) toggleTreeView() error {
 	// get path of currently selected file
 	path := self.getSelectedPath()
 
-	self.getViewModel().ToggleShowTree()
+	self.getContext().FileTreeViewModel.ToggleShowTree()
 
 	// find that same node in the new format and move the cursor to it
 	if path != "" {
-		self.getViewModel().ExpandToPath(path)
-		index, found := self.getViewModel().GetIndexForPath(path)
+		self.getContext().FileTreeViewModel.ExpandToPath(path)
+		index, found := self.getContext().FileTreeViewModel.GetIndexForPath(path)
 		if found {
 			self.getContext().GetPanelState().SetSelectedLineIdx(index)
 		}
diff --git a/pkg/gui/files_panel.go b/pkg/gui/files_panel.go
index db4bd3a54..1ec32e708 100644
--- a/pkg/gui/files_panel.go
+++ b/pkg/gui/files_panel.go
@@ -5,18 +5,12 @@ import (
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
 	"github.com/jesseduffield/lazygit/pkg/gui/filetree"
 	"github.com/jesseduffield/lazygit/pkg/gui/types"
-	"github.com/jesseduffield/lazygit/pkg/utils"
 )
 
 // list panel functions
 
 func (gui *Gui) getSelectedFileNode() *filetree.FileNode {
-	selectedLine := gui.State.Panels.Files.SelectedLineIdx
-	if selectedLine == -1 {
-		return nil
-	}
-
-	return gui.State.FileTreeViewModel.GetItemAtIndex(selectedLine)
+	return gui.State.Contexts.Files.GetSelectedFileNode()
 }
 
 func (gui *Gui) getSelectedFile() *models.File {
@@ -96,44 +90,6 @@ func (gui *Gui) promptToContinueRebase() error {
 	})
 }
 
-// Let's try to find our file again and move the cursor to that.
-// If we can't find our file, it was probably just removed by the user. In that
-// case, we go looking for where the next file has been moved to. Given that the
-// user could have removed a whole directory, we continue iterating through the old
-// nodes until we find one that exists in the new set of nodes, then move the cursor
-// to that.
-// prevNodes starts from our previously selected node because we don't need to consider anything above that
-func (gui *Gui) findNewSelectedIdx(prevNodes []*filetree.FileNode, currNodes []*filetree.FileNode) int {
-	getPaths := func(node *filetree.FileNode) []string {
-		if node == nil {
-			return nil
-		}
-		if node.File != nil && node.File.IsRename() {
-			return node.File.Names()
-		} else {
-			return []string{node.Path}
-		}
-	}
-
-	for _, prevNode := range prevNodes {
-		selectedPaths := getPaths(prevNode)
-
-		for idx, node := range currNodes {
-			paths := getPaths(node)
-
-			// If you started off with a rename selected, and now it's broken in two, we want you to jump to the new file, not the old file.
-			// This is because the new should be in the same position as the rename was meaning less cursor jumping
-			foundOldFileInRename := prevNode.File != nil && prevNode.File.IsRename() && node.Path == prevNode.File.PreviousName
-			foundNode := utils.StringArraysOverlap(paths, selectedPaths) && !foundOldFileInRename
-			if foundNode {
-				return idx
-			}
-		}
-	}
-
-	return -1
-}
-
 func (gui *Gui) onFocusFile() error {
 	gui.takeOverMergeConflictScrolling()
 	return nil
diff --git a/pkg/gui/filetree/file_node_test.go b/pkg/gui/filetree/file_node_test.go
index 8961015ac..02ae84f1d 100644
--- a/pkg/gui/filetree/file_node_test.go
+++ b/pkg/gui/filetree/file_node_test.go
@@ -142,13 +142,13 @@ func TestGetFile(t *testing.T) {
 	}{
 		{
 			name:      "valid case",
-			viewModel: NewFileTreeViewModel([]*models.File{{Name: "blah/one"}, {Name: "blah/two"}}, nil, false),
+			viewModel: NewFileTreeViewModel(func() []*models.File { return []*models.File{{Name: "blah/one"}, {Name: "blah/two"}} }, nil, false),
 			path:      "blah/two",
 			expected:  &models.File{Name: "blah/two"},
 		},
 		{
 			name:      "not found",
-			viewModel: NewFileTreeViewModel([]*models.File{{Name: "blah/one"}, {Name: "blah/two"}}, nil, false),
+			viewModel: NewFileTreeViewModel(func() []*models.File { return []*models.File{{Name: "blah/one"}, {Name: "blah/two"}} }, nil, false),
 			path:      "blah/three",
 			expected:  nil,
 		},
diff --git a/pkg/gui/filetree/file_tree_view_model.go b/pkg/gui/filetree/file_tree_view_model.go
index 01eb751e3..e3d01ef7d 100644
--- a/pkg/gui/filetree/file_tree_view_model.go
+++ b/pkg/gui/filetree/file_tree_view_model.go
@@ -19,7 +19,7 @@ const (
 )
 
 type FileTreeViewModel struct {
-	files          []*models.File
+	getFiles       func() []*models.File
 	tree           *FileNode
 	showTree       bool
 	log            *logrus.Entry
@@ -28,8 +28,9 @@ type FileTreeViewModel struct {
 	sync.RWMutex
 }
 
-func NewFileTreeViewModel(files []*models.File, log *logrus.Entry, showTree bool) *FileTreeViewModel {
+func NewFileTreeViewModel(getFiles func() []*models.File, log *logrus.Entry, showTree bool) *FileTreeViewModel {
 	viewModel := &FileTreeViewModel{
+		getFiles:       getFiles,
 		log:            log,
 		showTree:       showTree,
 		filter:         DisplayAll,
@@ -37,8 +38,6 @@ func NewFileTreeViewModel(files []*models.File, log *logrus.Entry, showTree bool
 		RWMutex:        sync.RWMutex{},
 	}
 
-	viewModel.SetFiles(files)
-
 	return viewModel
 }
 
@@ -51,11 +50,9 @@ func (self *FileTreeViewModel) ExpandToPath(path string) {
 }
 
 func (self *FileTreeViewModel) GetFilesForDisplay() []*models.File {
-	files := self.files
-
 	switch self.filter {
 	case DisplayAll:
-		return files
+		return self.getFiles()
 	case DisplayStaged:
 		return self.FilterFiles(func(file *models.File) bool { return file.HasStagedChanges })
 	case DisplayUnstaged:
@@ -69,7 +66,7 @@ func (self *FileTreeViewModel) GetFilesForDisplay() []*models.File {
 
 func (self *FileTreeViewModel) FilterFiles(test func(*models.File) bool) []*models.File {
 	result := make([]*models.File, 0)
-	for _, file := range self.files {
+	for _, file := range self.getFiles() {
 		if test(file) {
 			result = append(result, file)
 		}
@@ -93,7 +90,7 @@ func (self *FileTreeViewModel) GetItemAtIndex(index int) *FileNode {
 }
 
 func (self *FileTreeViewModel) GetFile(path string) *models.File {
-	for _, file := range self.files {
+	for _, file := range self.getFiles() {
 		if file.Name == path {
 			return file
 		}
@@ -120,13 +117,7 @@ func (self *FileTreeViewModel) GetItemsLength() int {
 }
 
 func (self *FileTreeViewModel) GetAllFiles() []*models.File {
-	return self.files
-}
-
-func (self *FileTreeViewModel) SetFiles(files []*models.File) {
-	self.files = files
-
-	self.SetTree()
+	return self.getFiles()
 }
 
 func (self *FileTreeViewModel) SetTree() {
diff --git a/pkg/gui/filetree/file_tree_view_model_test.go b/pkg/gui/filetree/file_tree_view_model_test.go
index 89b8e74df..a168573ea 100644
--- a/pkg/gui/filetree/file_tree_view_model_test.go
+++ b/pkg/gui/filetree/file_tree_view_model_test.go
@@ -73,7 +73,7 @@ func TestFilterAction(t *testing.T) {
 	for _, s := range scenarios {
 		s := s
 		t.Run(s.name, func(t *testing.T) {
-			mngr := &FileTreeViewModel{files: s.files, filter: s.filter}
+			mngr := &FileTreeViewModel{getFiles: s.files, filter: s.filter}
 			result := mngr.GetFilesForDisplay()
 			assert.EqualValues(t, s.expected, result)
 		})
diff --git a/pkg/gui/gui.go b/pkg/gui/gui.go
index 09d508e44..a71b45116 100644
--- a/pkg/gui/gui.go
+++ b/pkg/gui/gui.go
@@ -175,8 +175,8 @@ type PrevLayout struct {
 type GuiRepoState struct {
 	// the file panels (files and commit files) can render as a tree, so we have
 	// managers for them which handle rendering a flat list of files in tree form
-	FileTreeViewModel       *filetree.FileTreeViewModel
 	CommitFileTreeViewModel *filetree.CommitFileTreeViewModel
+	Files                   []*models.File
 	Submodules              []*models.SubmoduleConfig
 	Branches                []*models.Branch
 	Commits                 []*models.Commit
@@ -433,14 +433,13 @@ func (gui *Gui) resetState(filterPath string, reuseState bool) {
 	contexts := gui.contextTree()
 
 	screenMode := SCREEN_NORMAL
-	initialContext := contexts.Files
+	var initialContext types.IListContext = contexts.Files
 	if filterPath != "" {
 		screenMode = SCREEN_HALF
 		initialContext = contexts.BranchCommits
 	}
 
 	gui.State = &GuiRepoState{
-		FileTreeViewModel:       filetree.NewFileTreeViewModel(make([]*models.File, 0), gui.Log, showTree),
 		CommitFileTreeViewModel: filetree.NewCommitFileTreeViewModel(make([]*models.CommitFile, 0), gui.Log, showTree),
 		Commits:                 make([]*models.Commit, 0),
 		FilteredReflogCommits:   make([]*models.Commit, 0),
@@ -586,7 +585,7 @@ func (gui *Gui) setControllers() {
 		bisect:      controllers.NewBisectHelper(controllerCommon, gui.git),
 		suggestions: NewSuggestionsHelper(controllerCommon, getState, gui.refreshSuggestions),
 		files:       NewFilesHelper(controllerCommon, gui.git, osCommand),
-		workingTree: NewWorkingTreeHelper(func() *filetree.FileTreeViewModel { return gui.State.FileTreeViewModel }),
+		workingTree: NewWorkingTreeHelper(func() []*models.File { return gui.State.Files }),
 		tags:        controllers.NewTagsHelper(controllerCommon, gui.git),
 	}
 
@@ -609,12 +608,12 @@ func (gui *Gui) setControllers() {
 		),
 		Files: controllers.NewFilesController(
 			controllerCommon,
-			func() types.IListContext { return gui.State.Contexts.Files },
+			func() *context.WorkingTreeContext { return gui.State.Contexts.Files },
+			func() []*models.File { return gui.State.Files },
 			gui.git,
 			osCommand,
 			gui.getSelectedFileNode,
 			getContexts,
-			func() *filetree.FileTreeViewModel { return gui.State.FileTreeViewModel },
 			gui.enterSubmodule,
 			func() []*models.SubmoduleConfig { return gui.State.Submodules },
 			gui.getSetTextareaTextFn(func() *gocui.View { return gui.Views.CommitMessage }),
diff --git a/pkg/gui/list_context_config.go b/pkg/gui/list_context_config.go
index 82c5624bd..6f7203db0 100644
--- a/pkg/gui/list_context_config.go
+++ b/pkg/gui/list_context_config.go
@@ -28,21 +28,12 @@ func (gui *Gui) menuListContext() types.IListContext {
 	}
 }
 
-func (gui *Gui) filesListContext() types.IListContext {
-	return &ListContext{
-		BaseContext: context.NewBaseContext(context.NewBaseContextOpts{
-			ViewName:   "files",
-			WindowName: "files",
-			Key:        context.FILES_CONTEXT_KEY,
-			Kind:       types.SIDE_CONTEXT,
-		}),
-		GetItemsLength:  func() int { return gui.State.FileTreeViewModel.GetItemsLength() },
-		OnGetPanelState: func() types.IListPanelState { return gui.State.Panels.Files },
-		OnFocus:         OnFocusWrapper(gui.onFocusFile),
-		OnRenderToMain:  OnFocusWrapper(gui.withDiffModeCheck(gui.filesRenderToMain)),
-		Gui:             gui,
-		GetDisplayStrings: func(startIdx int, length int) [][]string {
-			lines := presentation.RenderFileTree(gui.State.FileTreeViewModel, gui.State.Modes.Diffing.Ref, gui.State.Submodules)
+func (gui *Gui) filesListContext() *context.WorkingTreeContext {
+	return context.NewWorkingTreeContext(
+		func() []*models.File { return gui.State.Files },
+		func() *gocui.View { return gui.Views.Files },
+		func(startIdx int, length int) [][]string {
+			lines := presentation.RenderFileTree(gui.State.Contexts.Files.FileTreeViewModel, gui.State.Modes.Diffing.Ref, gui.State.Submodules)
 			mappedLines := make([][]string, len(lines))
 			for i, line := range lines {
 				mappedLines[i] = []string{line}
@@ -50,11 +41,11 @@ func (gui *Gui) filesListContext() types.IListContext {
 
 			return mappedLines
 		},
-		SelectedItem: func() (types.ListItem, bool) {
-			item := gui.getSelectedFileNode()
-			return item, item != nil
-		},
-	}
+		OnFocusWrapper(gui.onFocusFile),
+		OnFocusWrapper(gui.withDiffModeCheck(gui.filesRenderToMain)),
+		nil,
+		gui.c,
+	)
 }
 
 func (gui *Gui) branchesListContext() types.IListContext {
diff --git a/pkg/gui/presentation/files_test.go b/pkg/gui/presentation/files_test.go
index c40f6247e..7a3a0ba6f 100644
--- a/pkg/gui/presentation/files_test.go
+++ b/pkg/gui/presentation/files_test.go
@@ -69,7 +69,7 @@ M  file1
 	for _, s := range scenarios {
 		s := s
 		t.Run(s.name, func(t *testing.T) {
-			viewModel := filetree.NewFileTreeViewModel(s.files, utils.NewDummyLog(), true)
+			viewModel := filetree.NewFileTreeViewModel(func() []*models.File { return s.files }, utils.NewDummyLog(), true)
 			for _, path := range s.collapsedPaths {
 				viewModel.ToggleCollapsed(path)
 			}
diff --git a/pkg/gui/refresh.go b/pkg/gui/refresh.go
index bf22b5ff6..6f8f16fa5 100644
--- a/pkg/gui/refresh.go
+++ b/pkg/gui/refresh.go
@@ -371,7 +371,8 @@ func (gui *Gui) refreshStateFiles() error {
 
 	selectedNode := gui.getSelectedFileNode()
 
-	prevNodes := gui.State.FileTreeViewModel.GetAllItems()
+	fileTreeViewModel := state.Contexts.Files.WorkingTreeViewModal
+	prevNodes := fileTreeViewModel.GetAllItems()
 	prevSelectedLineIdx := gui.State.Panels.Files.SelectedLineIdx
 
 	// If git thinks any of our files have inline merge conflicts, but they actually don't,
@@ -383,7 +384,7 @@ func (gui *Gui) refreshStateFiles() error {
 	// we call git status again.
 	pathsToStage := []string{}
 	prevConflictFileCount := 0
-	for _, file := range state.FileTreeViewModel.GetAllFiles() {
+	for _, file := range gui.State.Files {
 		if file.HasMergeConflicts {
 			prevConflictFileCount++
 		}
@@ -419,10 +420,10 @@ func (gui *Gui) refreshStateFiles() error {
 	}
 
 	// for when you stage the old file of a rename and the new file is in a collapsed dir
-	state.FileTreeViewModel.RWMutex.Lock()
+	fileTreeViewModel.RWMutex.Lock()
 	for _, file := range files {
 		if selectedNode != nil && selectedNode.Path != "" && file.PreviousName == selectedNode.Path {
-			state.FileTreeViewModel.ExpandToPath(file.Name)
+			fileTreeViewModel.ExpandToPath(file.Name)
 		}
 	}
 
@@ -432,46 +433,70 @@ func (gui *Gui) refreshStateFiles() error {
 	// extra state here to see if the user's set the filter themselves we can do that, but
 	// I'd prefer to maintain as little state as possible.
 	if conflictFileCount > 0 {
-		if state.FileTreeViewModel.GetFilter() == filetree.DisplayAll {
-			state.FileTreeViewModel.SetFilter(filetree.DisplayConflicted)
+		if fileTreeViewModel.GetFilter() == filetree.DisplayAll {
+			fileTreeViewModel.SetFilter(filetree.DisplayConflicted)
 		}
-	} else if state.FileTreeViewModel.GetFilter() == filetree.DisplayConflicted {
-		state.FileTreeViewModel.SetFilter(filetree.DisplayAll)
+	} else if fileTreeViewModel.GetFilter() == filetree.DisplayConflicted {
+		fileTreeViewModel.SetFilter(filetree.DisplayAll)
 	}
 
-	state.FileTreeViewModel.SetFiles(files)
-	state.FileTreeViewModel.RWMutex.Unlock()
+	state.Files = files
+	fileTreeViewModel.SetTree()
+	fileTreeViewModel.RWMutex.Unlock()
 
 	if err := gui.fileWatcher.addFilesToFileWatcher(files); err != nil {
 		return err
 	}
 
 	if selectedNode != nil {
-		newIdx := gui.findNewSelectedIdx(prevNodes[prevSelectedLineIdx:], state.FileTreeViewModel.GetAllItems())
+		newIdx := gui.findNewSelectedIdx(prevNodes[prevSelectedLineIdx:], fileTreeViewModel.GetAllItems())
 		if newIdx != -1 && newIdx != prevSelectedLineIdx {
-			newNode := state.FileTreeViewModel.GetItemAtIndex(newIdx)
-			// when not in tree mode, we show merge conflict files at the top, so you
-			// can work through them one by one without having to sift through a large
-			// set of files. If you have just fixed the merge conflicts of a file, we
-			// actually don't want to jump to that file's new position, because that
-			// file will now be ages away amidst the other files without merge
-			// conflicts: the user in this case would rather work on the next file
-			// with merge conflicts, which will have moved up to fill the gap left by
-			// the last file, meaning the cursor doesn't need to move at all.
-			leaveCursor := !state.FileTreeViewModel.InTreeMode() && newNode != nil &&
-				selectedNode.File != nil && selectedNode.File.HasMergeConflicts &&
-				newNode.File != nil && !newNode.File.HasMergeConflicts
-
-			if !leaveCursor {
-				state.Panels.Files.SelectedLineIdx = newIdx
-			}
+			state.Panels.Files.SelectedLineIdx = newIdx
 		}
 	}
 
-	gui.refreshSelectedLine(state.Panels.Files, state.FileTreeViewModel.GetItemsLength())
+	gui.refreshSelectedLine(state.Panels.Files, fileTreeViewModel.GetItemsLength())
 	return nil
 }
 
+// Let's try to find our file again and move the cursor to that.
+// If we can't find our file, it was probably just removed by the user. In that
+// case, we go looking for where the next file has been moved to. Given that the
+// user could have removed a whole directory, we continue iterating through the old
+// nodes until we find one that exists in the new set of nodes, then move the cursor
+// to that.
+// prevNodes starts from our previously selected node because we don't need to consider anything above that
+func (gui *Gui) findNewSelectedIdx(prevNodes []*filetree.FileNode, currNodes []*filetree.FileNode) int {
+	getPaths := func(node *filetree.FileNode) []string {
+		if node == nil {
+			return nil
+		}
+		if node.File != nil && node.File.IsRename() {
+			return node.File.Names()
+		} else {
+			return []string{node.Path}
+		}
+	}
+
+	for _, prevNode := range prevNodes {
+		selectedPaths := getPaths(prevNode)
+
+		for idx, node := range currNodes {
+			paths := getPaths(node)
+
+			// If you started off with a rename selected, and now it's broken in two, we want you to jump to the new file, not the old file.
+			// This is because the new should be in the same position as the rename was meaning less cursor jumping
+			foundOldFileInRename := prevNode.File != nil && prevNode.File.IsRename() && node.Path == prevNode.File.PreviousName
+			foundNode := utils.StringArraysOverlap(paths, selectedPaths) && !foundOldFileInRename
+			if foundNode {
+				return idx
+			}
+		}
+	}
+
+	return -1
+}
+
 // the reflogs panel is the only panel where we cache data, in that we only
 // load entries that have been created since we last ran the call. This means
 // we need to be more careful with how we use this, and to ensure we're emptying
diff --git a/pkg/gui/working_tree_helper.go b/pkg/gui/working_tree_helper.go
index ceca5e0e2..3b0162d75 100644
--- a/pkg/gui/working_tree_helper.go
+++ b/pkg/gui/working_tree_helper.go
@@ -2,21 +2,20 @@ package gui
 
 import (
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
-	"github.com/jesseduffield/lazygit/pkg/gui/filetree"
 )
 
 type WorkingTreeHelper struct {
-	getFileTreeViewModel func() *filetree.FileTreeViewModel
+	getFiles func() []*models.File
 }
 
-func NewWorkingTreeHelper(getFileTreeViewModel func() *filetree.FileTreeViewModel) *WorkingTreeHelper {
+func NewWorkingTreeHelper(getFiles func() []*models.File) *WorkingTreeHelper {
 	return &WorkingTreeHelper{
-		getFileTreeViewModel: getFileTreeViewModel,
+		getFiles: getFiles,
 	}
 }
 
 func (self *WorkingTreeHelper) AnyStagedFiles() bool {
-	files := self.getFileTreeViewModel().GetAllFiles()
+	files := self.getFiles()
 	for _, file := range files {
 		if file.HasStagedChanges {
 			return true
@@ -26,7 +25,7 @@ func (self *WorkingTreeHelper) AnyStagedFiles() bool {
 }
 
 func (self *WorkingTreeHelper) AnyTrackedFiles() bool {
-	files := self.getFileTreeViewModel().GetAllFiles()
+	files := self.getFiles()
 	for _, file := range files {
 		if file.Tracked {
 			return true
@@ -40,7 +39,7 @@ func (self *WorkingTreeHelper) IsWorkingTreeDirty() bool {
 }
 
 func (self *WorkingTreeHelper) FileForSubmodule(submodule *models.SubmoduleConfig) *models.File {
-	for _, file := range self.getFileTreeViewModel().GetAllFiles() {
+	for _, file := range self.getFiles() {
 		if file.IsSubmodule([]*models.SubmoduleConfig{submodule}) {
 			return file
 		}