move more view model logic into the files view model

17 files changed, 370 insertions, 301 deletions

diff --git a/pkg/gui/context/list_context_trait.go b/pkg/gui/context/list_context_trait.go
index 390059b2e..c1d45eb4e 100644
--- a/pkg/gui/context/list_context_trait.go
+++ b/pkg/gui/context/list_context_trait.go
@@ -11,7 +11,7 @@ import (
 
 type ListContextTrait struct {
 	base      types.IBaseContext
-	listTrait *ListTrait
+	list      types.IList
 	viewTrait *ViewTrait
 
 	takeFocus func() error
@@ -30,19 +30,19 @@ type ListContextTrait struct {
 }
 
 func (self *ListContextTrait) GetPanelState() types.IListPanelState {
-	return self.listTrait
+	return self.list
 }
 
 func (self *ListContextTrait) FocusLine() {
 	// we need a way of knowing whether we've rendered to the view yet.
-	self.viewTrait.FocusPoint(self.listTrait.GetSelectedLineIdx())
+	self.viewTrait.FocusPoint(self.list.GetSelectedLineIdx())
 	if self.RenderSelection {
 		min, max := self.viewTrait.ViewPortYBounds()
 		displayStrings := self.GetDisplayStrings(min, max)
 		content := utils.RenderDisplayStrings(displayStrings)
 		self.viewTrait.SetViewPortContent(content)
 	}
-	self.viewTrait.SetFooter(formatListFooter(self.listTrait.GetSelectedLineIdx(), self.listTrait.GetItemsLength()))
+	self.viewTrait.SetFooter(formatListFooter(self.list.GetSelectedLineIdx(), self.list.GetItemsLength()))
 }
 
 func formatListFooter(selectedLineIdx int, length int) string {
@@ -52,8 +52,8 @@ func formatListFooter(selectedLineIdx int, length int) string {
 // OnFocus assumes that the content of the context has already been rendered to the view. OnRender is the function which actually renders the content to the view
 func (self *ListContextTrait) HandleRender() error {
 	if self.GetDisplayStrings != nil {
-		self.listTrait.RefreshSelectedIdx()
-		content := utils.RenderDisplayStrings(self.GetDisplayStrings(0, self.listTrait.GetItemsLength()))
+		self.list.RefreshSelectedIdx()
+		content := utils.RenderDisplayStrings(self.GetDisplayStrings(0, self.list.GetItemsLength()))
 		self.viewTrait.SetContent(content)
 		self.c.Render()
 	}
@@ -112,10 +112,12 @@ func (self *ListContextTrait) scroll(scrollFunc func()) error {
 }
 
 func (self *ListContextTrait) handleLineChange(change int) error {
-	before := self.listTrait.GetSelectedLineIdx()
-	self.listTrait.MoveSelectedLine(change)
-	after := self.listTrait.GetSelectedLineIdx()
+	before := self.list.GetSelectedLineIdx()
+	self.list.MoveSelectedLine(change)
+	after := self.list.GetSelectedLineIdx()
 
+	// doing this check so that if we're holding the up key at the start of the list
+	// we're not constantly re-rendering the main view.
 	if before != after {
 		return self.HandleFocus()
 	}
@@ -132,15 +134,15 @@ func (self *ListContextTrait) HandleNextPage() error {
 }
 
 func (self *ListContextTrait) HandleGotoTop() error {
-	return self.handleLineChange(-self.listTrait.GetItemsLength())
+	return self.handleLineChange(-self.list.GetItemsLength())
 }
 
 func (self *ListContextTrait) HandleGotoBottom() error {
-	return self.handleLineChange(self.listTrait.GetItemsLength())
+	return self.handleLineChange(self.list.GetItemsLength())
 }
 
 func (self *ListContextTrait) HandleClick(onClick func() error) error {
-	prevSelectedLineIdx := self.listTrait.GetSelectedLineIdx()
+	prevSelectedLineIdx := self.list.GetSelectedLineIdx()
 	// because we're handling a click, we need to determine the new line idx based
 	// on the view itself.
 	newSelectedLineIdx := self.viewTrait.SelectedLineIdx()
@@ -150,17 +152,16 @@ func (self *ListContextTrait) HandleClick(onClick func() error) error {
 
 	// we need to focus the view
 	if !alreadyFocused {
-
 		if err := self.takeFocus(); err != nil {
 			return err
 		}
 	}
 
-	if newSelectedLineIdx > self.listTrait.GetItemsLength()-1 {
+	if newSelectedLineIdx > self.list.GetItemsLength()-1 {
 		return nil
 	}
 
-	self.listTrait.SetSelectedLineIdx(newSelectedLineIdx)
+	self.list.SetSelectedLineIdx(newSelectedLineIdx)
 
 	if prevSelectedLineIdx == newSelectedLineIdx && alreadyFocused && onClick != nil {
 		return onClick()
@@ -169,7 +170,7 @@ func (self *ListContextTrait) HandleClick(onClick func() error) error {
 }
 
 func (self *ListContextTrait) OnSearchSelect(selectedLineIdx int) error {
-	self.listTrait.SetSelectedLineIdx(selectedLineIdx)
+	self.list.SetSelectedLineIdx(selectedLineIdx)
 	return self.HandleFocus()
 }
 
diff --git a/pkg/gui/context/list_trait.go b/pkg/gui/context/list_trait.go
deleted file mode 100644
index 27b0b5345..000000000
--- a/pkg/gui/context/list_trait.go
+++ /dev/null
@@ -1,32 +0,0 @@
-package context
-
-import "github.com/jesseduffield/lazygit/pkg/gui/types"
-
-type HasLength interface {
-	GetItemsLength() int
-}
-
-type ListTrait struct {
-	selectedIdx int
-	HasLength
-}
-
-var _ types.IListPanelState = (*ListTrait)(nil)
-
-func (self *ListTrait) GetSelectedLineIdx() int {
-	return self.selectedIdx
-}
-
-func (self *ListTrait) SetSelectedLineIdx(value int) {
-	self.selectedIdx = clamp(value, 0, self.GetItemsLength()-1)
-}
-
-// moves the cursor up or down by the given amount
-func (self *ListTrait) MoveSelectedLine(value int) {
-	self.SetSelectedLineIdx(self.selectedIdx + value)
-}
-
-// to be called when the model might have shrunk so that our selection is not not out of bounds
-func (self *ListTrait) RefreshSelectedIdx() {
-	self.SetSelectedLineIdx(self.selectedIdx)
-}
diff --git a/pkg/gui/context/tags_context.go b/pkg/gui/context/tags_context.go
index b79e4d31f..c7b468972 100644
--- a/pkg/gui/context/tags_context.go
+++ b/pkg/gui/context/tags_context.go
@@ -3,6 +3,7 @@ package context
 import (
 	"github.com/jesseduffield/gocui"
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
+	"github.com/jesseduffield/lazygit/pkg/gui/context/traits"
 	"github.com/jesseduffield/lazygit/pkg/gui/types"
 )
 
@@ -39,7 +40,7 @@ func NewTagsContext(
 	viewTrait := NewViewTrait(getView)
 	listContextTrait := &ListContextTrait{
 		base:      baseContext,
-		listTrait: list.ListTrait,
+		list:      list,
 		viewTrait: viewTrait,
 
 		GetDisplayStrings: getDisplayStrings,
@@ -62,7 +63,7 @@ func NewTagsContext(
 }
 
 type TagsViewModel struct {
-	*ListTrait
+	*traits.ListCursor
 	getModel func() []*models.Tag
 }
 
@@ -88,19 +89,7 @@ func NewTagsViewModel(getModel func() []*models.Tag) *TagsViewModel {
 		getModel: getModel,
 	}
 
-	self.ListTrait = &ListTrait{
-		selectedIdx: 0,
-		HasLength:   self,
-	}
+	self.ListCursor = traits.NewListCursor(self)
 
 	return self
 }
-
-func clamp(x int, min int, max int) int {
-	if x < min {
-		return min
-	} else if x > max {
-		return max
-	}
-	return x
-}
diff --git a/pkg/gui/context/traits/list_cursor.go b/pkg/gui/context/traits/list_cursor.go
new file mode 100644
index 000000000..9423ad89c
--- /dev/null
+++ b/pkg/gui/context/traits/list_cursor.go
@@ -0,0 +1,43 @@
+package traits
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/gui/types"
+	"github.com/jesseduffield/lazygit/pkg/utils"
+)
+
+type HasLength interface {
+	GetItemsLength() int
+}
+
+type ListCursor struct {
+	selectedIdx int
+	list        HasLength
+}
+
+func NewListCursor(list HasLength) *ListCursor {
+	return &ListCursor{selectedIdx: 0, list: list}
+}
+
+var _ types.IListCursor = (*ListCursor)(nil)
+
+func (self *ListCursor) GetSelectedLineIdx() int {
+	return self.selectedIdx
+}
+
+func (self *ListCursor) SetSelectedLineIdx(value int) {
+	self.selectedIdx = utils.Clamp(value, 0, self.list.GetItemsLength()-1)
+}
+
+// moves the cursor up or down by the given amount
+func (self *ListCursor) MoveSelectedLine(delta int) {
+	self.SetSelectedLineIdx(self.selectedIdx + delta)
+}
+
+// to be called when the model might have shrunk so that our selection is not not out of bounds
+func (self *ListCursor) RefreshSelectedIdx() {
+	self.SetSelectedLineIdx(self.selectedIdx)
+}
+
+func (self *ListCursor) GetItemsLength() int {
+	return self.list.GetItemsLength()
+}
diff --git a/pkg/gui/context/working_tree_context.go b/pkg/gui/context/working_tree_context.go
index 4df16fd11..68f197259 100644
--- a/pkg/gui/context/working_tree_context.go
+++ b/pkg/gui/context/working_tree_context.go
@@ -5,11 +5,10 @@ import (
 	"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
+	*filetree.FileTreeViewModel
 	*BaseContext
 	*ListContextTrait
 }
@@ -37,11 +36,11 @@ func NewWorkingTreeContext(
 	self := &WorkingTreeContext{}
 	takeFocus := func() error { return c.PushContext(self) }
 
-	list := NewWorkingTreeViewModal(getModel, c.Log, c.UserConfig.Gui.ShowFileTree)
+	viewModel := filetree.NewFileTreeViewModel(getModel, c.Log, c.UserConfig.Gui.ShowFileTree)
 	viewTrait := NewViewTrait(getView)
 	listContextTrait := &ListContextTrait{
 		base:      baseContext,
-		listTrait: list.ListTrait,
+		list:      viewModel,
 		viewTrait: viewTrait,
 
 		GetDisplayStrings: getDisplayStrings,
@@ -58,44 +57,12 @@ func NewWorkingTreeContext(
 
 	self.BaseContext = baseContext
 	self.ListContextTrait = listContextTrait
-	self.WorkingTreeViewModal = list
+	self.FileTreeViewModel = viewModel
 
 	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()
+func (self *WorkingTreeContext) GetSelectedItem() (types.ListItem, bool) {
+	item := self.FileTreeViewModel.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 a27197948..9209a96b3 100644
--- a/pkg/gui/controllers/files_controller.go
+++ b/pkg/gui/controllers/files_controller.go
@@ -653,20 +653,8 @@ func (self *FilesController) handleToggleDirCollapsed() error {
 }
 
 func (self *FilesController) toggleTreeView() error {
-	// get path of currently selected file
-	path := self.getSelectedPath()
-
 	self.getContext().FileTreeViewModel.ToggleShowTree()
 
-	// find that same node in the new format and move the cursor to it
-	if path != "" {
-		self.getContext().FileTreeViewModel.ExpandToPath(path)
-		index, found := self.getContext().FileTreeViewModel.GetIndexForPath(path)
-		if found {
-			self.getContext().GetPanelState().SetSelectedLineIdx(index)
-		}
-	}
-
 	return self.c.PostRefreshUpdate(self.getContext())
 }
 
diff --git a/pkg/gui/filetree/file_node_test.go b/pkg/gui/filetree/file_node_test.go
index 02ae84f1d..c7649bd16 100644
--- a/pkg/gui/filetree/file_node_test.go
+++ b/pkg/gui/filetree/file_node_test.go
@@ -136,19 +136,19 @@ func TestCompress(t *testing.T) {
 func TestGetFile(t *testing.T) {
 	scenarios := []struct {
 		name      string
-		viewModel *FileTreeViewModel
+		viewModel *FileTree
 		path      string
 		expected  *models.File
 	}{
 		{
 			name:      "valid case",
-			viewModel: NewFileTreeViewModel(func() []*models.File { return []*models.File{{Name: "blah/one"}, {Name: "blah/two"}} }, nil, false),
+			viewModel: NewFileTree(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(func() []*models.File { return []*models.File{{Name: "blah/one"}, {Name: "blah/two"}} }, nil, false),
+			viewModel: NewFileTree(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.go b/pkg/gui/filetree/file_tree.go
new file mode 100644
index 000000000..504707b28
--- /dev/null
+++ b/pkg/gui/filetree/file_tree.go
@@ -0,0 +1,174 @@
+package filetree
+
+import (
+	"fmt"
+	"sync"
+
+	"github.com/jesseduffield/lazygit/pkg/commands/models"
+	"github.com/sirupsen/logrus"
+)
+
+type FileTreeDisplayFilter int
+
+const (
+	DisplayAll FileTreeDisplayFilter = iota
+	DisplayStaged
+	DisplayUnstaged
+	// this shows files with merge conflicts
+	DisplayConflicted
+)
+
+type IFileTree interface {
+	InTreeMode() bool
+	ExpandToPath(path string)
+	FilterFiles(test func(*models.File) bool) []*models.File
+	SetFilter(filter FileTreeDisplayFilter)
+	ToggleShowTree()
+
+	GetItemAtIndex(index int) *FileNode
+	GetFile(path string) *models.File
+	GetIndexForPath(path string) (int, bool)
+	GetAllItems() []*FileNode
+	GetItemsLength() int
+	GetAllFiles() []*models.File
+
+	SetTree()
+	IsCollapsed(path string) bool
+	ToggleCollapsed(path string)
+	Tree() INode
+	CollapsedPaths() CollapsedPaths
+	GetFilter() FileTreeDisplayFilter
+}
+
+type FileTree struct {
+	getFiles       func() []*models.File
+	tree           *FileNode
+	showTree       bool
+	log            *logrus.Entry
+	filter         FileTreeDisplayFilter
+	collapsedPaths CollapsedPaths
+
+	sync.RWMutex
+}
+
+func NewFileTree(getFiles func() []*models.File, log *logrus.Entry, showTree bool) *FileTree {
+	return &FileTree{
+		getFiles:       getFiles,
+		log:            log,
+		showTree:       showTree,
+		filter:         DisplayAll,
+		collapsedPaths: CollapsedPaths{},
+		RWMutex:        sync.RWMutex{},
+	}
+}
+
+func (self *FileTree) InTreeMode() bool {
+	return self.showTree
+}
+
+func (self *FileTree) ExpandToPath(path string) {
+	self.collapsedPaths.ExpandToPath(path)
+}
+
+func (self *FileTree) getFilesForDisplay() []*models.File {
+	switch self.filter {
+	case DisplayAll:
+		return self.getFiles()
+	case DisplayStaged:
+		return self.FilterFiles(func(file *models.File) bool { return file.HasStagedChanges })
+	case DisplayUnstaged:
+		return self.FilterFiles(func(file *models.File) bool { return file.HasUnstagedChanges })
+	case DisplayConflicted:
+		return self.FilterFiles(func(file *models.File) bool { return file.HasMergeConflicts })
+	default:
+		panic(fmt.Sprintf("Unexpected files display filter: %d", self.filter))
+	}
+}
+
+func (self *FileTree) FilterFiles(test func(*models.File) bool) []*models.File {
+	result := make([]*models.File, 0)
+	for _, file := range self.getFiles() {
+		if test(file) {
+			result = append(result, file)
+		}
+	}
+	return result
+}
+
+func (self *FileTree) SetFilter(filter FileTreeDisplayFilter) {
+	self.filter = filter
+	self.SetTree()
+}
+
+func (self *FileTree) ToggleShowTree() {
+	self.showTree = !self.showTree
+	self.SetTree()
+}
+
+func (self *FileTree) GetItemAtIndex(index int) *FileNode {
+	// need to traverse the three depth first until we get to the index.
+	return self.tree.GetNodeAtIndex(index+1, self.collapsedPaths) // ignoring root
+}
+
+func (self *FileTree) GetFile(path string) *models.File {
+	for _, file := range self.getFiles() {
+		if file.Name == path {
+			return file
+		}
+	}
+
+	return nil
+}
+
+func (self *FileTree) GetIndexForPath(path string) (int, bool) {
+	index, found := self.tree.GetIndexForPath(path, self.collapsedPaths)
+	return index - 1, found
+}
+
+// note: this gets all items when the filter is taken into consideration. There may
+// be hidden files that aren't included here. Files off the screen however will
+// be included
+func (self *FileTree) GetAllItems() []*FileNode {
+	if self.tree == nil {
+		return nil
+	}
+
+	return self.tree.Flatten(self.collapsedPaths)[1:] // ignoring root
+}
+
+func (self *FileTree) GetItemsLength() int {
+	return self.tree.Size(self.collapsedPaths) - 1 // ignoring root
+}
+
+func (self *FileTree) GetAllFiles() []*models.File {
+	return self.getFiles()
+}
+
+func (self *FileTree) SetTree() {
+	filesForDisplay := self.getFilesForDisplay()
+	if self.showTree {
+		self.tree = BuildTreeFromFiles(filesForDisplay)
+	} else {
+		self.tree = BuildFlatTreeFromFiles(filesForDisplay)
+	}
+}
+
+func (self *FileTree) IsCollapsed(path string) bool {
+	return self.collapsedPaths.IsCollapsed(path)
+}
+
+func (self *FileTree) ToggleCollapsed(path string) {
+	self.collapsedPaths.ToggleCollapsed(path)
+}
+
+func (self *FileTree) Tree() INode {
+	return self.tree
+}
+
+func (self *FileTree) CollapsedPaths() CollapsedPaths {
+	return self.collapsedPaths
+}
+
+func (self *FileTree) GetFilter() FileTreeDisplayFilter {
+	return self.filter
+}
diff --git a/pkg/gui/filetree/file_tree_view_model_test.go b/pkg/gui/filetree/file_tree_test.go
index a168573ea..32c110425 100644
--- a/pkg/gui/filetree/file_tree_view_model_test.go
+++ b/pkg/gui/filetree/file_tree_test.go
@@ -73,8 +73,8 @@ func TestFilterAction(t *testing.T) {
 	for _, s := range scenarios {
 		s := s
 		t.Run(s.name, func(t *testing.T) {
-			mngr := &FileTreeViewModel{getFiles: s.files, filter: s.filter}
-			result := mngr.GetFilesForDisplay()
+			mngr := &FileTree{getFiles: func() []*models.File { return s.files }, filter: s.filter}
+			result := mngr.getFilesForDisplay()
 			assert.EqualValues(t, s.expected, result)
 		})
 	}
diff --git a/pkg/gui/filetree/file_tree_view_model.go b/pkg/gui/filetree/file_tree_view_model.go
index e3d01ef7d..814d6eaac 100644
--- a/pkg/gui/filetree/file_tree_view_model.go
+++ b/pkg/gui/filetree/file_tree_view_model.go
@@ -1,150 +1,139 @@
 package filetree
 
 import (
-	"fmt"
 	"sync"
 
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
+	"github.com/jesseduffield/lazygit/pkg/gui/context/traits"
+	"github.com/jesseduffield/lazygit/pkg/gui/types"
+	"github.com/jesseduffield/lazygit/pkg/utils"
 	"github.com/sirupsen/logrus"
 )
 
-type FileTreeDisplayFilter int
-
-const (
-	DisplayAll FileTreeDisplayFilter = iota
-	DisplayStaged
-	DisplayUnstaged
-	// this shows files with merge conflicts
-	DisplayConflicted
-)
+type IFileTreeViewModel interface {
+	IFileTree
+	types.IListCursor
+}
 
+// This combines our FileTree struct with a cursor that retains information about
+// which item is selected. It also contains logic for repositioning that cursor
+// after the files are refreshed
 type FileTreeViewModel struct {
-	getFiles       func() []*models.File
-	tree           *FileNode
-	showTree       bool
-	log            *logrus.Entry
-	filter         FileTreeDisplayFilter
-	collapsedPaths CollapsedPaths
 	sync.RWMutex
+	IFileTree
+	types.IListCursor
 }
 
+var _ IFileTreeViewModel = &FileTreeViewModel{}
+
 func NewFileTreeViewModel(getFiles func() []*models.File, log *logrus.Entry, showTree bool) *FileTreeViewModel {
-	viewModel := &FileTreeViewModel{
-		getFiles:       getFiles,
-		log:            log,
-		showTree:       showTree,
-		filter:         DisplayAll,
-		collapsedPaths: CollapsedPaths{},
-		RWMutex:        sync.RWMutex{},
+	fileTree := NewFileTree(getFiles, log, showTree)
+	listCursor := traits.NewListCursor(fileTree)
+	return &FileTreeViewModel{
+		IFileTree:   fileTree,
+		IListCursor: listCursor,
 	}
-
-	return viewModel
 }
 
-func (self *FileTreeViewModel) InTreeMode() bool {
-	return self.showTree
-}
+func (self *FileTreeViewModel) GetSelectedFileNode() *FileNode {
+	if self.GetItemsLength() == 0 {
+		return nil
+	}
 
-func (self *FileTreeViewModel) ExpandToPath(path string) {
-	self.collapsedPaths.ExpandToPath(path)
+	return self.GetItemAtIndex(self.GetSelectedLineIdx())
 }
 
-func (self *FileTreeViewModel) GetFilesForDisplay() []*models.File {
-	switch self.filter {
-	case DisplayAll:
-		return self.getFiles()
-	case DisplayStaged:
-		return self.FilterFiles(func(file *models.File) bool { return file.HasStagedChanges })
-	case DisplayUnstaged:
-		return self.FilterFiles(func(file *models.File) bool { return file.HasUnstagedChanges })
-	case DisplayConflicted:
-		return self.FilterFiles(func(file *models.File) bool { return file.HasMergeConflicts })
-	default:
-		panic(fmt.Sprintf("Unexpected files display filter: %d", self.filter))
-	}
-}
+func (self *FileTreeViewModel) SetTree() {
+	newFiles := self.GetAllFiles()
+	selectedNode := self.GetSelectedFileNode()
 
-func (self *FileTreeViewModel) FilterFiles(test func(*models.File) bool) []*models.File {
-	result := make([]*models.File, 0)
-	for _, file := range self.getFiles() {
-		if test(file) {
-			result = append(result, file)
+	// for when you stage the old file of a rename and the new file is in a collapsed dir
+	for _, file := range newFiles {
+		if selectedNode != nil && selectedNode.Path != "" && file.PreviousName == selectedNode.Path {
+			self.ExpandToPath(file.Name)
 		}
 	}
-	return result
-}
-
-func (self *FileTreeViewModel) SetFilter(filter FileTreeDisplayFilter) {
-	self.filter = filter
-	self.SetTree()
-}
 
-func (self *FileTreeViewModel) ToggleShowTree() {
-	self.showTree = !self.showTree
-	self.SetTree()
-}
+	prevNodes := self.GetAllItems()
+	prevSelectedLineIdx := self.GetSelectedLineIdx()
 
-func (self *FileTreeViewModel) GetItemAtIndex(index int) *FileNode {
-	// need to traverse the three depth first until we get to the index.
-	return self.tree.GetNodeAtIndex(index+1, self.collapsedPaths) // ignoring root
-}
+	self.IFileTree.SetTree()
 
-func (self *FileTreeViewModel) GetFile(path string) *models.File {
-	for _, file := range self.getFiles() {
-		if file.Name == path {
-			return file
+	if selectedNode != nil {
+		newNodes := self.GetAllItems()
+		newIdx := self.findNewSelectedIdx(prevNodes[prevSelectedLineIdx:], newNodes)
+		if newIdx != -1 && newIdx != prevSelectedLineIdx {
+			self.SetSelectedLineIdx(newIdx)
 		}
 	}
 
-	return nil
-}
-
-func (self *FileTreeViewModel) GetIndexForPath(path string) (int, bool) {
-	index, found := self.tree.GetIndexForPath(path, self.collapsedPaths)
-	return index - 1, found
+	self.RefreshSelectedIdx()
 }
 
-func (self *FileTreeViewModel) GetAllItems() []*FileNode {
-	if self.tree == nil {
-		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 (self *FileTreeViewModel) findNewSelectedIdx(prevNodes []*FileNode, currNodes []*FileNode) int {
+	getPaths := func(node *FileNode) []string {
+		if node == nil {
+			return nil
+		}
+		if node.File != nil && node.File.IsRename() {
+			return node.File.Names()
+		} else {
+			return []string{node.Path}
+		}
 	}
 
-	return self.tree.Flatten(self.collapsedPaths)[1:] // ignoring root
-}
-
-func (self *FileTreeViewModel) GetItemsLength() int {
-	return self.tree.Size(self.collapsedPaths) - 1 // ignoring root
-}
+	for _, prevNode := range prevNodes {
+		selectedPaths := getPaths(prevNode)
 
-func (self *FileTreeViewModel) GetAllFiles() []*models.File {
-	return self.getFiles()
-}
+		for idx, node := range currNodes {
+			paths := getPaths(node)
 
-func (self *FileTreeViewModel) SetTree() {
-	filesForDisplay := self.GetFilesForDisplay()
-	if self.showTree {
-		self.tree = BuildTreeFromFiles(filesForDisplay)
-	} else {
-		self.tree = BuildFlatTreeFromFiles(filesForDisplay)
+			// 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
+			}
+		}
 	}
-}
 
-func (self *FileTreeViewModel) IsCollapsed(path string) bool {
-	return self.collapsedPaths.IsCollapsed(path)
+	return -1
 }
 
-func (self *FileTreeViewModel) ToggleCollapsed(path string) {
-	self.collapsedPaths.ToggleCollapsed(path)
+func (self *FileTreeViewModel) SetFilter(filter FileTreeDisplayFilter) {
+	self.IFileTree.SetFilter(filter)
+	self.IListCursor.SetSelectedLineIdx(0)
 }
 
-func (self *FileTreeViewModel) Tree() INode {
-	return self.tree
-}
+// If we're going from flat to tree we want to select the same file.
+// If we're going from tree to flat and we have a file selected we want to select that.
+// If instead we've selected a directory we need to select the first file in that directory.
+func (self *FileTreeViewModel) ToggleShowTree() {
+	selectedNode := self.GetSelectedFileNode()
 
-func (self *FileTreeViewModel) CollapsedPaths() CollapsedPaths {
-	return self.collapsedPaths
-}
+	self.IFileTree.ToggleShowTree()
 
-func (self *FileTreeViewModel) GetFilter() FileTreeDisplayFilter {
-	return self.filter
+	if selectedNode == nil {
+		return
+	}
+	path := selectedNode.Path
+
+	if self.InTreeMode() {
+		self.ExpandToPath(path)
+	} else if len(selectedNode.Children) > 0 {
+		path = selectedNode.GetLeaves()[0].Path
+	}
+
+	index, found := self.GetIndexForPath(path)
+	if found {
+		self.SetSelectedLineIdx(index)
+	}
 }
diff --git a/pkg/gui/gui.go b/pkg/gui/gui.go
index a71b45116..bb515be4d 100644
--- a/pkg/gui/gui.go
+++ b/pkg/gui/gui.go
@@ -334,7 +334,6 @@ type suggestionsPanelState struct {
 }
 
 type panelStates struct {
-	Files          *filePanelState
 	Branches       *branchPanelState
 	Remotes        *remotePanelState
 	RemoteBranches *remoteBranchesState
@@ -448,7 +447,6 @@ func (gui *Gui) resetState(filterPath string, reuseState bool) {
 		BisectInfo:              git_commands.NewNullBisectInfo(),
 		Panels: &panelStates{
 			// TODO: work out why some of these are -1 and some are 0. Last time I checked there was a good reason but I'm less certain now
-			Files:          &filePanelState{listPanelState{SelectedLineIdx: -1}},
 			Submodules:     &submodulePanelState{listPanelState{SelectedLineIdx: -1}},
 			Branches:       &branchPanelState{listPanelState{SelectedLineIdx: 0}},
 			Remotes:        &remotePanelState{listPanelState{SelectedLineIdx: 0}},
diff --git a/pkg/gui/mergeconflicts/state.go b/pkg/gui/mergeconflicts/state.go
index d84f05545..b40b979e9 100644
--- a/pkg/gui/mergeconflicts/state.go
+++ b/pkg/gui/mergeconflicts/state.go