Merge branch 'main' into uv-migration-fr

misrasaurabh1 · web-flow · commit c33d1e96be56 · 2025-06-20T17:55:00.000-07:00
diff --git a/code_to_optimize/code_directories/simple_tracer_e2e/workload.py b/code_to_optimize/code_directories/simple_tracer_e2e/workload.py
@@ -21,6 +21,45 @@ def test_threadpool() -> None:
     for r in result:
         print(r)
 
+class AlexNet:
+    def __init__(self, num_classes=1000):
+        self.num_classes = num_classes
+        self.features_size = 256 * 6 * 6
+
+    def forward(self, x):
+        features = self._extract_features(x)
+        
+        output = self._classify(features)
+        return output
+
+    def _extract_features(self, x):
+        result = []
+        for i in range(len(x)):
+            pass
+
+        return result
+
+    def _classify(self, features):
+        total = sum(features)
+        return [total % self.num_classes for _ in features]
+
+class SimpleModel:
+    @staticmethod
+    def predict(data):
+        return [x * 2 for x in data]
+    
+    @classmethod
+    def create_default(cls):
+        return cls()
+
+def test_models():
+    model = AlexNet(num_classes=10)
+    input_data = [1, 2, 3, 4, 5]
+    result = model.forward(input_data)
+
+    model2 = SimpleModel.create_default()
+    prediction = model2.predict(input_data)
 
 if __name__ == "__main__":
     test_threadpool()
+    test_models()
diff --git a/codeflash/code_utils/edit_generated_tests.py b/codeflash/code_utils/edit_generated_tests.py
@@ -1,14 +1,22 @@
+from __future__ import annotations
+
+import ast
 import os
 import re
 from pathlib import Path
+from textwrap import dedent
+from typing import TYPE_CHECKING, Union
 
 import libcst as cst
 
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.time_utils import format_perf, format_time
-from codeflash.models.models import GeneratedTests, GeneratedTestsList, InvocationId
+from codeflash.models.models import GeneratedTests, GeneratedTestsList
 from codeflash.result.critic import performance_gain
