|
1 | 1 | from __future__ import annotations |
2 | 2 |
|
| 3 | +import ast |
| 4 | +from collections import defaultdict |
3 | 5 | from dataclasses import dataclass, field |
| 6 | +from pathlib import Path |
4 | 7 | from typing import Optional |
5 | 8 |
|
6 | 9 | import libcst as cst |
7 | 10 |
|
| 11 | +from codeflash.cli_cmds.console import logger |
| 12 | +from codeflash.code_utils.code_replacer import replace_function_definitions_in_module |
| 13 | +from codeflash.models.models import CodeOptimizationContext, FunctionSource |
| 14 | + |
8 | 15 |
|
9 | 16 | @dataclass |
10 | 17 | class UsageInfo: |
@@ -480,3 +487,210 @@ def print_definitions(definitions: dict[str, UsageInfo]) -> None: |
480 | 487 | print(f" Used by qualified function: {info.used_by_qualified_function}") |
481 | 488 | print(f" Dependencies: {', '.join(sorted(info.dependencies)) if info.dependencies else 'None'}") |
482 | 489 | print() |
| 490 | + |
| 491 | + |
| 492 | +def revert_unused_helper_functions( |
| 493 | + project_root, unused_helpers: list[FunctionSource], original_helper_code: dict[Path, str] |
| 494 | +) -> None: |
| 495 | + """Revert unused helper functions back to their original definitions. |
| 496 | +
|
| 497 | + Args: |
| 498 | + unused_helpers: List of unused helper functions to revert |
| 499 | + original_helper_code: Dictionary mapping file paths to their original code |
| 500 | +
|
| 501 | + """ |
| 502 | + if not unused_helpers: |
| 503 | + return |
| 504 | + |
| 505 | + logger.info(f"Reverting {len(unused_helpers)} unused helper function(s) to original definitions") |
| 506 | + |
| 507 | + # Group unused helpers by file path |
| 508 | + unused_helpers_by_file = defaultdict(list) |
| 509 | + for helper in unused_helpers: |
| 510 | + unused_helpers_by_file[helper.file_path].append(helper) |
| 511 | + |
| 512 | + # For each file, revert the unused helper functions to their original definitions |
| 513 | + for file_path, helpers_in_file in unused_helpers_by_file.items(): |
| 514 | + if file_path in original_helper_code: |
| 515 | + try: |
| 516 | + # Read current file content |
| 517 | + current_code = file_path.read_text(encoding="utf8") |
| 518 | + |
| 519 | + # Get original code for this file |
| 520 | + original_code = original_helper_code[file_path] |
| 521 | + |
| 522 | + # Use the code replacer to selectively revert only the unused helper functions |
| 523 | + helper_names = [helper.qualified_name for helper in helpers_in_file] |
| 524 | + reverted_code = replace_function_definitions_in_module( |
| 525 | + function_names=helper_names, |
| 526 | + optimized_code=original_code, # Use original code as the "optimized" code to revert |
| 527 | + module_abspath=file_path, |
| 528 | + preexisting_objects=set(), # Empty set since we're reverting |
| 529 | + project_root_path=project_root, |
| 530 | + ) |
| 531 | + |
| 532 | + if reverted_code: |
| 533 | + logger.debug(f"Reverted unused helpers in {file_path}: {', '.join(helper_names)}") |
| 534 | + |
| 535 | + except Exception as e: |
| 536 | + logger.error(f"Error reverting unused helpers in {file_path}: {e}") |
| 537 | + |
| 538 | + |
| 539 | +def _analyze_imports_in_optimized_code( |
| 540 | + optimized_ast: ast.AST, code_context: CodeOptimizationContext |
| 541 | +) -> dict[str, set[str]]: |
| 542 | + """Analyze import statements in optimized code to map imported names to qualified helper names. |
| 543 | +
|
| 544 | + Args: |
| 545 | + optimized_ast: The AST of the optimized code |
| 546 | + code_context: The code optimization context containing helper functions |
| 547 | +
|
| 548 | + Returns: |
| 549 | + Dictionary mapping imported names to sets of possible qualified helper names |
| 550 | +
|
| 551 | + """ |
| 552 | + imported_names_map = defaultdict(set) |
| 553 | + |
| 554 | + # Create a lookup of helper functions by their simple names and file paths |
| 555 | + helpers_by_name = defaultdict(list) |
| 556 | + helpers_by_file = defaultdict(list) |
| 557 | + |
| 558 | + for helper in code_context.helper_functions: |
| 559 | + if helper.jedi_definition.type != "class": |
| 560 | + helpers_by_name[helper.only_function_name].append(helper) |
| 561 | + module_name = helper.file_path.stem |
| 562 | + helpers_by_file[module_name].append(helper) |
| 563 | + |
| 564 | + # Analyze import statements in the optimized code |
| 565 | + for node in ast.walk(optimized_ast): |
| 566 | + if isinstance(node, ast.ImportFrom): |
| 567 | + # Handle "from module import function" statements |
| 568 | + if node.module: |
| 569 | + module_name = node.module |
| 570 | + for alias in node.names: |
| 571 | + imported_name = alias.asname if alias.asname else alias.name |
| 572 | + original_name = alias.name |
| 573 | + |
| 574 | + # Find helpers that match this import |
| 575 | + for helper in helpers_by_file.get(module_name, []): |
| 576 | + if helper.only_function_name == original_name: |
| 577 | + imported_names_map[imported_name].add(helper.qualified_name) |
| 578 | + imported_names_map[imported_name].add(helper.fully_qualified_name) |
| 579 | + |
| 580 | + elif isinstance(node, ast.Import): |
| 581 | + # Handle "import module" statements |
| 582 | + for alias in node.names: |
| 583 | + imported_name = alias.asname if alias.asname else alias.name |
| 584 | + module_name = alias.name |
| 585 | + |
| 586 | + # For "import module" statements, functions would be called as module.function |
| 587 | + for helper in helpers_by_file.get(module_name, []): |
| 588 | + full_call = f"{imported_name}.{helper.only_function_name}" |
| 589 | + imported_names_map[full_call].add(helper.qualified_name) |
| 590 | + imported_names_map[full_call].add(helper.fully_qualified_name) |
| 591 | + |
| 592 | + return dict(imported_names_map) |
| 593 | + |
| 594 | + |
| 595 | +def detect_unused_helper_functions( |
| 596 | + function_to_optimize, code_context: CodeOptimizationContext, optimized_code: str |
| 597 | +) -> list[FunctionSource]: |
| 598 | + """Detect helper functions that are no longer called by the optimized entrypoint function. |
| 599 | +
|
| 600 | + Args: |
| 601 | + code_context: The code optimization context containing helper functions |
| 602 | + optimized_code: The optimized code to analyze |
| 603 | +
|
| 604 | + Returns: |
| 605 | + List of FunctionSource objects representing unused helper functions |
| 606 | +
|
| 607 | + """ |
| 608 | + try: |
| 609 | + # Parse the optimized code to analyze function calls and imports |
| 610 | + optimized_ast = ast.parse(optimized_code) |
| 611 | + |
| 612 | + # Find the optimized entrypoint function |
| 613 | + entrypoint_function_ast = None |
| 614 | + for node in ast.walk(optimized_ast): |
| 615 | + if isinstance(node, ast.FunctionDef) and node.name == function_to_optimize.function_name: |
| 616 | + entrypoint_function_ast = node |
| 617 | + break |
| 618 | + |
| 619 | + if not entrypoint_function_ast: |
| 620 | + logger.debug(f"Could not find entrypoint function {function_to_optimize.function_name} in optimized code") |
| 621 | + return [] |
| 622 | + |
| 623 | + # First, analyze imports to build a mapping of imported names to their original qualified names |
| 624 | + imported_names_map = _analyze_imports_in_optimized_code(optimized_ast, code_context) |
| 625 | + |
| 626 | + # Extract all function calls in the entrypoint function |
| 627 | + called_function_names = set() |
| 628 | + for node in ast.walk(entrypoint_function_ast): |
| 629 | + if isinstance(node, ast.Call): |
| 630 | + if isinstance(node.func, ast.Name): |
| 631 | + # Regular function call: function_name() |
| 632 | + called_name = node.func.id |
| 633 | + called_function_names.add(called_name) |
| 634 | + # Also add the qualified name if this is an imported function |
| 635 | + if called_name in imported_names_map: |
| 636 | + called_function_names.update(imported_names_map[called_name]) |
| 637 | + elif isinstance(node.func, ast.Attribute): |
| 638 | + # Method call: obj.method() or self.method() or module.function() |
| 639 | + if isinstance(node.func.value, ast.Name): |
| 640 | + if node.func.value.id == "self": |
| 641 | + # self.method_name() -> add both method_name and ClassName.method_name |
| 642 | + called_function_names.add(node.func.attr) |
| 643 | + # For class methods, also add the qualified name |
| 644 | + if hasattr(function_to_optimize, "parents") and function_to_optimize.parents: |
| 645 | + class_name = function_to_optimize.parents[0].name |
| 646 | + called_function_names.add(f"{class_name}.{node.func.attr}") |
| 647 | + else: |
| 648 | + # obj.method() or module.function() |
| 649 | + attr_name = node.func.attr |
| 650 | + called_function_names.add(attr_name) |
| 651 | + called_function_names.add(f"{node.func.value.id}.{attr_name}") |
| 652 | + # Check if this is a module.function call that maps to a helper |
| 653 | + full_call = f"{node.func.value.id}.{attr_name}" |
| 654 | + if full_call in imported_names_map: |
| 655 | + called_function_names.update(imported_names_map[full_call]) |
| 656 | + # Handle nested attribute access like obj.attr.method() |
| 657 | + else: |
| 658 | + called_function_names.add(node.func.attr) |
| 659 | + |
| 660 | + logger.debug(f"Functions called in optimized entrypoint: {called_function_names}") |
| 661 | + logger.debug(f"Imported names mapping: {imported_names_map}") |
| 662 | + |
| 663 | + # Find helper functions that are no longer called |
| 664 | + unused_helpers = [] |
| 665 | + for helper_function in code_context.helper_functions: |
| 666 | + if helper_function.jedi_definition.type != "class": |
| 667 | + # Check if the helper function is called using multiple name variants |
| 668 | + helper_qualified_name = helper_function.qualified_name |
| 669 | + helper_simple_name = helper_function.only_function_name |
| 670 | + helper_fully_qualified_name = helper_function.fully_qualified_name |
| 671 | + |
| 672 | + # Create a set of all possible names this helper might be called by |
| 673 | + possible_call_names = {helper_qualified_name, helper_simple_name, helper_fully_qualified_name} |
| 674 | + |
| 675 | + # For cross-file helpers, also consider module-based calls |
| 676 | + if helper_function.file_path != function_to_optimize.file_path: |
| 677 | + # Add potential module.function combinations |
| 678 | + module_name = helper_function.file_path.stem |
| 679 | + possible_call_names.add(f"{module_name}.{helper_simple_name}") |
| 680 | + |
| 681 | + # Check if any of the possible names are in the called functions |
| 682 | + is_called = bool(possible_call_names.intersection(called_function_names)) |
| 683 | + |
| 684 | + if not is_called: |
| 685 | + unused_helpers.append(helper_function) |
| 686 | + logger.debug(f"Helper function {helper_qualified_name} is not called in optimized code") |
| 687 | + logger.debug(f" Checked names: {possible_call_names}") |
| 688 | + else: |
| 689 | + logger.debug(f"Helper function {helper_qualified_name} is still called in optimized code") |
| 690 | + logger.debug(f" Called via: {possible_call_names.intersection(called_function_names)}") |
| 691 | + |
| 692 | + return unused_helpers |
| 693 | + |
| 694 | + except Exception as e: |
| 695 | + logger.debug(f"Error detecting unused helper functions: {e}") |
| 696 | + return [] |
0 commit comments