prevent duplicates for new global statements

Author: mohammedahmed18

codeflash/code_utils/code_extractor.py

@@ -338,20 +338,28 @@ def delete___future___aliased_imports(module_code: str) -> str:
 
 
 def add_global_assignments(src_module_code: str, dst_module_code: str) -> str:
-    non_assignment_global_statements = extract_global_statements(src_module_code)
-
-    # Find the last import line in target
-    last_import_line = find_last_import_line(dst_module_code)
-
-    # Parse the target code
-    target_module = cst.parse_module(dst_module_code)
-
-    # Create transformer to insert non_assignment_global_statements
-    transformer = ImportInserter(non_assignment_global_statements, last_import_line)
-    #
-    # # Apply transformation
-    modified_module = target_module.visit(transformer)
-    dst_module_code = modified_module.code
+    new_added_global_statements = extract_global_statements(src_module_code)
+    existing_global_statements = extract_global_statements(dst_module_code)
+
+    unique_global_statements = [
+        stmt
+        for stmt in new_added_global_statements
+        if not any(stmt.deep_equals(existing_stmt) for existing_stmt in existing_global_statements)
+    ]
+
+    if unique_global_statements:
+        # Find the last import line in target
+        last_import_line = find_last_import_line(dst_module_code)
+
+        # Parse the target code
+        target_module = cst.parse_module(dst_module_code)
+
+        # Create transformer to insert new statements
+        transformer = ImportInserter(unique_global_statements, last_import_line)
+        #
+        # # Apply transformation
+        modified_module = target_module.visit(transformer)
+        dst_module_code = modified_module.code
 
     # Parse the code
     original_module = cst.parse_module(dst_module_code)

codeflash/code_utils/code_replacer.py

@@ -412,7 +412,6 @@ def replace_function_definitions_in_module(
     module_abspath: Path,
     preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]],
     project_root_path: Path,
-    global_assignments_added_before: bool = False,  # noqa: FBT001, FBT002
 ) -> bool:
     source_code: str = module_abspath.read_text(encoding="utf8")
     code_to_apply = get_optimized_code_for_module(module_abspath.relative_to(project_root_path), optimized_code)
@@ -422,7 +421,7 @@ def replace_function_definitions_in_module(
         # becuase of an "edge case" where the optimized code intoduced a new import and a global assignment using that import
         # and that import wasn't used before, so it was ignored when calling AddImportsVisitor.add_needed_import inside replace_functions_and_add_imports (because the global assignment wasn't added yet)
         # this was added at https://github.com/codeflash-ai/codeflash/pull/448
-        add_global_assignments(code_to_apply, source_code) if not global_assignments_added_before else source_code,
+        add_global_assignments(code_to_apply, source_code),
         function_names,
         code_to_apply,
         module_abspath,

codeflash/context/unused_definition_remover.py

@@ -537,7 +537,6 @@ def revert_unused_helper_functions(
                     module_abspath=file_path,
                     preexisting_objects=set(),  # Empty set since we're reverting
                     project_root_path=project_root,
-                    global_assignments_added_before=True,  # since we revert helpers functions after applying the optimization, we know that the file already has global assignments added, otherwise they would be added twice.
                 )
 
                 if reverted_code:

tests/test_code_replacement.py

@@ -1707,7 +1707,6 @@ def new_function2(value):
 """
     expected_code = """import numpy as np
 
-print("Hello world")
 a=2
 print("Hello world")
 def some_fn():
@@ -1783,7 +1782,6 @@ def new_function2(value):
 """
     expected_code = """import numpy as np
 
-print("Hello world")
 print("Hello world")
 def some_fn():
     a=np.zeros(10)
@@ -1862,7 +1860,6 @@ def new_function2(value):
 """
     expected_code = """import numpy as np
 
-print("Hello world")
 a=3
 print("Hello world")
 def some_fn():
@@ -1940,7 +1937,6 @@ def new_function2(value):
 """
     expected_code = """import numpy as np
 
-print("Hello world")
 a=2
 print("Hello world")
 def some_fn():
@@ -2019,7 +2015,6 @@ def new_function2(value):
 """
     expected_code = """import numpy as np
 
-print("Hello world")
 a=3
 print("Hello world")
 def some_fn():
@@ -2104,7 +2099,6 @@ def new_function2(value):
 """
     expected_code = """import numpy as np
 
-print("Hello world")
 if 2<3:
     a=4
 else: