⚡️ Speed up method QualifiedFunctionUsageMarker._expand_qualified_functions by 53% in PR #161 (benchmark-docs)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #161

If you approve this dependent PR, these changes will be merged into the original PR branch benchmark-docs.

This PR will be automatically closed if the original PR is merged.

---

📄 53% (0.53x) speedup for QualifiedFunctionUsageMarker._expand_qualified_functions in codeflash/context/unused_definition_remover.py

⏱️ Runtime : 1.98 microseconds → 1.29 microsecond (best of 96 runs)

📝 Explanation and details

To optimize this program, we'll focus on reducing the runtime of the _expand_qualified_functions method. Let's analyze the main performance issues based on the provided profiling results.

1. The loop for name in self.definitions is called 22,095,676 times, which is significantly higher than the outer loop iterations (3,396), suggesting inefficiency in handling self.definitions.

2. The name.startswith(f"{class_name}.__") and name.endswith("__") checks are done multiple times and each check is quite expensive within the high number of iterations.

Optimizations.

1. Use more efficient data structures:

   • Convert self.definitions to a preprocessed set or dictionary to quickly check for dunder methods.

2. Preprocess the definitions only once.

   • Instead of checking name.startswith(f"{class_name}.__") and name.endswith("__") inside the loop, preprocess the self.definitions to filter and classify dunder methods by class names.

Optimized Code.

Explanation.

1. We preprocess the definitions once in the _preprocess_definitions function to categorize dunder methods by their class names.
2. Reuse this preprocessed data in the _expand_qualified_functions function to check and expand dunder methods more efficiently.

This significantly reduces the complexity of the loops and the number of checks required during the expansion process.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 29 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	✅ 2 Passed
📊 Tests Coverage

🌀 Generated Regression Tests Details

from __future__ import annotations

# imports
import pytest  # used for our unit tests
from codeflash.context.unused_definition_remover import \
    QualifiedFunctionUsageMarker


# Mock class for UsageInfo
class UsageInfo:
    pass

# unit tests
def test_single_qualified_function_without_class_methods():
    definitions = {}
    qualified_function_names = {"foo"}
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)

def test_single_qualified_function_with_class_methods():
    definitions = {"MyClass.__init__": UsageInfo(), "MyClass.__str__": UsageInfo()}
    qualified_function_names = {"MyClass.my_method"}
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)

def test_multiple_qualified_functions():
    definitions = {"MyClass.__init__": UsageInfo(), "OtherClass.__str__": UsageInfo()}
    qualified_function_names = {"MyClass.my_method", "OtherClass.other_method"}
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)

def test_empty_inputs():
    definitions = {}
    qualified_function_names = set()
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)

def test_qualified_function_names_with_no_corresponding_definitions():
    definitions = {"SomeClass.__init__": UsageInfo()}
    qualified_function_names = {"NonExistentClass.method"}
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)

def test_qualified_function_names_that_are_only_dunder_methods():
    definitions = {"MyClass.__init__": UsageInfo(), "MyClass.__str__": UsageInfo()}
    qualified_function_names = {"MyClass.__init__", "MyClass.__str__"}
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)

def test_nested_classes():
    definitions = {"OuterClass.__init__": UsageInfo(), "OuterClass.InnerClass.__init__": UsageInfo()}
    qualified_function_names = {"OuterClass.InnerClass.method"}
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)

def test_inheritance_with_overridden_methods():
    definitions = {"BaseClass.__init__": UsageInfo(), "DerivedClass.__init__": UsageInfo()}
    qualified_function_names = {"BaseClass.method", "DerivedClass.method"}
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)

def test_class_and_method_names_with_underscores():
    definitions = {"My_Class.__init__": UsageInfo(), "My_Class.__str__": UsageInfo()}
    qualified_function_names = {"My_Class.my_method"}
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)

def test_class_and_method_names_with_numbers():
    definitions = {"Class123.__init__": UsageInfo(), "Class123.__str__": UsageInfo()}
    qualified_function_names = {"Class123.method456"}
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)

def test_large_number_of_qualified_functions():
    definitions = {f"Class{i}.__init__": UsageInfo() for i in range(1000)}
    qualified_function_names = {f"Class{i}.method{i}" for i in range(1000)}
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)
    expected_expanded = set(qualified_function_names)
    for i in range(1000):
        expected_expanded.add(f"Class{i}")
        expected_expanded.add(f"Class{i}.__init__")

def test_large_number_of_definitions():
    definitions = {f"Class{i}.__init__": UsageInfo() for i in range(1000)}
    qualified_function_names = {"Class0.method0"}
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)
    expected_expanded = {"Class0.method0", "Class0", "Class0.__init__"}

def test_stress_test_with_maximum_input_size():
    definitions = {f"Class{i}.__init__": UsageInfo() for i in range(1000)}
    qualified_function_names = {f"Class{i}.method{i}" for i in range(1000)}
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)
    expected_expanded = set(qualified_function_names)
    for i in range(1000):
        expected_expanded.add(f"Class{i}")
        expected_expanded.add(f"Class{i}.__init__")

