works, optimization list length will be incorrectly displayed as we are dynamically popping and appending to the list, right now appending to just the control opt candidates, need to refactor that

Author: aseembits93

code_to_optimize/bubble_sort_deps.py

@@ -1,11 +1,15 @@
-from code_to_optimize.bubble_sort_dep1_helper import dep1_comparer
-from code_to_optimize.bubble_sort_dep2_swap import dep2_swap
-
-
 def sorter_deps(arr):
-    for i in range(len(arr)):
-        for j in range(len(arr) - 1):
-            if dep1_comparer(arr, j):
-                dep2_swap(arr, j)
+    n = len(arr)
+    for i in range(n):
+        # We use a flag to check if the array is already sorted
+        swapped = False
+        # Reduce the range of j, since the last i elements are already sorted
+        for j in range(n - 1 - i):
+            if arr[j] > arr[j + 1]:
+                # Swap without a helper function
+                arr[j], arr[j + 1] = arr[j + 1], arr[j]
+                swapped = True
+        # If no elements were swapped in the inner loop, break
+        if not swapped:
+            break
     return arr
-

code_to_optimize/bubble_sort_nested_classmethod.py

@@ -0,0 +1,8 @@
+from code_to_optimize.bubble_sort_in_nested_class import WrapperClass
+from line_profiler import profile as codeflash_line_profile
+
+
+@codeflash_line_profile
+def sort_classmethod(x):
+    y = WrapperClass.BubbleSortClass()
+    return y.sorter(x)

code_to_optimize/tests/pytest/test_bubble_sort_deps.py

@@ -0,0 +1,15 @@
+from code_to_optimize.bubble_sort_deps import sorter_deps
+
+
+def test_sort():
+    input = [5, 4, 3, 2, 1, 0]
+    output = sorter_deps(input)
+    assert output == [0, 1, 2, 3, 4, 5]
+
+    input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0]
+    output = sorter_deps(input)
+    assert output == [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]
+
+    input = list(reversed(range(5000)))
+    output = sorter_deps(input)
+    assert output == list(range(5000))

codeflash/code_utils/lprof_utils.py renamed to codeflash/code_utils/line_profile_utils.py

@@ -1,14 +1,18 @@
+"""Adapted from line_profiler (https://github.com/pyutils/line_profiler) written by Enthought, Inc. (BSD License)"""
+from collections import defaultdict
+
 import isort
 import libcst as cst
 from pathlib import Path
-from typing import Union
+from typing import Union, List
+from libcst import ImportFrom, ImportAlias, Name
 
 from codeflash.code_utils.code_utils import get_run_tmp_file
 
 
-class DecoratorAdder(cst.CSTTransformer):
+class LineProfilerDecoratorAdder(cst.CSTTransformer):
     """Transformer that adds a decorator to a function with a specific qualified name."""
-
+    #Todo we don't support nested functions yet so they can only be inside classes, dont use qualified names, instead use the structure
     def __init__(self, qualified_name: str, decorator_name: str):
         """
         Initialize the transformer.
@@ -21,7 +25,7 @@ def __init__(self, qualified_name: str, decorator_name: str):
         self.qualified_name_parts = qualified_name.split(".")
         self.decorator_name = decorator_name
 
-        # Track our current context path
+        # Track our current context path, only add when we encounter a class
         self.context_stack = []
 
     def visit_ClassDef(self, node: cst.ClassDef) -> None:
@@ -83,6 +87,64 @@ def _is_target_decorator(self, decorator_node: Union[cst.Name, cst.Attribute, cs
             return decorator_node.func.value == self.decorator_name
         return False
 
+class ProfileEnableTransformer(cst.CSTTransformer):
+    def __init__(self,filename):
+        # Flag to track if we found the import statement
+        self.found_import = False
+        # Track indentation of the import statement
+        self.import_indentation = None
+        self.filename = filename
+
+    def leave_ImportFrom(self, original_node: cst.ImportFrom, updated_node: cst.ImportFrom) -> cst.ImportFrom:
+        # Check if this is the line profiler import statement
+        if (isinstance(original_node.module, cst.Name) and
+                original_node.module.value == "line_profiler" and
+                any(name.name.value == "profile" and
+                    (not name.asname or name.asname.name.value == "codeflash_line_profile")
+                    for name in original_node.names)):
+
+            self.found_import = True
+            # Get the indentation from the original node
+            if hasattr(original_node, "leading_lines"):
+                leading_whitespace = original_node.leading_lines[-1].whitespace if original_node.leading_lines else ""
+                self.import_indentation = leading_whitespace
+
+        return updated_node
+
+    def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module:
+        if not self.found_import:
+            return updated_node
+
+        # Create a list of statements from the original module
+        new_body = list(updated_node.body)
+
+        # Find the index of the import statement
+        import_index = None
+        for i, stmt in enumerate(new_body):
+            if isinstance(stmt, cst.SimpleStatementLine):
+                for small_stmt in stmt.body:
+                    if isinstance(small_stmt, cst.ImportFrom):
+                        if (isinstance(small_stmt.module, cst.Name) and
+                                small_stmt.module.value == "line_profiler" and
+                                any(name.name.value == "profile" and
+                                    (not name.asname or name.asname.name.value == "codeflash_line_profile")
+                                    for name in small_stmt.names)):
+                            import_index = i
+                            break
+                if import_index is not None:
+                    break
+
+        if import_index is not None:
+            # Create the new enable statement to insert after the import
+            enable_statement = cst.parse_statement(
+                f"codeflash_line_profile.enable(output_prefix='{self.filename}')"
+            )
+
+            # Insert the new statement after the import statement
+            new_body.insert(import_index + 1, enable_statement)
+
+        # Create a new module with the updated body
+        return updated_node.with_changes(body=new_body)
 
 def add_decorator_to_qualified_function(module, qualified_name, decorator_name):
     """
