Merge pull request #709 from codeflash-ai/codeflash/optimize-pr687-2025-09-03T05.27.10

⚡️ Speed up method `CommentMapper.visit_FunctionDef` by 84% in PR #687 (`granular-async-instrumentation`)

Author: KRRT7

codeflash/code_utils/edit_generated_tests.py

@@ -70,8 +70,9 @@ def _process_function_def_common(self, node: ast.FunctionDef | ast.AsyncFunction
                 j = len(line_node.body) - 1
                 while j >= 0:
                     compound_line_node: ast.stmt = line_node.body[j]
-                    internal_node: ast.AST
-                    for internal_node in ast.walk(compound_line_node):
+                    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