Fix StackViewer: "Set Time Range" resets "Goto Items by callees"

src/PerfView.Tests/StackViewer/StackWindowTests.cs

@@ -277,6 +277,78 @@ public Task TestSetTimeRangeWithSpaceAsync()
             return TestSetTimeRangeWithSpaceImplAsync(CultureInfo.CurrentCulture);
         }
 
+        [WpfFact]
+        [WorkItem(2179, "https://github.com/Microsoft/perfview/issues/2179")]
+        [UseCulture("en-US")]
+        public Task TestSetTimeRangeAfterGotoCalleesAsync()
+        {
+            Func<Task<StackWindow>> setupAsync = async () =>
+            {
+                await JoinableTaskFactory.SwitchToMainThreadAsync();
+
+                var file = new TimeRangeFile();
+                await OpenAsync(JoinableTaskFactory, file, GuiApp.MainWindow, GuiApp.MainWindow.StatusBar).ConfigureAwait(true);
+                var stackSource = file.GetStackSource();
+                return stackSource.Viewer;
+            };
+
+            Func<StackWindow, Task> cleanupAsync = async stackWindow =>
+            {
+                await JoinableTaskFactory.SwitchToMainThreadAsync();
+
+                stackWindow.Close();
+            };
+
+            Func<StackWindow, Task> testDriverAsync = async stackWindow =>
+            {
+                await JoinableTaskFactory.SwitchToMainThreadAsync();
+
+                // First, navigate to the call tree tab
+                stackWindow.CallTreeTab.IsSelected = true;
+                await WaitForUIAsync(stackWindow.Dispatcher, CancellationToken.None);
+
+                // Find a node in the call tree that is not the root node.
+                var callTreeNode = stackWindow.m_callTreeView.Root.Callees[0];
+                Assert.NotNull(callTreeNode);
+
+                // Set focus to this node
+                stackWindow.SetFocus(callTreeNode.Name);
+
+                // Use "Goto Items in Callees" command (Shift+F10)
+                stackWindow.CalleesTab.IsSelected = true;
+                await WaitForUIAsync(stackWindow.Dispatcher, CancellationToken.None);
+
+                // Remember current focus after going to callees
+                string calleeFocusName = stackWindow.FocusName;
+                Assert.NotNull(calleeFocusName);
+                Assert.NotEqual("ROOT", calleeFocusName);
+
+                // Now execute Set Time Range command
+                var byNameView = stackWindow.m_byNameView;
+
+                // Use min and max values from the stacksource
+                double minTime = stackWindow.StackSource.GetSampleByIndex(0).TimeRelativeMSec;
+                double maxTime = 0;
+                for (int i = 0; i < stackWindow.StackSource.SampleIndexLimit; i++)
+                {
+                    maxTime = Math.Max(maxTime, stackWindow.StackSource.GetSampleByIndex((StackSourceSampleIndex)i).TimeRelativeMSec);
+                }
+
+                stackWindow.StartTextBox.Text = minTime.ToString("n3");
+                stackWindow.EndTextBox.Text = maxTime.ToString("n3");
+                stackWindow.Update();
+
+                // Wait for any background processing to complete
+                await stackWindow.StatusBar.WaitForWorkCompleteAsync().ConfigureAwait(true);
+
+                // Verify focus was maintained
+                Assert.Equal(calleeFocusName, stackWindow.FocusName);
+                //}
+            };
+
+            return RunUITestAsync(setupAsync, testDriverAsync, cleanupAsync);
+        }
+
         private Task TestSetTimeRangeWithSpaceImplAsync(CultureInfo culture)
         {
             Func<Task<StackWindow>> setupAsync = async () =>
@@ -304,6 +376,14 @@ private Task TestSetTimeRangeWithSpaceImplAsync(CultureInfo culture)
                 var row = byNameView.FindIndex(node => node.FirstTimeRelativeMSec > 0 && node.FirstTimeRelativeMSec < node.LastTimeRelativeMSec);
                 CallTreeNodeBase selected = byNameView[row];
 
+                // Set focus to a specific node
+                string testFocusName = selected.Name;
+                stackWindow.SetFocus(testFocusName);
+
+                // Verify the focus is set correctly
+                Assert.Equal(testFocusName, stackWindow.FocusName);
+
+                // Now execute Set Time Range command
                 var selectedCells = stackWindow.ByNameDataGrid.Grid.SelectedCells;
                 selectedCells.Clear();
                 selectedCells.Add(new DataGridCellInfo(byNameView[row], stackWindow.ByNameDataGrid.FirstTimeColumn));
@@ -314,8 +394,12 @@ private Task TestSetTimeRangeWithSpaceImplAsync(CultureInfo culture)
                 // Wait for any background processing to complete
                 await stackWindow.StatusBar.WaitForWorkCompleteAsync().ConfigureAwait(true);
 
+                // Verify time range was set correctly
                 Assert.Equal(selected.FirstTimeRelativeMSec.ToString("n3", culture), stackWindow.StartTextBox.Text);
                 Assert.Equal(selected.LastTimeRelativeMSec.ToString("n3", culture), stackWindow.EndTextBox.Text);
+
+                // Verify focus was maintained
+                Assert.Equal(testFocusName, stackWindow.FocusName);
             };
 
             return RunUITestAsync(setupAsync, testDriverAsync, cleanupAsync);

src/PerfView/StackViewer/StackWindow.xaml.cs

@@ -371,6 +371,7 @@ public void SetStackSource(StackSource newSource, Action onComplete = null)
 
                 StatusBar.EndWork(delegate ()
                 {
+                    var selectedNodeName = FocusName ?? "ROOT";
                     var oldCallTree = m_callTree;
                     m_callTree = newCallTree;
 
@@ -395,8 +396,7 @@ public void SetStackSource(StackSource newSource, Action onComplete = null)
                         ByNameDataGrid.Grid.Columns[i].SortDirection = direction;
                     }
 
-                    // SignalPropertyChange the Caller-Callee Tab
-                    SetFocus(m_callTree.Root);
+                    SetFocus(selectedNodeName);
 
                     ByNameDataGrid.Focus();