def test_performance_with_deeply_nested_class_structures():
    definitions = {
        "OuterClass.__init__": UsageInfo(),
        "OuterClass.InnerClass1.__init__": UsageInfo(),
        "OuterClass.InnerClass1.InnerClass2.__init__": UsageInfo()
    }
    qualified_function_names = {"OuterClass.InnerClass1.InnerClass2.method"}
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)
    expected_expanded = {
        "OuterClass.InnerClass1.InnerClass2.method",
        "OuterClass",
        "OuterClass.__init__",
        "OuterClass.InnerClass1",
        "OuterClass.InnerClass1.__init__",
        "OuterClass.InnerClass1.InnerClass2",
        "OuterClass.InnerClass1.InnerClass2.__init__"
    }


def test_malformed_qualified_function_names():
    definitions = {}
    qualified_function_names = {"IncompleteClassName."}
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)

def test_consistent_output_for_same_input():
    definitions = {"MyClass.__init__": UsageInfo(), "MyClass.__str__": UsageInfo()}
    qualified_function_names = {"MyClass.my_method"}
    marker1 = QualifiedFunctionUsageMarker(definitions, qualified_function_names)
    marker2 = QualifiedFunctionUsageMarker(definitions, qualified_function_names)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from __future__ import annotations

# imports
import pytest  # used for our unit tests
from codeflash.context.unused_definition_remover import \
    QualifiedFunctionUsageMarker


class UsageInfo:
    pass
from codeflash.context.unused_definition_remover import \
    QualifiedFunctionUsageMarker

# unit tests

def test_single_class_method():
    definitions = {"MyClass": UsageInfo(), "MyClass.my_method": UsageInfo()}
    qualified_function_names = {"MyClass.my_method"}
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)

def test_multiple_class_methods():
    definitions = {"MyClass": UsageInfo(), "MyClass.method1": UsageInfo(), "MyClass.method2": UsageInfo()}
    qualified_function_names = {"MyClass.method1", "MyClass.method2"}
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)

def test_class_with_dunder_methods():
    definitions = {"MyClass": UsageInfo(), "MyClass.my_method": UsageInfo(), "MyClass.__init__": UsageInfo(), "MyClass.__str__": UsageInfo()}
    qualified_function_names = {"MyClass.my_method"}
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)

def test_empty_qualified_function_names():
    definitions = {"MyClass": UsageInfo(), "MyClass.my_method": UsageInfo()}
    qualified_function_names = set()
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)

def test_non_existent_class_methods():
    definitions = {"MyClass": UsageInfo(), "MyClass.my_method": UsageInfo()}
    qualified_function_names = {"NonExistentClass.method"}
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)

def test_nested_class_methods():
    definitions = {"OuterClass": UsageInfo(), "OuterClass.InnerClass": UsageInfo(), "OuterClass.InnerClass.method": UsageInfo()}
    qualified_function_names = {"OuterClass.InnerClass.method"}
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)

def test_multiple_levels_of_nesting():
    definitions = {"OuterClass": UsageInfo(), "OuterClass.method": UsageInfo(), "OuterClass.InnerClass": UsageInfo(), "OuterClass.InnerClass.method": UsageInfo()}
    qualified_function_names = {"OuterClass.InnerClass.method", "OuterClass.method"}
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)

def test_large_number_of_methods_and_classes():
    definitions = {f"Class{i}": UsageInfo() for i in range(100)}
    definitions.update({f"Class{i}.method{j}": UsageInfo() for i in range(100) for j in range(10)})
    definitions.update({f"Class{i}.__init__": UsageInfo() for i in range(100)})
    qualified_function_names = {f"Class{i}.method{j}" for i in range(100) for j in range(10)}
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)
    expected = qualified_function_names | {f"Class{i}" for i in range(100)} | {f"Class{i}.__init__" for i in range(100)}

def test_methods_with_special_characters():
    definitions = {"MyClass": UsageInfo(), "MyClass.my_method$": UsageInfo(), "MyClass._private_method": UsageInfo()}
    qualified_function_names = {"MyClass.my_method$", "MyClass._private_method"}
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)

def test_class_methods_and_standalone_functions():
    definitions = {"MyClass": UsageInfo(), "MyClass.my_method": UsageInfo(), "standalone_function": UsageInfo()}
    qualified_function_names = {"MyClass.my_method", "standalone_function"}
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)


def test_valid_and_invalid_class_methods():
    definitions = {"ValidClass": UsageInfo(), "ValidClass.method": UsageInfo()}
    qualified_function_names = {"ValidClass.method", "InvalidClass.method"}
    marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from codeflash.context.unused_definition_remover import QualifiedFunctionUsageMarker
from codeflash.context.unused_definition_remover import UsageInfo

def test_QualifiedFunctionUsageMarker__expand_qualified_functions():
    QualifiedFunctionUsageMarker._expand_qualified_functions(QualifiedFunctionUsageMarker({'.__': UsageInfo('', used_by_qualified_function=False, dependencies={''})}, {'\x01', '.', '\x00'}))

To edit these changes git checkout codeflash/optimize-pr161-2025-04-19T20.05.06 and push.