gh-149430: Fix edge-cases in profiling.sampling outputs

Lib/profiling/sampling/_heatmap_assets/heatmap.js

@@ -84,7 +84,7 @@ function showNavigationMenu(button, items, title) {
 
         item.appendChild(funcDiv);
         item.appendChild(createElement('div', 'callee-menu-file', linkData.file));
-        item.addEventListener('click', () => window.location.href = linkData.link);
+        item.addEventListener('click', () => navigateToLine(linkData.link));
         menu.appendChild(item);
     });
 
@@ -105,7 +105,7 @@ function handleNavigationClick(button, e) {
 
     const navData = button.getAttribute('data-nav');
     if (navData) {
-        window.location.href = JSON.parse(navData).link;
+        navigateToLine(JSON.parse(navData).link);
         return;
     }
 
@@ -117,11 +117,29 @@ function handleNavigationClick(button, e) {
     }
 }
 
+function restartLineHighlight(target) {
+    target.style.animation = 'none';
+    // Force style recalculation so restoring the animation restarts it.
+    void target.offsetWidth;
+    target.style.animation = '';
+}
+
+function navigateToLine(link) {
+    const url = new URL(link, window.location.href);
+
+    if (url.href === window.location.href) {
+        scrollToTargetLine();
+    } else {
+        window.location.href = link;
+    }
+}
+
 function scrollToTargetLine() {
     if (!window.location.hash) return;
     const target = document.querySelector(window.location.hash);
     if (target) {
         target.scrollIntoView({ behavior: 'smooth', block: 'start' });
+        restartLineHighlight(target);
     }
 }
 

Lib/profiling/sampling/cli.py

@@ -160,7 +160,9 @@ def _build_child_profiler_args(args):
         child_args.extend(["--mode", mode])
 
     # Format options (skip pstats as it's the default)
-    if args.format != "pstats":
+    if args.format == "diff_flamegraph":
+        child_args.extend(["--diff-flamegraph", args.diff_baseline])
+    elif args.format != "pstats":
         child_args.append(f"--{args.format}")
 
     return child_args

Lib/profiling/sampling/stack_collector.py

@@ -698,6 +698,8 @@ def _add_elided_metadata(self, node, baseline_stats, scale, path):
         func_key = self._extract_func_key(node, self._baseline_collector._string_table)
         current_path = path + (func_key,) if func_key else path
 
+        baseline_self = 0
+        baseline_total = 0
         if func_key and current_path in baseline_stats:
             baseline_data = baseline_stats[current_path]
             baseline_self = baseline_data["self"] * scale

Lib/test/test_profiling/test_sampling_profiler/test_children.py

@@ -109,6 +109,39 @@ def _wait_for_process_ready(proc, timeout):
     return proc.poll() is None
 
 
+@unittest.skipIf(
+    _build_child_profiler_args is None,
+    "profiling.sampling.cli unavailable",
+)
+class TestChildProfilerArgBuilder(unittest.TestCase):
+    """Tests for child profiler CLI argument construction."""
+
+    def test_build_child_profiler_args_diff_flamegraph(self):
+        """Test child args use the real --diff-flamegraph flag."""
+        args = argparse.Namespace(
+            sample_interval_usec=1000,
+            duration=None,
+            all_threads=False,
+            realtime_stats=False,
+            native=False,
+            gc=True,
+            opcodes=False,
+            async_aware=False,
+            mode="wall",
+            format="diff_flamegraph",
+            diff_baseline="baseline.bin",
+        )
+
+        child_args = _build_child_profiler_args(args)
+
+        self.assertIn("--diff-flamegraph", child_args)
+        self.assertNotIn("--diff_flamegraph", child_args)
+
+        flag_index = child_args.index("--diff-flamegraph")
+        self.assertGreater(len(child_args), flag_index + 1)
+        self.assertEqual(child_args[flag_index + 1], "baseline.bin")
+
+
 @requires_remote_subprocess_debugging()
 class TestGetChildPids(unittest.TestCase):
     """Tests for the get_child_pids function."""

Lib/test/test_profiling/test_sampling_profiler/test_collectors.py

@@ -17,6 +17,7 @@
         FlamegraphCollector,
     )
     from profiling.sampling.gecko_collector import GeckoCollector
+    from profiling.sampling.heatmap_collector import _TemplateLoader
     from profiling.sampling.collector import extract_lineno, normalize_location
     from profiling.sampling.opcode_utils import get_opcode_info, format_opcode
     from profiling.sampling.constants import (
@@ -81,6 +82,18 @@ def test_mock_frame_info_with_empty_and_unicode_values(self):
         self.assertEqual(frame.location.lineno, 999999)
         self.assertEqual(frame.funcname, long_funcname)
 
+    def test_heatmap_navigation_restarts_line_highlight(self):
+        """Test heatmap navigation can replay target line highlights."""
+        loader = _TemplateLoader()
+
+        self.assertIn(".code-line:target", loader.file_css)
+        self.assertIn("function restartLineHighlight(target)", loader.file_js)
+        self.assertIn("target.style.animation = 'none'", loader.file_js)
+        self.assertIn("void target.offsetWidth", loader.file_js)
+        self.assertIn("url.href === window.location.href", loader.file_js)
+        self.assertIn("navigateToLine(JSON.parse(navData).link)", loader.file_js)
+        self.assertIn("navigateToLine(linkData.link)", loader.file_js)
+
     def test_pstats_collector_with_extreme_intervals_and_empty_data(self):
         """Test PstatsCollector handles zero/large intervals, empty frames, None thread IDs, and duplicate frames."""
         # Test with zero interval
@@ -1402,6 +1415,39 @@ def test_diff_flamegraph_elided_stacks(self):
         self.assertGreater(child["baseline"], 0)
         self.assertAlmostEqual(child["diff"], -child["baseline"])
 
+    def test_diff_flamegraph_elided_top_level_root(self):
+        """Elided top-level roots do not crash metadata generation."""
+        baseline_frames_1 = [
+            MockInterpreterInfo(0, [
+                MockThreadInfo(1, [
+                    MockFrameInfo("file.py", 10, "kept_leaf"),
+                    MockFrameInfo("file.py", 20, "kept_root"),
+                ])
+            ])
+        ]
+        baseline_frames_2 = [
+            MockInterpreterInfo(0, [
+                MockThreadInfo(1, [
+                    MockFrameInfo("file.py", 30, "old_leaf"),
+                    MockFrameInfo("file.py", 40, "old_root"),
+                ])
+            ])
+        ]
+
+        diff = make_diff_collector_with_mock_baseline([
+            baseline_frames_1,
+            baseline_frames_2,
+        ])
+        diff.collect(baseline_frames_1)
+
+        data = diff._convert_to_flamegraph_format()
+        elided = data["stats"]["elided_flamegraph"]
+        elided_strings = elided.get("strings", [])
+        children = elided.get("children", [])
+
+        self.assertEqual(len(children), 1)
+        self.assertIn("old_root", resolve_name(children[0], elided_strings))
+
     def test_diff_flamegraph_function_matched_despite_line_change(self):
         """Functions match by (filename, funcname), ignoring lineno."""
         baseline_frames = [