Fix topbar UI being lost on undo w/ mode change

3 files changed, 48 insertions, 0 deletions

diff --git a/source/blender/editors/undo/ed_undo.c b/source/blender/editors/undo/ed_undo.c
index 1d5c57a39b4..8b4292dedbd 100644
--- a/source/blender/editors/undo/ed_undo.c
+++ b/source/blender/editors/undo/ed_undo.c
@@ -214,6 +214,8 @@ static int ed_undo_step(bContext *C, int step, const char *undoname, ReportList
 	WM_event_add_notifier(C, NC_WINDOW, NULL);
 	WM_event_add_notifier(C, NC_WM | ND_UNDO, NULL);
 
+	WM_toolsystem_refresh_active(C);
+
 	Main *bmain = CTX_data_main(C);
 	WM_toolsystem_refresh_screen_all(bmain);
 
diff --git a/source/blender/windowmanager/WM_toolsystem.h b/source/blender/windowmanager/WM_toolsystem.h
index fd61e5c9699..9d9278668a9 100644
--- a/source/blender/windowmanager/WM_toolsystem.h
+++ b/source/blender/windowmanager/WM_toolsystem.h
@@ -106,6 +106,8 @@ void WM_toolsystem_ref_properties_ensure_ex(
 void WM_toolsystem_ref_properties_init_for_keymap(
         struct bToolRef *tref, struct PointerRNA *dst_ptr, struct PointerRNA *src_ptr, struct wmOperatorType *ot);
 
+void WM_toolsystem_refresh_active(struct bContext *C);
+
 void WM_toolsystem_refresh_screen_area(struct WorkSpace *workspace, struct ViewLayer *view_layer, struct ScrArea *sa);
 void WM_toolsystem_refresh_screen_all(struct Main *bmain);
 
diff --git a/source/blender/windowmanager/intern/wm_toolsystem.c b/source/blender/windowmanager/intern/wm_toolsystem.c
index 04548ef1315..565fade50bb 100644
--- a/source/blender/windowmanager/intern/wm_toolsystem.c
+++ b/source/blender/windowmanager/intern/wm_toolsystem.c
@@ -484,6 +484,50 @@ bool WM_toolsystem_key_from_context(
 	return false;
 }
 
+/**
+ * Use to update the active tool (shown in the top bar) in the least disruptive way.
+ *
+ * This is a little involved since there may be multiple valid active tools depending on the mode and space type.
+ *
+ * Used when undoing since the active mode may have changed.
+ */
+void WM_toolsystem_refresh_active(bContext *C)
+{
+	Main *bmain = CTX_data_main(C);
+	for (wmWindowManager *wm = bmain->wm.first; wm; wm = wm->id.next) {
+		for (wmWindow *win = wm->windows.first; win; win = win->next) {
+			WorkSpace *workspace = WM_window_get_active_workspace(win);
+			bScreen *screen = WM_window_get_active_screen(win);
+			ViewLayer *view_layer = WM_window_get_active_view_layer(win);
+			int mode_other = 0;
+			enum { UNSET = -1, CHANGE = 0, MATCH = 1 } mode_match = UNSET;
+			/* Could skip loop for modes that don't depend on space type. */
+			for (ScrArea *sa = screen->areabase.first; sa; sa = sa->next) {
+				/* Don't change the space type of the active tool, only update it's mode. */
+				if (sa->spacetype == workspace->tools_space_type) {
+					const int mode = WM_toolsystem_mode_from_spacetype(view_layer, sa, sa->spacetype);
+					if (workspace->tools_mode == mode) {
+						mode_match = MATCH;
+						break;
+					}
+					else if (mode_match == -1) {
+						mode_match = CHANGE;
+						mode_other = mode;
+					}
+				}
+			}
+
+			if (mode_match == CHANGE) {
+				const bToolKey tkey = {
+					.space_type = workspace->tools_space_type,
+					.mode = mode_other,
+				};
+				toolsystem_reinit_ensure_toolref(C, workspace, &tkey, NULL);
+			}
+		}
+	}
+}
+
 void WM_toolsystem_refresh_screen_area(WorkSpace *workspace, ViewLayer *view_layer, ScrArea *sa)
 {
 	sa->runtime.tool = NULL;