Optimize CommentMapper.visit_FunctionDef

The optimized code achieves a 10% speedup through several targeted micro-optimizations that reduce overhead in the tight loops:

**Key optimizations:**

1. **Local variable caching for frequent lookups**: Storing `self.original_runtimes` and `self.optimized_runtimes` in local variables (`original_runtimes`, `optimized_runtimes`) eliminates repeated attribute lookups in the inner loops, which is expensive due to Python's method resolution.

2. **Eliminated unnecessary list allocations**: The original code created `nodes_to_check = [compound_line_node]` and extended it with `getattr(compound_line_node, "body", [])` on every iteration. The optimized version uses `hasattr()` to check for a body attribute first, then iterates directly over `compound_line_node.body` when present, avoiding list creation and extension.

3. **Reduced redundant attribute access**: Added `nbody = node.body` and `line_body = line_node.body` to cache commonly accessed attributes, reducing repeated lookups.

4. **Streamlined f-string usage**: Consistent use of f-strings (`f"{i}_{j}"` instead of `str(i) + "_" + str(j)`) provides marginal performance improvements for string concatenation.

The optimizations are most effective for **large-scale test cases** - the annotated tests show the biggest improvements (13.4% faster) occur with functions containing many compound statements (50 for-loops with 10 statements each), where the eliminated allocations and attribute lookups compound significantly. Simple functions show minimal or slight performance regression due to the overhead of additional local variable assignments, but complex functions with nested loops see substantial gains.

Author: codeflash-ai[bot]

codeflash/code_utils/edit_generated_tests.py

@@ -61,28 +61,39 @@ def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> ast.AsyncFunctio
 
     def _process_function_def_common(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> None:
         self.context_stack.append(node.name)
-        i = len(node.body) - 1
+        nbody = node.body
+        i = len(nbody) - 1
         test_qualified_name = ".".join(self.context_stack)
         key = test_qualified_name + "#" + str(self.abs_path)
+        original_runtimes = self.original_runtimes
+        optimized_runtimes = self.optimized_runtimes
+
         while i >= 0:
-            line_node = node.body[i]
+            line_node = nbody[i]
             if isinstance(line_node, (ast.With, ast.For, ast.While, ast.If)):
-                j = len(line_node.body) - 1
+                line_body = line_node.body
+                j = len(line_body) - 1
                 while j >= 0:
-                    compound_line_node: ast.stmt = line_node.body[j]
-                    nodes_to_check = [compound_line_node]
-                    nodes_to_check.extend(getattr(compound_line_node, "body", []))
-                    for internal_node in nodes_to_check:
-                        if isinstance(internal_node, (ast.stmt, ast.Assign)):
-                            inv_id = str(i) + "_" + str(j)
-                            match_key = key + "#" + inv_id
-                            if match_key in self.original_runtimes and match_key in self.optimized_runtimes:
-                                self.results[internal_node.lineno] = self.get_comment(match_key)
+                    compound_line_node: ast.stmt = line_body[j]
+                    # Fast-path: most ast.stmt don't have a body
+                    if hasattr(compound_line_node, "body"):
+                        for internal_node in compound_line_node.body:
+                            if isinstance(internal_node, (ast.stmt, ast.Assign)):
+                                inv_id = f"{i}_{j}"
+                                match_key = f"{key}#{inv_id}"
+                                if match_key in original_runtimes and match_key in optimized_runtimes:
+                                    self.results[internal_node.lineno] = self.get_comment(match_key)
+                    # Always check the compound_line_node itself
+                    if isinstance(compound_line_node, (ast.stmt, ast.Assign)):
+                        inv_id = f"{i}_{j}"
+                        match_key = f"{key}#{inv_id}"
+                        if match_key in original_runtimes and match_key in optimized_runtimes:
+                            self.results[compound_line_node.lineno] = self.get_comment(match_key)
                     j -= 1
             else:
                 inv_id = str(i)
-                match_key = key + "#" + inv_id
-                if match_key in self.original_runtimes and match_key in self.optimized_runtimes:
+                match_key = f"{key}#{inv_id}"
+                if match_key in original_runtimes and match_key in optimized_runtimes:
                     self.results[line_node.lineno] = self.get_comment(match_key)
             i -= 1
         self.context_stack.pop()