gh-142927: Fix heatmap caller navigation for interior lines

pablogsal · pablogsal · commit 57f1d48355a6 · 2025-12-25T19:13:26.000Z
The heatmap was only showing caller buttons on function definition
lines, not on interior lines within a function. This happened because
callers were recorded against the function definition line but looked
up by the current line number when building navigation buttons.

Added a line_to_function mapping to track which function each sampled
line belongs to. When building navigation buttons, callers are now
looked up via the function definition line so all lines in a function
show who calls that function. Callees remain line-specific since only
actual call sites should show what they call. Added tests covering
root, middle, and leaf frame behavior in call stacks.
diff --git a/Lib/profiling/sampling/heatmap_collector.py b/Lib/profiling/sampling/heatmap_collector.py
@@ -472,6 +472,10 @@ def __init__(self, *args, **kwargs):
         self.callers_graph = collections.defaultdict(set)
         self.function_definitions = {}
 
+        # Map each sampled line to its function for proper caller lookup
+        # (filename, lineno) -> funcname
+        self.line_to_function = {}
+
         # Edge counting for call path analysis
         self.edge_samples = collections.Counter()
 
@@ -596,6 +600,10 @@ def _record_line_sample(self, filename, lineno, funcname, is_leaf=False,
         if funcname and (filename, funcname) not in self.function_definitions:
             self.function_definitions[(filename, funcname)] = lineno
 
+        # Map this line to its function for caller/callee navigation
+        if funcname:
+            self.line_to_function[(filename, lineno)] = funcname
+
     def _record_bytecode_sample(self, filename, lineno, opcode,
                                 end_lineno=None, col_offset=None, end_col_offset=None,
                                 weight=1):
@@ -1150,13 +1158,37 @@ def _format_specialization_color(self, spec_pct: int) -> str:
         return f"rgba({r}, {g}, {b}, {alpha})"
 
     def _build_navigation_buttons(self, filename: str, line_num: int) -> str:
-        """Build navigation buttons for callers/callees."""
+        """Build navigation buttons for callers/callees.
+
+        - Callers: All lines in a function show who calls this function
+        - Callees: Only actual call site lines show what they call
+        """
         line_key = (filename, line_num)
-        caller_list = self._deduplicate_by_function(self.callers_graph.get(line_key, set()))
+
+        # Find which function this line belongs to
+        funcname = self.line_to_function.get(line_key)
+
+        # Get callers: look up by function definition line, not current line
+        # This ensures all lines in a function show who calls this function
+        if funcname:
+            func_def_line = self.function_definitions.get((filename, funcname), line_num)
+            func_def_key = (filename, func_def_line)
+            caller_list = self._deduplicate_by_function(self.callers_graph.get(func_def_key, set()))
+        else:
+            caller_list = self._deduplicate_by_function(self.callers_graph.get(line_key, set()))
+
+        # Get callees: only show for actual call site lines (not every line in function)
         callee_list = self._deduplicate_by_function(self.call_graph.get(line_key, set()))
 
         # Get edge counts for each caller/callee
-        callers_with_counts = self._get_edge_counts(line_key, caller_list, is_caller=True)
+        # For callers, use the function definition key for edge lookup
+        if funcname:
+            func_def_line = self.function_definitions.get((filename, funcname), line_num)
+            caller_edge_key = (filename, func_def_line)
+        else:
+            caller_edge_key = line_key
+        callers_with_counts = self._get_edge_counts(caller_edge_key, caller_list, is_caller=True)
+        # For callees, use the actual line key since that's where the call happens
         callees_with_counts = self._get_edge_counts(line_key, callee_list, is_caller=False)
 
         # Build navigation buttons with counts
diff --git a/Lib/test/test_profiling/test_heatmap.py b/Lib/test/test_profiling/test_heatmap.py
@@ -367,6 +367,246 @@ def test_process_frames_with_file_samples_dict(self):
         self.assertEqual(collector.file_samples['test.py'][10], 1)
 
 
+class TestHeatmapCollectorNavigationButtons(unittest.TestCase):
+    """Test navigation button behavior for caller/callee relationships.
+
+    For every call stack:
+    - Root frames (entry points): only DOWN arrow (▼) for callees
+    - Middle frames: both UP (▲) and DOWN (▼) arrows
+    - Leaf frames: only UP arrow (▲) for callers
+
+    All lines within a function should show who calls that function.
+    """
+
+    def setUp(self):
+        """Set up test directory and collector."""
+        self.test_dir = tempfile.mkdtemp()
+        self.addCleanup(shutil.rmtree, self.test_dir)
+
+    def _create_collector_with_stack(self, frames_list):
+        """Helper to create a collector and process frame stacks."""
+        collector = HeatmapCollector(sample_interval_usec=100)
+        for frames in frames_list:
+            collector.process_frames(frames, thread_id=1)
+        return collector
+
+    def _export_and_get_html(self, collector, filename):
+        """Export and return HTML content for a specific file."""
+        output_path = os.path.join(self.test_dir, 'nav_test')
+        with captured_stdout(), captured_stderr():
+            collector.export(output_path)
+
+        # Find the HTML file for the given source file
+        for html_file in os.listdir(output_path):
+            if html_file.startswith('file_') and html_file.endswith('.html'):
+                html_path = os.path.join(output_path, html_file)
+                with open(html_path, 'r', encoding='utf-8') as f:
+                    content = f.read()
+                if filename in content:
+                    return content
+        return None
+
+    def test_leaf_frame_has_only_caller_button(self):
+        """Test that leaf frames (top of stack) only have caller buttons."""
+        # Create a simple call stack: caller.py:10 -> leaf.py:5
+        # leaf.py:5 is the leaf (where execution happens)
+        frames = [
+            ('leaf.py', (5, 5, -1, -1), 'leaf_func', None),      # Leaf
+            ('caller.py', (10, 10, -1, -1), 'caller_func', None)  # Caller
+        ]
+        collector = self._create_collector_with_stack([frames])
+
+        # Check that leaf has callers recorded but no callees
+        leaf_key = ('leaf.py', 5)
+        caller_key = ('caller.py', 10)
+
+        # Leaf should be in callers_graph (someone calls it)
+        self.assertIn(leaf_key, collector.callers_graph)
+
+        # Leaf should NOT be in call_graph (it doesn't call anyone)
+        self.assertNotIn(leaf_key, collector.call_graph)
+
+        # Caller should be in call_graph (it calls the leaf)
+        self.assertIn(caller_key, collector.call_graph)
+
+    def test_root_frame_has_only_callee_button(self):
+        """Test that root frames (bottom of stack) only have callee buttons."""
+        # Create a simple call stack: root.py:20 -> child.py:10
+        # root.py:20 is the root (entry point, no one calls it)
+        frames = [
+            ('child.py', (10, 10, -1, -1), 'child_func', None),  # Child (leaf)
+            ('root.py', (20, 20, -1, -1), 'root_func', None)     # Root
+        ]
+        collector = self._create_collector_with_stack([frames])
+
+        root_key = ('root.py', 20)
+        child_key = ('child.py', 10)
+
+        # Root should be in call_graph (it calls the child)
+        self.assertIn(root_key, collector.call_graph)
+
+        # Root should NOT be in callers_graph (no one calls it in this profile)
+        self.assertNotIn(root_key, collector.callers_graph)
+
+        # Child's definition line should be in callers_graph
+        child_def_key = ('child.py', 10)
+        self.assertIn(child_def_key, collector.callers_graph)
+
+    def test_middle_frame_has_both_buttons(self):
+        """Test that middle frames have both caller and callee buttons."""
+        # Create a 3-level stack: root.py:30 -> middle.py:20 -> leaf.py:10
+        frames = [
+            ('leaf.py', (10, 10, -1, -1), 'leaf_func', None),
+            ('middle.py', (20, 20, -1, -1), 'middle_func', None),
+            ('root.py', (30, 30, -1, -1), 'root_func', None)
+        ]
+        collector = self._create_collector_with_stack([frames])
+
+        middle_key = ('middle.py', 20)
+
+        # Middle should be in call_graph (it calls leaf)
+        self.assertIn(middle_key, collector.call_graph)
+
+        # Middle's definition should be in callers_graph (root calls it)
+        self.assertIn(middle_key, collector.callers_graph)
+
+    def test_all_lines_in_function_see_callers(self):
+        """Test that all lines in a function show who calls the function."""
+        # Simulate a function with multiple sampled lines
+        # Stack 1: caller.py:100 calls func at line 10, executing at line 12
+        frames1 = [
+            ('module.py', (12, 12, -1, -1), 'my_func', None),  # Interior line
+            ('caller.py', (100, 100, -1, -1), 'caller', None)
+        ]
+        # Stack 2: same caller, function executing at line 15
+        frames2 = [
+            ('module.py', (15, 15, -1, -1), 'my_func', None),  # Different interior line
+            ('caller.py', (100, 100, -1, -1), 'caller', None)
+        ]
+        # Stack 3: function definition line
+        frames3 = [
+            ('module.py', (10, 10, -1, -1), 'my_func', None),  # Definition line
+            ('caller.py', (100, 100, -1, -1), 'caller', None)
+        ]
+
+        collector = self._create_collector_with_stack([frames1, frames2, frames3])
+
+        # All lines should map to the same function
+        self.assertEqual(collector.line_to_function[('module.py', 10)], 'my_func')
+        self.assertEqual(collector.line_to_function[('module.py', 12)], 'my_func')
+        self.assertEqual(collector.line_to_function[('module.py', 15)], 'my_func')
+
+        # Function definition should have callers recorded
+        # The first seen line becomes the definition (line 12 in this case)
+        func_def = collector.function_definitions[('module.py', 'my_func')]
+        func_def_key = ('module.py', func_def)
+        self.assertIn(func_def_key, collector.callers_graph)
+
+    def test_multiple_callers_recorded(self):
+        """Test that multiple callers are recorded for a function."""
+        # Two different callers call the same function
+        frames1 = [
+            ('target.py', (10, 10, -1, -1), 'target_func', None),
+            ('caller1.py', (20, 20, -1, -1), 'caller1', None)
+        ]
+        frames2 = [
+            ('target.py', (10, 10, -1, -1), 'target_func', None),
+            ('caller2.py', (30, 30, -1, -1), 'caller2', None)
+        ]
+
+        collector = self._create_collector_with_stack([frames1, frames2])
+
+        target_key = ('target.py', 10)
+        callers = collector.callers_graph[target_key]
+
+        # Should have 2 different callers
+        self.assertEqual(len(callers), 2)
+
+        # Extract caller files
+        caller_files = {c[0] for c in callers}
+        self.assertIn('caller1.py', caller_files)
+        self.assertIn('caller2.py', caller_files)
+
+    def test_multiple_callees_from_same_line(self):
+        """Test that multiple callees from the same call site are recorded."""
+        # Same line calls different functions in different samples
+        frames1 = [
+            ('target1.py', (10, 10, -1, -1), 'func1', None),
+            ('caller.py', (20, 20, -1, -1), 'caller', None)
+        ]
+        frames2 = [
+            ('target2.py', (15, 15, -1, -1), 'func2', None),
+            ('caller.py', (20, 20, -1, -1), 'caller', None)
+        ]
+
+        collector = self._create_collector_with_stack([frames1, frames2])
+
+        caller_key = ('caller.py', 20)
+        callees = collector.call_graph[caller_key]
+
+        # Should have 2 different callees
+        self.assertEqual(len(callees), 2)
+
+        # Extract callee files
+        callee_files = {c[0] for c in callees}
+        self.assertIn('target1.py', callee_files)
+        self.assertIn('target2.py', callee_files)
+
+    def test_edge_samples_counted(self):
+        """Test that edge samples are counted correctly."""
+        frames = [
+            ('callee.py', (10, 10, -1, -1), 'callee', None),
+            ('caller.py', (20, 20, -1, -1), 'caller', None)
+        ]
+
+        collector = self._create_collector_with_stack([frames, frames, frames])
+
+        # Edge should have 3 samples
+        caller_key = ('caller.py', 20)
+        callee_key = ('callee.py', 10)
+        edge_key = (caller_key, callee_key)
+
+        self.assertEqual(collector.edge_samples[edge_key], 3)
+
+    def test_deep_call_stack_relationships(self):
+        """Test navigation relationships in a deep call stack."""
+        # 5-level stack: root -> A -> B -> C -> leaf
+        frames = [
+            ('leaf.py', (5, 5, -1, -1), 'leaf', None),
+            ('c.py', (10, 10, -1, -1), 'func_c', None),
+            ('b.py', (15, 15, -1, -1), 'func_b', None),
+            ('a.py', (20, 20, -1, -1), 'func_a', None),
+            ('root.py', (25, 25, -1, -1), 'root', None),
+        ]
+
+        collector = self._create_collector_with_stack([frames])
+
+        # Root: only callees (calls A), no callers
+        root_key = ('root.py', 25)
+        self.assertIn(root_key, collector.call_graph)
+        self.assertNotIn(root_key, collector.callers_graph)
+
+        # A: both (called by root, calls B)
+        a_key = ('a.py', 20)
+        self.assertIn(a_key, collector.call_graph)
+        self.assertIn(a_key, collector.callers_graph)
+
+        # B: both (called by A, calls C)
+        b_key = ('b.py', 15)
+        self.assertIn(b_key, collector.call_graph)
+        self.assertIn(b_key, collector.callers_graph)
+
+        # C: both (called by B, calls leaf)
+        c_key = ('c.py', 10)
+        self.assertIn(c_key, collector.call_graph)
+        self.assertIn(c_key, collector.callers_graph)
+
+        # Leaf: only callers (called by C), no callees
+        leaf_key = ('leaf.py', 5)
+        self.assertNotIn(leaf_key, collector.call_graph)
+        self.assertIn(leaf_key, collector.callers_graph)
+
+
 class TestHeatmapCollectorExport(unittest.TestCase):
     """Test HeatmapCollector.export() method."""
 