@@ -99,48 +161,22 @@ def add_decorator_to_qualified_function(module, qualified_name, decorator_name):
     # Parse the source code into a CST
 
     # Apply our transformer
-    transformer = DecoratorAdder(qualified_name, decorator_name)
+    transformer = LineProfilerDecoratorAdder(qualified_name, decorator_name)
     modified_module = module.visit(transformer)
 
     # Convert the modified CST back to source code
     return modified_module
 
-def add_profile_enable(original_code: str, db_file: str) -> str:
+def add_profile_enable(original_code: str, line_profile_output_file: str) -> str:
+    # todo modify by using a libcst transformer
     module = cst.parse_module(original_code)
-    found_index = -1
-
-    for idx, statement in enumerate(module.body):
-        if isinstance(statement, cst.SimpleStatementLine):
-            for stmt in statement.body:
-                if isinstance(stmt, cst.ImportFrom):
-                    if stmt.module and stmt.module.value == 'line_profiler':
-                        for name in stmt.names:
-                            if isinstance(name, cst.ImportAlias):
-                                if name.name.value == 'profile' and name.asname is None:
-                                    found_index = idx
-                                    break
-                        if found_index != -1:
-                            break
-        if found_index != -1:
-            break
-
-    if found_index == -1:
-        return original_code  # or raise an exception if the import is not found
-
-    # Create the new line to insert
-    new_line = f"profile.enable(output_prefix='{db_file}')\n"
-    new_statement = cst.parse_statement(new_line)
-
-    # Insert the new statement into the module's body
-    new_body = list(module.body)
-    new_body.insert(found_index + 1, new_statement)
-    modified_module = module.with_changes(body=new_body)
-
+    transformer = ProfileEnableTransformer(line_profile_output_file)
+    modified_module = module.visit(transformer)
     return modified_module.code
 
 
 class ImportAdder(cst.CSTTransformer):
-    def __init__(self, import_statement='from line_profiler import profile'):
+    def __init__(self, import_statement):
         self.import_statement = import_statement
         self.has_import = False
 
@@ -166,45 +202,36 @@ def visit_ImportFrom(self, node):
 
 
 def add_decorator_imports(function_to_optimize, code_context):
+    """Adds a profile decorator to a function in a Python file and all its helper functions."""
     #self.function_to_optimize, file_path_to_helper_classes, self.test_cfg.tests_root
-    #todo change function signature to get filepaths of fn, helpers and db
-    # modify libcst parser to visit with qualified name
-    file_paths = list()
-    fn_list = list()
-    db_file = get_run_tmp_file(Path("baseline"))
-    file_paths.append(function_to_optimize.file_path)
-    fn_list.append(function_to_optimize.qualified_name)
+    #grouped iteration, file to fns to optimize, from line_profiler import profile as codeflash_line_profile
+    file_paths = defaultdict(list)
+    line_profile_output_file = get_run_tmp_file(Path("baseline_lprof"))
+    file_paths[function_to_optimize.file_path].append(function_to_optimize.qualified_name)
     for elem in code_context.helper_functions:
-        file_paths.append(elem.file_path)
-        fn_list.append(elem.qualified_name)
-    """Adds a decorator to a function in a Python file."""
-    for file_path, fn_name in zip(file_paths, fn_list):
+        file_paths[elem.file_path].append(elem.qualified_name)
+    for file_path,fns_present in file_paths.items():
         #open file
         with open(file_path, "r", encoding="utf-8") as file:
             file_contents = file.read()
         # parse to cst
         module_node = cst.parse_module(file_contents)
-        # add decorator
-        module_node = add_decorator_to_qualified_function(module_node, fn_name, 'profile')
+        for fn_name in fns_present:
+            # add decorator
+            module_node = add_decorator_to_qualified_function(module_node, fn_name, 'codeflash_line_profile')
         # add imports
         # Create a transformer to add the import
-        transformer = ImportAdder("from line_profiler import profile")
+        transformer = ImportAdder("from line_profiler import profile as codeflash_line_profile")
         # Apply the transformer to add the import
         module_node = module_node.visit(transformer)
         modified_code = isort.code(module_node.code, float_to_top=True)
         # write to file
         with open(file_path, "w", encoding="utf-8") as file:
             file.write(modified_code)
-    #Adding profile.enable line for changing the savepath of the data, do this only for the main file and not the helper files, can use libcst seems like an overkill, will go just with some simple string manipulation
-    with open(file_paths[0],'r') as f:
+    #Adding profile.enable line for changing the savepath of the data, do this only for the main file and not the helper files
+    with open(function_to_optimize.file_path,'r') as f:
         file_contents = f.read()
-    modified_code = add_profile_enable(file_contents,db_file)
-    with open(file_paths[0],'w') as f:
+    modified_code = add_profile_enable(file_contents,str(line_profile_output_file))
+    with open(function_to_optimize.file_path,'w') as f:
         f.write(modified_code)
-    return db_file
-
-
-def prepare_lprofiler_files(prefix: str = "") -> tuple[Path]:
-    """Prepare line profiler output file."""
-    lprofiler_database_file = get_run_tmp_file(Path(prefix))
-    return lprofiler_database_file
+    return line_profile_output_file