From 9d63078becb734e9131ff330915b53b462ff96d9 Mon Sep 17 00:00:00 2001
From: leo <longshuang@msn.cn>
Date: Mon, 16 Mar 2026 18:39:03 +0800
Subject: [PATCH] enhance: manually update local branch tree after checking out
 a exiting local branch (#2169)

Since the remote branch tree did not changed after checkout, and the changes in the local branch tree are completely predictable, update it manually.

Signed-off-by: leo <longshuang@msn.cn>
---
 src/ViewModels/Checkout.cs   | 31 ++++++++++++++---------------
 src/ViewModels/Repository.cs | 38 +++++++++++++++++++++++++++++++++++-
 src/Views/Checkout.axaml     |  2 +-
 3 files changed, 53 insertions(+), 18 deletions(-)

diff --git a/src/ViewModels/Checkout.cs b/src/ViewModels/Checkout.cs
index ee760331..8da89a13 100644
--- a/src/ViewModels/Checkout.cs
+++ b/src/ViewModels/Checkout.cs
@@ -4,9 +4,9 @@ namespace SourceGit.ViewModels
 {
     public class Checkout : Popup
     {
-        public string Branch
+        public string BranchName
         {
-            get;
+            get => _branch.Name;
         }
 
         public bool DiscardLocalChanges
@@ -15,10 +15,10 @@ namespace SourceGit.ViewModels
             set;
         }
 
-        public Checkout(Repository repo, string branch)
+        public Checkout(Repository repo, Models.Branch branch)
         {
             _repo = repo;
-            Branch = branch;
+            _branch = branch;
             DiscardLocalChanges = false;
         }
 
@@ -30,9 +30,10 @@ namespace SourceGit.ViewModels
         public override async Task<bool> Sure()
         {
             using var lockWatcher = _repo.LockWatcher();
-            ProgressDescription = $"Checkout '{Branch}' ...";
+            var branchName = BranchName;
+            ProgressDescription = $"Checkout '{branchName}' ...";
 
-            var log = _repo.CreateLog($"Checkout '{Branch}'");
+            var log = _repo.CreateLog($"Checkout '{branchName}'");
             Use(log);
 
             if (_repo.CurrentBranch is { IsDetachedHead: true })
@@ -70,7 +71,7 @@ namespace SourceGit.ViewModels
 
             succ = await new Commands.Checkout(_repo.FullPath)
                 .Use(log)
-                .BranchAsync(Branch, DiscardLocalChanges);
+                .BranchAsync(branchName, DiscardLocalChanges);
 
             if (succ)
             {
@@ -80,21 +81,19 @@ namespace SourceGit.ViewModels
                     await new Commands.Stash(_repo.FullPath)
                         .Use(log)
                         .PopAsync("stash@{0}");
+
+                _repo.FastRefreshBranchesAfterCheckout(_branch);
+            }
+            else
+            {
+                _repo.MarkWorkingCopyDirtyManually();
             }
 
             log.Complete();
-
-            var b = _repo.Branches.Find(x => x.IsLocal && x.Name == Branch);
-            if (b != null && _repo.HistoryFilterMode == Models.FilterMode.Included)
-                _repo.SetBranchFilterMode(b, Models.FilterMode.Included, false, false);
-
-            _repo.MarkBranchesDirtyManually();
-
-            ProgressDescription = "Waiting for branch updated...";
-            await Task.Delay(400);
             return succ;
         }
 
         private readonly Repository _repo = null;
+        private readonly Models.Branch _branch = null;
     }
 }
diff --git a/src/ViewModels/Repository.cs b/src/ViewModels/Repository.cs
index 7123a2e0..6fa078a9 100644
--- a/src/ViewModels/Repository.cs
+++ b/src/ViewModels/Repository.cs
@@ -788,6 +788,42 @@ namespace SourceGit.ViewModels
             return _watcher?.Lock();
         }
 
+        public void FastRefreshBranchesAfterCheckout(Models.Branch checkouted)
+        {
+            _watcher?.MarkBranchUpdated();
+            _watcher?.MarkWorkingCopyUpdated();
+
+            if (_currentBranch.IsDetachedHead)
+            {
+                _branches.Remove(_currentBranch);
+            }
+            else
+            {
+                _currentBranch.IsCurrent = false;
+                _currentBranch.WorktreePath = null;
+            }
+
+            checkouted.IsCurrent = true;
+            checkouted.WorktreePath = FullPath;
+            if (_historyFilterMode == Models.FilterMode.Included)
+                SetBranchFilterMode(checkouted, Models.FilterMode.Included, false, false);
+
+            List<Models.Branch> locals = [];
+            foreach (var b in _branches)
+            {
+                if (b.IsLocal)
+                    locals.Add(b);
+            }
+
+            var builder = BuildBranchTree(locals, []);
+            LocalBranchTrees = builder.Locals;
+            CurrentBranch = checkouted;
+
+            RefreshCommits();
+            RefreshWorkingCopyChanges();
+            RefreshWorktrees();
+        }
+
         public void MarkBranchesDirtyManually()
         {
             _watcher?.MarkBranchUpdated();
@@ -1304,7 +1340,7 @@ namespace SourceGit.ViewModels
 
             if (branch.IsLocal)
             {
-                await ShowAndStartPopupAsync(new Checkout(this, branch.Name));
+                await ShowAndStartPopupAsync(new Checkout(this, branch));
             }
             else
             {
diff --git a/src/Views/Checkout.axaml b/src/Views/Checkout.axaml
index 0d2f2eca..26fafe81 100644
--- a/src/Views/Checkout.axaml
+++ b/src/Views/Checkout.axaml
@@ -30,7 +30,7 @@
                  Text="{DynamicResource Text.Checkout.Target}"/>
       <StackPanel Grid.Row="0" Grid.Column="1" Orientation="Horizontal">
         <Path Width="14" Height="14" Margin="4,0" Data="{StaticResource Icons.Branch}"/>
-        <TextBlock Text="{Binding Branch}"/>
+        <TextBlock Text="{Binding BranchName}"/>
       </StackPanel>
 
       <TextBlock Grid.Row="1" Grid.Column="0"