Merge pull request #1641 from codeflash-ai/fix-importlens-bugs

fix: importlens optimization bugs

Author: KRRT7

codeflash/languages/python/context/code_context_extractor.py

@@ -368,7 +368,7 @@ def process_file_context(
     try:
         all_names = primary_qualified_names | secondary_qualified_names
         code_without_unused_defs = remove_unused_definitions_by_function_names(original_code, all_names)
-        code_context = parse_code_and_prune_cst(
+        pruned_module = parse_code_and_prune_cst(
             code_without_unused_defs,
             code_context_type,
             primary_qualified_names,
@@ -379,11 +379,13 @@ def process_file_context(
         logger.debug(f"Error while getting read-only code: {e}")
         return None
 
-    if code_context.strip():
-        if code_context_type != CodeContextType.HASHING:
+    if pruned_module.code.strip():
+        if code_context_type == CodeContextType.HASHING:
+            code_context = ast.unparse(ast.parse(pruned_module.code))
+        else:
             code_context = add_needed_imports_from_module(
                 src_module_code=original_code,
-                dst_module_code=code_context,
+                dst_module_code=pruned_module,
                 src_path=file_path,
                 dst_path=file_path,
                 project_root=project_root_path,
@@ -1280,8 +1282,8 @@ def parse_code_and_prune_cst(
     target_functions: set[str],
     helpers_of_helper_functions: set[str] = set(),  # noqa: B006
     remove_docstrings: bool = False,
-) -> str:
-    """Create a read-only version of the code by parsing and filtering the code to keep only class contextual information, and other module scoped variables."""
+) -> cst.Module:
+    """Parse and filter the code CST, returning the pruned Module."""
     module = cst.parse_module(code)
     defs_with_usages = collect_top_level_defs_with_usages(module, target_functions | helpers_of_helper_functions)
 
@@ -1317,11 +1319,8 @@ def parse_code_and_prune_cst(
     if not found_target:
         raise ValueError("No target functions found in the provided code")
     if filtered_node and isinstance(filtered_node, cst.Module):
-        code = str(filtered_node.code)
-        if code_context_type == CodeContextType.HASHING:
-            code = ast.unparse(ast.parse(code))  # Makes it standard
-        return code
-    return ""
+        return filtered_node
+    raise ValueError("Pruning produced no module")
 
 
 def prune_cst(

codeflash/languages/python/static_analysis/code_extractor.py

@@ -684,7 +684,7 @@ def resolve_star_import(module_name: str, project_root: Path) -> set[str]:
 
 def add_needed_imports_from_module(
     src_module_code: str,
-    dst_module_code: str,
+    dst_module_code: str | cst.Module,
     src_path: Path,
     dst_path: Path,
     project_root: Path,
@@ -696,6 +696,8 @@ def add_needed_imports_from_module(
     if not helper_functions_fqn:
         helper_functions_fqn = {f.fully_qualified_name for f in (helper_functions or [])}
 
+    dst_code_fallback = dst_module_code if isinstance(dst_module_code, str) else dst_module_code.code
+
     src_module_and_package: ModuleNameAndPackage = calculate_module_and_package(project_root, src_path)
     dst_module_and_package: ModuleNameAndPackage = calculate_module_and_package(project_root, dst_path)
 
@@ -715,15 +717,19 @@ def add_needed_imports_from_module(
         cst.parse_module(src_module_code).visit(gatherer)
     except Exception as e:
         logger.error(f"Error parsing source module code: {e}")
-        return dst_module_code
+        return dst_code_fallback
 
     dotted_import_collector = DottedImportCollector()
-    try:
-        parsed_dst_module = cst.parse_module(dst_module_code)
+    if isinstance(dst_module_code, cst.Module):
+        parsed_dst_module = dst_module_code
         parsed_dst_module.visit(dotted_import_collector)
-    except cst.ParserSyntaxError as e:
-        logger.exception(f"Syntax error in destination module code: {e}")
-        return dst_module_code  # Return the original code if there's a syntax error
+    else:
+        try:
+            parsed_dst_module = cst.parse_module(dst_module_code)
+            parsed_dst_module.visit(dotted_import_collector)
+        except cst.ParserSyntaxError as e:
+            logger.exception(f"Syntax error in destination module code: {e}")
+            return dst_code_fallback
 
     try:
         for mod in gatherer.module_imports:
@@ -768,7 +774,7 @@ def add_needed_imports_from_module(
                     RemoveImportsVisitor.remove_unused_import(dst_context, mod, obj)
     except Exception as e:
         logger.exception(f"Error adding imports to destination module code: {e}")
-        return dst_module_code
+        return dst_code_fallback
 
     for mod, asname in gatherer.module_aliases.items():
         if not asname:
@@ -796,7 +802,7 @@ def add_needed_imports_from_module(
         return transformed_module.code.lstrip("\n")
     except Exception as e:
         logger.exception(f"Error adding imports to destination module code: {e}")
-        return dst_module_code
+        return dst_code_fallback
 
 
 def get_code(functions_to_optimize: list[FunctionToOptimize]) -> tuple[str | None, set[tuple[str, str]]]:

codeflash/optimization/function_optimizer.py

@@ -1023,6 +1023,7 @@ def process_single_candidate(
         eval_ctx: CandidateEvaluationContext,
         exp_type: str,
         function_references: str,
+        normalized_original: str,
     ) -> BestOptimization | None:
         """Process a single optimization candidate.
 
@@ -1033,8 +1034,24 @@ def process_single_candidate(
         get_run_tmp_file(Path(f"test_return_values_{candidate_index}.bin")).unlink(missing_ok=True)
         get_run_tmp_file(Path(f"test_return_values_{candidate_index}.sqlite")).unlink(missing_ok=True)
 
-        logger.info(f"h3|Optimization candidate {candidate_index}/{total_candidates}:")
         candidate = candidate_node.candidate
+
+        normalized_code = normalize_code(candidate.source_code.flat.strip())
+
+        if normalized_code == normalized_original:
+            logger.info(f"h3|Candidate {candidate_index}/{total_candidates}: Identical to original code, skipping.")
+            console.rule()
+            return None
+
+        if normalized_code in eval_ctx.ast_code_to_id:
+            logger.info(
+                f"h3|Candidate {candidate_index}/{total_candidates}: Duplicate of a previous candidate, skipping."
+            )
+            eval_ctx.handle_duplicate_candidate(candidate, normalized_code, code_context)
+            console.rule()
+            return None
+
+        logger.info(f"h3|Optimization candidate {candidate_index}/{total_candidates}:")
         # Use correct extension based on language
         ext = self.language_support.file_extensions[0]
         code_print(
@@ -1062,13 +1079,6 @@ def process_single_candidate(
             )
             return None
 
-        # Check for duplicate candidates
-        normalized_code = normalize_code(candidate.source_code.flat.strip())
-        if normalized_code in eval_ctx.ast_code_to_id:
-            logger.info("Current candidate has been encountered before in testing, Skipping optimization candidate.")
-            eval_ctx.handle_duplicate_candidate(candidate, normalized_code, code_context)
-            return None
-
         eval_ctx.register_new_candidate(normalized_code, candidate, code_context)
 
         # Run the optimized candidate
@@ -1238,6 +1248,7 @@ def determine_best_candidate(
             self.future_adaptive_optimizations,
         )
         candidate_index = 0
+        normalized_original = normalize_code(code_context.read_writable_code.flat.strip())
 
         # Process candidates using queue-based approach
         while not processor.is_done():
@@ -1259,6 +1270,7 @@ def determine_best_candidate(
                     eval_ctx=eval_ctx,
                     exp_type=exp_type,
                     function_references=function_references,
+                    normalized_original=normalized_original,
                 )
             except KeyboardInterrupt as e:
                 logger.exception(f"Optimization interrupted: {e}")