-from codeflash.verification.verification_utils import TestConfig
+
+if TYPE_CHECKING:
+    from codeflash.models.models import InvocationId
+    from codeflash.verification.verification_utils import TestConfig
 
 
 def remove_functions_from_generated_tests(
@@ -36,6 +44,94 @@ def remove_functions_from_generated_tests(
     return GeneratedTestsList(generated_tests=new_generated_tests)
 
 
+class CfoVisitor(ast.NodeVisitor):
+    """AST visitor that finds all assignments to a variable named 'codeflash_output'.
+
+    and reports their location relative to the function they're in.
+    """
+
+    def __init__(self, source_code: str) -> None:
+        self.source_lines = source_code.splitlines()
+        self.results: list[int] = []  # map actual line number to line number in ast
+
+    def _is_codeflash_output_target(self, target: Union[ast.expr, list]) -> bool:  # type: ignore[type-arg]
+        """Check if the assignment target is the variable 'codeflash_output'."""
+        if isinstance(target, ast.Name):
+            return target.id == "codeflash_output"
+        if isinstance(target, (ast.Tuple, ast.List)):
+            # Handle tuple/list unpacking: a, codeflash_output, b = values
+            return any(self._is_codeflash_output_target(elt) for elt in target.elts)
+        if isinstance(target, (ast.Subscript, ast.Attribute)):
+            # Not a simple variable assignment
+            return False
+        return False
+
+    def _record_assignment(self, node: ast.AST) -> None:
+        """Record an assignment to codeflash_output."""
+        relative_line = node.lineno - 1  # type: ignore[attr-defined]
+        self.results.append(relative_line)
+
+    def visit_Assign(self, node: ast.Assign) -> None:
+        """Visit assignment statements: codeflash_output = value."""
+        for target in node.targets:
+            if self._is_codeflash_output_target(target):
+                self._record_assignment(node)
+                break
+        self.generic_visit(node)
+
+    def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
+        """Visit annotated assignments: codeflash_output: int = value."""
+        if self._is_codeflash_output_target(node.target):
+            self._record_assignment(node)
+        self.generic_visit(node)
+
+    def visit_AugAssign(self, node: ast.AugAssign) -> None:
+        """Visit augmented assignments: codeflash_output += value."""
+        if self._is_codeflash_output_target(node.target):
+            self._record_assignment(node)
+        self.generic_visit(node)
+
+    def visit_NamedExpr(self, node: ast.NamedExpr) -> None:
+        """Visit walrus operator: (codeflash_output := value)."""
+        if isinstance(node.target, ast.Name) and node.target.id == "codeflash_output":
+            self._record_assignment(node)
+        self.generic_visit(node)
+
+    def visit_For(self, node: ast.For) -> None:
+        """Visit for loops: for codeflash_output in iterable."""
+        if self._is_codeflash_output_target(node.target):
+            self._record_assignment(node)
+        self.generic_visit(node)
+
+    def visit_comprehension(self, node: ast.comprehension) -> None:
+        """Visit comprehensions: [x for codeflash_output in iterable]."""
+        if self._is_codeflash_output_target(node.target):
+            # Comprehensions don't have line numbers, so we skip recording
+            pass
+        self.generic_visit(node)
+
+    def visit_With(self, node: ast.With) -> None:
+        """Visit with statements: with expr as codeflash_output."""
+        for item in node.items:
+            if item.optional_vars and self._is_codeflash_output_target(item.optional_vars):
+                self._record_assignment(node)
+                break
+        self.generic_visit(node)
+
+    def visit_ExceptHandler(self, node: ast.ExceptHandler) -> None:
+        """Visit except handlers: except Exception as codeflash_output."""
+        if node.name == "codeflash_output":
+            self._record_assignment(node)
+        self.generic_visit(node)
+
+
+def find_codeflash_output_assignments(source_code: str) -> list[int]:
+    tree = ast.parse(source_code)
+    visitor = CfoVisitor(source_code)
+    visitor.visit(tree)
+    return visitor.results
+
+
 def add_runtime_comments_to_generated_tests(
     test_cfg: TestConfig,
     generated_tests: GeneratedTestsList,
@@ -49,11 +145,15 @@ def add_runtime_comments_to_generated_tests(
 
     # TODO: reduce for loops to one
     class RuntimeCommentTransformer(cst.CSTTransformer):
-        def __init__(self, test: GeneratedTests, tests_root: Path, rel_tests_root: Path) -> None:
+        def __init__(self, module: cst.Module, test: GeneratedTests, tests_root: Path, rel_tests_root: Path) -> None:
+            super().__init__()
             self.test = test
             self.context_stack: list[str] = []
             self.tests_root = tests_root
             self.rel_tests_root = rel_tests_root
+            self.module = module
+            self.cfo_locs: list[int] = []
+            self.cfo_idx_loc_to_look_at: int = -1
 
         def visit_ClassDef(self, node: cst.ClassDef) -> None:
             # Track when we enter a class
@@ -65,6 +165,13 @@ def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef
             return updated_node
 
         def visit_FunctionDef(self, node: cst.FunctionDef) -> None:
+            # convert function body to ast normalized string and find occurrences of codeflash_output
+            body_code = dedent(self.module.code_for_node(node.body))
+            normalized_body_code = ast.unparse(ast.parse(body_code))
+            self.cfo_locs = sorted(
+                find_codeflash_output_assignments(normalized_body_code)
+            )  # sorted in order we will encounter them
+            self.cfo_idx_loc_to_look_at = -1
             self.context_stack.append(node.name.value)
 
         def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef:  # noqa: ARG002
@@ -91,10 +198,12 @@ def leave_SimpleStatementLine(
 
             if codeflash_assignment_found:
                 # Find matching test cases by looking for this test function name in the test results
+                self.cfo_idx_loc_to_look_at += 1
                 matching_original_times = []
                 matching_optimized_times = []
-                # TODO : will not work if there are multiple test cases with the same name, match filename + test class + test function name
+                # TODO : will not work if there are multiple test cases with the same name, match filename + test class + test function name + invocationid
                 for invocation_id, runtimes in original_runtimes.items():
+                    # get position here and match in if condition
                     qualified_name = (
                         invocation_id.test_class_name + "." + invocation_id.test_function_name  # type: ignore[operator]
                         if invocation_id.test_class_name
@@ -105,13 +214,19 @@ def leave_SimpleStatementLine(
                         .with_suffix(".py")
                         .relative_to(self.rel_tests_root)
                     )
-                    if qualified_name == ".".join(self.context_stack) and rel_path in [
-                        self.test.behavior_file_path.relative_to(self.tests_root),
-                        self.test.perf_file_path.relative_to(self.tests_root),
-                    ]:
+                    if (
+                        qualified_name == ".".join(self.context_stack)
+                        and rel_path
+                        in [
+                            self.test.behavior_file_path.relative_to(self.tests_root),
+                            self.test.perf_file_path.relative_to(self.tests_root),
+                        ]
+                        and int(invocation_id.iteration_id.split("_")[0]) == self.cfo_locs[self.cfo_idx_loc_to_look_at]  # type:ignore[union-attr]
+                    ):
                         matching_original_times.extend(runtimes)
 
                 for invocation_id, runtimes in optimized_runtimes.items():
+                    # get position here and match in if condition
                     qualified_name = (
                         invocation_id.test_class_name + "." + invocation_id.test_function_name  # type: ignore[operator]
                         if invocation_id.test_class_name
@@ -122,10 +237,15 @@ def leave_SimpleStatementLine(
                         .with_suffix(".py")
                         .relative_to(self.rel_tests_root)
                     )
-                    if qualified_name == ".".join(self.context_stack) and rel_path in [
-                        self.test.behavior_file_path.relative_to(self.tests_root),
-                        self.test.perf_file_path.relative_to(self.tests_root),
-                    ]:
+                    if (
+                        qualified_name == ".".join(self.context_stack)
+                        and rel_path
+                        in [
+                            self.test.behavior_file_path.relative_to(self.tests_root),
+                            self.test.perf_file_path.relative_to(self.tests_root),
+                        ]
+                        and int(invocation_id.iteration_id.split("_")[0]) == self.cfo_locs[self.cfo_idx_loc_to_look_at]  # type:ignore[union-attr]
+                    ):
                         matching_optimized_times.extend(runtimes)
 
                 if matching_original_times and matching_optimized_times:
@@ -161,9 +281,8 @@ def leave_SimpleStatementLine(
         try:
             # Parse the test source code
             tree = cst.parse_module(test.generated_original_test_source)
-
             # Transform the tree to add runtime comments
-            transformer = RuntimeCommentTransformer(test, tests_root, rel_tests_root)
+            transformer = RuntimeCommentTransformer(tree, test, tests_root, rel_tests_root)
             modified_tree = tree.visit(transformer)
 
             # Convert back to source code
diff --git a/codeflash/discovery/discover_unit_tests.py b/codeflash/discovery/discover_unit_tests.py
@@ -151,6 +151,15 @@ def __init__(self, function_names_to_find: set[str]) -> None:
         self.has_dynamic_imports: bool = False
         self.wildcard_modules: set[str] = set()
 
+        # Precompute function_names for prefix search
+        # For prefix match, store mapping from prefix-root to candidates for O(1) matching
+        self._exact_names = function_names_to_find
+        self._prefix_roots = {}
+        for name in function_names_to_find:
+            if "." in name:
+                root = name.split(".", 1)[0]
+                self._prefix_roots.setdefault(root, []).append(name)
+
     def visit_Import(self, node: ast.Import) -> None:
         """Handle 'import module' statements."""
         if self.found_any_target_function:
@@ -181,30 +190,46 @@ def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
         if self.found_any_target_function:
             return
 
-        if not node.module:
+        mod = node.module
+        if not mod:
             return
 
+        fnames = self._exact_names
+        proots = self._prefix_roots
+
         for alias in node.names:
-            if alias.name == "*":
-                self.wildcard_modules.add(node.module)
-            else:
-                imported_name = alias.asname if alias.asname else alias.name
-                self.imported_modules.add(imported_name)
-
-                # Check for dynamic import functions
-                if node.module == "importlib" and alias.name == "import_module":
-                    self.has_dynamic_imports = True
-
-                # Check if imported name is a target qualified name
-                if alias.name in self.function_names_to_find:
-                    self.found_any_target_function = True
-                    self.found_qualified_name = alias.name
-                    return
-                # Check if module.name forms a target qualified name
-                qualified_name = f"{node.module}.{alias.name}"
-                if qualified_name in self.function_names_to_find:
+            aname = alias.name
+            if aname == "*":
+                self.wildcard_modules.add(mod)
+                continue
+
+            imported_name = alias.asname if alias.asname else aname
+            self.imported_modules.add(imported_name)
+
+            # Fast check for dynamic import
+            if mod == "importlib" and aname == "import_module":
+                self.has_dynamic_imports = True
+
+            qname = f"{mod}.{aname}"
+
+            # Fast exact match check
+            if aname in fnames:
+                self.found_any_target_function = True
+                self.found_qualified_name = aname
+                return
+            if qname in fnames:
+                self.found_any_target_function = True
+                self.found_qualified_name = qname
+                return
+
+            # Fast prefix match: only for relevant roots
+            prefix = qname + "."
+            # Only bother if one of the targets startswith the prefix-root
+            candidates = proots.get(qname, ())
+            for target_func in candidates:
+                if target_func.startswith(prefix):
                     self.found_any_target_function = True
-                    self.found_qualified_name = qualified_name
+                    self.found_qualified_name = target_func
                     return
 
     def visit_Attribute(self, node: ast.Attribute) -> None:
diff --git a/codeflash/tracer.py b/codeflash/tracer.py
@@ -263,7 +263,7 @@ def __exit__(
             if self.function_count[
                 str(function.file_name)
                 + ":"
-                + (function.class_name + ":" if function.class_name else "")
+                + (function.class_name + "." if function.class_name else "")
                 + function.function_name
             ]
             > 0
@@ -353,7 +353,7 @@ def tracer_logic(self, frame: FrameType, event: str) -> None:  # noqa: PLR0911
             return
         if function_qualified_name not in self.function_count:
             # seeing this function for the first time
-            self.function_count[function_qualified_name] = 0
+            self.function_count[function_qualified_name] = 1
             file_valid = filter_files_optimized(
                 file_path=file_name,
                 tests_root=Path(self.config["tests_root"]),
diff --git a/tests/scripts/end_to_end_test_tracer_replay.py b/tests/scripts/end_to_end_test_tracer_replay.py
@@ -8,7 +8,7 @@ def run_test(expected_improvement_pct: int) -> bool:
     config = TestConfig(
         trace_mode=True,
         min_improvement_x=0.1,
-        expected_unit_tests=1,
+        expected_unit_tests=7,
         coverage_expectations=[
             CoverageExpectation(function_name="funcA", expected_coverage=100.0, expected_lines=[5, 6, 7, 8, 10, 13])
         ],
diff --git a/tests/scripts/end_to_end_test_utilities.py b/tests/scripts/end_to_end_test_utilities.py
@@ -206,8 +206,8 @@ def run_trace_test(cwd: pathlib.Path, config: TestConfig, expected_improvement_p
         return False
 
     functions_traced = re.search(r"Traced (\d+) function calls successfully and replay test created at - (.*)$", stdout)
-    if not functions_traced or int(functions_traced.group(1)) != 4:
-        logging.error("Expected 4 traced functions")
+    if not functions_traced or int(functions_traced.group(1)) != 13:
+        logging.error("Expected 13 traced functions")
         return False
 
     replay_test_path = pathlib.Path(functions_traced.group(2))
diff --git a/tests/test_add_runtime_comments.py b/tests/test_add_runtime_comments.py
