⚡️ Speed up method NumericalUsageChecker.visit_Attribute by 228% in PR #1051 (detect-numerical-code)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1051

If you approve this dependent PR, these changes will be merged into the original PR branch detect-numerical-code.

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

---

📄 228% (2.28x) speedup for NumericalUsageChecker.visit_Attribute in codeflash/code_utils/code_extractor.py

⏱️ Runtime : 1.12 milliseconds → 341 microseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves a 227% speedup (from 1.12ms to 341μs) through two key optimizations:

1. Recursive to Iterative Conversion in _get_root_name (~46% of original time saved)

The original implementation used recursion to traverse attribute chains like np.linalg.norm, making a new function call for each level. The line profiler shows return self._get_root_name(node.value) consumed 46.9% of the function's time with 725 recursive calls.

The optimized version replaces recursion with a while True loop that iteratively descends through Attribute nodes by reassigning node = node.value. This eliminates Python's function call overhead (stack frame creation, argument passing, return handling). The performance impact is visible: the iterative version completes in 1.65ms vs 2.07ms for the original _get_root_name.

2. Conditional generic_visit Call (~47% of visit_Attribute time saved)

The original code unconditionally called self.generic_visit(node) for every non-matching attribute (509 calls taking 46.9% of visit_Attribute time). However, generic_visit recursively traverses child nodes—unnecessary work when _get_root_name returns a valid name (the attribute chain already led to a Name node with no further children to visit).

The optimized version only calls generic_visit(node) when root_name is None, which occurs for non-standard AST node types (like Call, Subscript, BinOp). This reduces generic_visit calls from 509 to just 5 in the profiled run, cutting this operation's contribution from 46.9% to 3.1% of total time.

Test Case Performance Patterns

The optimization is most effective for:

• Deeply nested attributes: test_large_scale_deeply_nested_chain shows 172% speedup (28.3μs → 10.4μs) as recursion elimination scales with chain depth
• Many non-matching attributes: test_large_scale_library_not_in_large_set shows 335% speedup by avoiding unnecessary generic_visit traversals
• Sequential visits: test_large_scale_sequential_visits shows 329% speedup (379μs → 88.5μs) for 200 visits, demonstrating consistent per-call improvement

The optimization has minimal impact on edge cases with non-Name/non-Attribute root nodes (like Subscript, BinOp) where generic_visit must still be called, but these are rare in typical AST traversals.

Impact Assessment

These optimizations are particularly valuable when NumericalUsageChecker is used in hot paths analyzing large codebases with many attribute accesses, as the per-call overhead reduction compounds across thousands of AST nodes.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

import ast

from codeflash.code_utils.code_extractor import NumericalUsageChecker


# Helper to parse code and run the visitor on the module AST.
def _run_checker_on_code(code: str, numerical_names: set[str]) -> bool:
    """Parse the given code into an AST, run NumericalUsageChecker on it,
    and return whether any numerical usage was found.
    """
    tree = ast.parse(code)
    checker = NumericalUsageChecker(numerical_names=numerical_names)
    checker.visit(tree)  # triggers visit_Attribute when Attributes are present
    return checker.found_numerical


# -------------------------
# Basic Test Cases
# -------------------------


def test_basic_direct_numerical_attribute_detects_np_array():
    # Basic scenario: direct usage like `np.array(...)` should be detected.
    code = "import numpy as np\nx = np.array([1, 2, 3])"


def test_basic_nested_attribute_detects_deep_chain():
    # Basic scenario with nested attributes: `np.linalg.norm(...)` should be detected.
    code = "import numpy as np\nresult = np.linalg.norm(vec)"


def test_non_numerical_attribute_does_not_trigger_detection():
    # Basic negative test: attribute access on a non-numerical root should not trigger detection.
    code = "class Foo: pass\nobj = Foo()\nobj.attribute"  # attribute access on 'obj'


# -------------------------
# Edge Test Cases
# -------------------------


def test_attribute_with_call_value_does_not_infer_root_name():
    # Edge: Attribute whose value is a Call (e.g., get_np().array) should NOT be interpreted
    # as having a root Name for purposes of _get_root_name. Even if get_np is in numerical_names,
    # the _get_root_name should return None for a Call node and thus not mark found_numerical.
    code = "def get_np():\n    return None\nget_np().array"


def test_case_sensitivity_of_root_name_matching():
    # Edge: name matching should be case-sensitive. 'NP' should not match 'np'.
    code = "NP.array([1, 2, 3])"


def test_long_identifier_root_detection():
    # Edge: very long identifier names should still be handled correctly.
    long_name = "n" * 200  # 200-character name (within reasonable test limits)
    code = f"{long_name}.someattr"


def test_attribute_chain_root_detection_multiple_levels():
    # Edge: Attribute chains of arbitrary depth should resolve back to the true root Name.
    # e.g., a.b.c.d -> root is 'a'
    code = "a.b.c.d"  # root 'a'


def test_visit_attribute_short_circuits_when_already_found():
    # Edge: visit_Attribute should return immediately when found_numerical is already True.
    # Create a checker manually and prime it as if numerical usage was already found.
    checker = NumericalUsageChecker(numerical_names={"np"})
    checker.found_numerical = True  # simulate prior detection
    # Build an Attribute node for `obj.attr` directly from parsing
    node_module = ast.parse("obj.attr")
    # Extract the Attribute node (Module -> Expr -> Attribute)
    attr_node = node_module.body[0].value
    # Call visit_Attribute directly; it should return quickly and leave found_numerical True.
    checker.visit_Attribute(attr_node)  # 460ns -> 441ns (4.31% faster)


# -------------------------
# Large Scale Test Cases
# -------------------------


def test_large_scale_many_attributes_with_single_numerical():
    # Large-scale test: build many attribute expressions to ensure performance and correct detection.
    # Respect the constraint: do not exceed 1000 steps / elements. We'll use 500.
    parts = []
    count = 500  # number of non-numerical attribute expressions
    for i in range(count):
        parts.append(f"a{i}.x")  # many non-numerical attribute usages
    # Insert one numerical usage in the middle to ensure the visitor finds it among many nodes.
    parts.insert(count // 2, "np.special_function()")
    code = "\n".join(parts)


def test_large_scale_no_numericals_all_attributes_ignored():
    # Large-scale negative test: many attributes but none correspond to numerical roots.
    parts = []
    count = 400  # keep under the 1000-element rule
    for i in range(count):
        parts.append(f"obj_{i}.prop")
    code = "\n".join(parts)


# -------------------------
# Additional targeted cases to increase mutation-resilience
# -------------------------


def test_attribute_on_attribute_where_inner_root_is_numerical():
    # This ensures _get_root_name properly walks down nested Attribute.value chains.
    # Example: num.alias.func -> root is 'num'
    code = "num.alias.somefunc"


def test_attribute_returning_none_for_non_name_roots_prevents_false_positive():
    # For attributes whose value is a Constant or Subscript, _get_root_name should return None.
    # We'll create an attribute value that is a subscript expression: `arr[0].attr`
    code = "arr[0].attr"


def test_mixed_code_with_import_alias_and_direct_name():
    # Mixed scenario: numerical library imported under alias and direct usage by full name.
    code = (
        "import numpy as np\n"
        "x = np.array([1])\n"
        "y = numpy.core.func()"  # direct 'numpy' usage (no import), still an attribute chain
    )


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import ast

from codeflash.code_utils.code_extractor import NumericalUsageChecker

# ============================================================================
# BASIC TEST CASES
# ============================================================================


def test_basic_single_numerical_attribute():
    """Test detection of a simple numerical library attribute access (e.g., np.array)."""
    # Create a checker with 'np' as a numerical library
    checker = NumericalUsageChecker(numerical_names={"np"})

    # Create AST node for np.array
    node = ast.Attribute(value=ast.Name(id="np", ctx=ast.Load()), attr="array", ctx=ast.Load())

    # Visit the attribute node
    checker.visit_Attribute(node)  # 1.24μs -> 1.14μs (8.76% faster)


def test_basic_no_numerical_attribute():
    """Test that non-numerical attributes are not detected."""
    # Create a checker with 'np' as a numerical library
    checker = NumericalUsageChecker(numerical_names={"np"})

    # Create AST node for pd.DataFrame (where pd is not in numerical_names)
    node = ast.Attribute(value=ast.Name(id="pd", ctx=ast.Load()), attr="DataFrame", ctx=ast.Load())

    # Visit the attribute node
    checker.visit_Attribute(node)  # 5.42μs -> 1.07μs (406% faster)


def test_basic_multiple_numerical_libraries():
    """Test detection with multiple numerical libraries registered."""
    # Create a checker with multiple numerical libraries
    checker = NumericalUsageChecker(numerical_names={"np", "tf", "torch"})

    # Test with torch
    node = ast.Attribute(value=ast.Name(id="torch", ctx=ast.Load()), attr="tensor", ctx=ast.Load())
    checker.visit_Attribute(node)  # 1.20μs -> 1.11μs (8.28% faster)


def test_basic_nested_attribute_access():
    """Test detection with nested attribute access (e.g., np.random.rand)."""
    # Create a checker with 'np' as numerical library
    checker = NumericalUsageChecker(numerical_names={"np"})

    # Create AST node for np.random.rand
    # Structure: Attribute(value=Attribute(value=Name('np'), attr='random'), attr='rand')
    inner_attr = ast.Attribute(value=ast.Name(id="np", ctx=ast.Load()), attr="random", ctx=ast.Load())
    node = ast.Attribute(value=inner_attr, attr="rand", ctx=ast.Load())

    # Visit the attribute node
    checker.visit_Attribute(node)  # 1.32μs -> 1.17μs (12.9% faster)


def test_basic_initialization_with_empty_set():
    """Test initialization with an empty set of numerical names."""
    # Create a checker with empty numerical names set
    checker = NumericalUsageChecker(numerical_names=set())

    # Create AST node for np.array
    node = ast.Attribute(value=ast.Name(id="np", ctx=ast.Load()), attr="array", ctx=ast.Load())

    # Visit the attribute node
    checker.visit_Attribute(node)  # 5.04μs -> 1.08μs (366% faster)


def test_basic_exact_match_required():
    """Test that root name matching is exact (case-sensitive)."""
    # Create a checker with 'np' as numerical library
    checker = NumericalUsageChecker(numerical_names={"np"})

    # Create AST node for NP.array (uppercase, should not match)
    node = ast.Attribute(value=ast.Name(id="NP", ctx=ast.Load()), attr="array", ctx=ast.Load())

    # Visit the attribute node
    checker.visit_Attribute(node)  # 4.80μs -> 1.04μs (361% faster)


# ============================================================================
# EDGE TEST CASES
# ============================================================================


def test_edge_early_exit_on_found():
    """Test that once found_numerical is True, further visits return early."""
    # Create a checker with 'np' as numerical library
    checker = NumericalUsageChecker(numerical_names={"np"})

    # Set found_numerical to True
    checker.found_numerical = True

    # Create AST node
    node = ast.Attribute(value=ast.Name(id="np", ctx=ast.Load()), attr="array", ctx=ast.Load())

    # Visit the attribute node
    checker.visit_Attribute(node)  # 391ns -> 381ns (2.62% faster)


def test_edge_deeply_nested_attributes():
    """Test detection with deeply nested attribute chains (e.g., a.b.c.d.e.f)."""
    # Create a checker with 'root' as numerical library
    checker = NumericalUsageChecker(numerical_names={"root"})

    # Build deeply nested structure: root.a.b.c.d.e
    node = ast.Name(id="root", ctx=ast.Load())
    for attr_name in ["a", "b", "c", "d", "e"]:
        node = ast.Attribute(value=node, attr=attr_name, ctx=ast.Load())

    # Visit the outermost attribute node
    checker.visit_Attribute(node)  # 1.69μs -> 1.48μs (14.2% faster)


def test_edge_subscript_expression():
    """Test attribute access on a subscript expression (e.g., array[0].attr)."""
    # Create a checker with 'np' as numerical library
    checker = NumericalUsageChecker(numerical_names={"np"})

    # Create AST node for something[index].method
    # The root is not 'np' directly, but a subscript of 'np'
    subscript = ast.Subscript(value=ast.Name(id="np", ctx=ast.Load()), slice=ast.Constant(value=0), ctx=ast.Load())
    node = ast.Attribute(value=subscript, attr="method", ctx=ast.Load())

    # Visit the attribute node
    checker.visit_Attribute(node)  # 11.1μs -> 11.2μs (1.51% slower)


def test_edge_function_call_expression():
    """Test attribute access on a function call result (e.g., func().method)."""
    # Create a checker with 'np' as numerical library
    checker = NumericalUsageChecker(numerical_names={"np"})

    # Create AST node for func().attr
    call = ast.Call(func=ast.Name(id="func", ctx=ast.Load()), args=[], keywords=[])
    node = ast.Attribute(value=call, attr="method", ctx=ast.Load())

    # Visit the attribute node
    checker.visit_Attribute(node)  # 7.05μs -> 6.67μs (5.71% faster)


def test_edge_binary_operation_expression():
    """Test attribute access on a binary operation (e.g., (a + b).attr)."""
    # Create a checker with 'np' as numerical library
    checker = NumericalUsageChecker(numerical_names={"np"})

    # Create AST node for (a + b).attr
    binop = ast.BinOp(left=ast.Name(id="a", ctx=ast.Load()), op=ast.Add(), right=ast.Name(id="b", ctx=ast.Load()))
    node = ast.Attribute(value=binop, attr="method", ctx=ast.Load())

    # Visit the attribute node
    checker.visit_Attribute(node)  # 7.12μs -> 7.45μs (4.44% slower)


def test_edge_single_character_names():
    """Test with single-character library names."""
    # Create a checker with single-character names
    checker = NumericalUsageChecker(numerical_names={"x", "y", "z"})

    # Create AST node for y.method
    node = ast.Attribute(value=ast.Name(id="y", ctx=ast.Load()), attr="method", ctx=ast.Load())

    # Visit the attribute node
    checker.visit_Attribute(node)  # 1.13μs -> 1.00μs (13.0% faster)


def test_edge_very_long_attribute_name():
    """Test with very long attribute names."""
    # Create a checker
    checker = NumericalUsageChecker(numerical_names={"np"})

    # Create AST node with a very long attribute name
    long_attr = "a" * 500
    node = ast.Attribute(value=ast.Name(id="np", ctx=ast.Load()), attr=long_attr, ctx=ast.Load())

    # Visit the attribute node
    checker.visit_Attribute(node)  # 1.02μs -> 1.04μs (1.92% slower)


def test_edge_special_characters_in_library_name():
    """Test with special characters in numerical library names."""
    # Create a checker with underscore in name
    checker = NumericalUsageChecker(numerical_names={"_np", "np_", "__np__"})

    # Test with _np
    node = ast.Attribute(value=ast.Name(id="_np", ctx=ast.Load()), attr="array", ctx=ast.Load())
    checker.visit_Attribute(node)  # 1.10μs -> 1.03μs (6.78% faster)


def test_edge_none_root_name():
    """Test when _get_root_name returns None for certain node types."""
    # Create a checker
    checker = NumericalUsageChecker(numerical_names={"np"})

    # Create an attribute node with a non-standard value type
    # Using a Constant node which should return None from _get_root_name
    node = ast.Attribute(value=ast.Constant(value=42), attr="method", ctx=ast.Load())

    # Visit the attribute node (should not crash)
    checker.visit_Attribute(node)  # 8.20μs -> 8.12μs (0.862% faster)


def test_edge_attribute_on_unpack_expression():
    """Test attribute access on unpacking expressions."""
    # Create a checker
    checker = NumericalUsageChecker(numerical_names={"np"})

    # Create an attribute on a starred expression
    starred = ast.Starred(value=ast.Name(id="np", ctx=ast.Load()), ctx=ast.Load())
    node = ast.Attribute(value=starred, attr="method", ctx=ast.Load())

    # Visit the attribute node
    checker.visit_Attribute(node)  # 6.25μs -> 6.36μs (1.73% slower)


# ============================================================================
# LARGE SCALE TEST CASES
# ============================================================================


def test_large_scale_many_numerical_libraries():
    """Test with a large set of numerical library names."""
    # Create a checker with 500 different library names
    numerical_libs = {f"lib_{i}" for i in range(500)}
    checker = NumericalUsageChecker(numerical_names=numerical_libs)

    # Test with a library from the middle of the set
    node = ast.Attribute(value=ast.Name(id="lib_250", ctx=ast.Load()), attr="method", ctx=ast.Load())
    checker.visit_Attribute(node)  # 1.25μs -> 1.21μs (3.38% faster)


def test_large_scale_library_not_in_large_set():
    """Test that non-existent library is not detected in large set."""
    # Create a checker with 500 library names
    numerical_libs = {f"lib_{i}" for i in range(500)}
    checker = NumericalUsageChecker(numerical_names=numerical_libs)

    # Test with a library not in the set
    node = ast.Attribute(value=ast.Name(id="lib_1000", ctx=ast.Load()), attr="method", ctx=ast.Load())
    checker.visit_Attribute(node)  # 4.79μs -> 1.10μs (335% faster)


def test_large_scale_sequential_visits():
    """Test many sequential attribute visits with detection on specific visit."""
    # Create a checker with 'target' as numerical library
    checker = NumericalUsageChecker(numerical_names={"target"})

    # Visit 200 non-matching attributes
    for i in range(200):
        node = ast.Attribute(value=ast.Name(id=f"lib_{i}", ctx=ast.Load()), attr="method", ctx=ast.Load())
        checker.visit_Attribute(node)  # 379μs -> 88.5μs (329% faster)

    # Now visit a matching attribute
    node = ast.Attribute(value=ast.Name(id="target", ctx=ast.Load()), attr="method", ctx=ast.Load())
    checker.visit_Attribute(node)  # 531ns -> 501ns (5.99% faster)

    # Visit more attributes after detection (should return early)
    for i in range(200, 250):
        node = ast.Attribute(value=ast.Name(id=f"lib_{i}", ctx=ast.Load()), attr="method", ctx=ast.Load())
        checker.visit_Attribute(node)  # 7.39μs -> 7.07μs (4.57% faster)


def test_large_scale_deeply_nested_chain():
    """Test with a very long chain of nested attributes."""
    # Create a checker with 'root' as numerical library
    checker = NumericalUsageChecker(numerical_names={"root"})

    # Build a chain of 100 nested attributes
    node = ast.Name(id="root", ctx=ast.Load())
    for i in range(100):
        node = ast.Attribute(value=node, attr=f"attr_{i}", ctx=ast.Load())

    # Visit the outermost attribute
    checker.visit_Attribute(node)  # 28.3μs -> 10.4μs (172% faster)


def test_large_scale_mixed_matching_libraries():
    """Test with alternating matching and non-matching libraries in large set."""
    # Create a checker with every other library as numerical
    matching_libs = {f"lib_{i}" for i in range(0, 400, 2)}
    checker = NumericalUsageChecker(numerical_names=matching_libs)

    # Visit libraries in sequence, checking detection pattern
    for i in range(200):
        node = ast.Attribute(value=ast.Name(id=f"lib_{i}", ctx=ast.Load()), attr="method", ctx=ast.Load())
        checker.visit_Attribute(node)  # 1.32μs -> 1.17μs (12.8% faster)

        # Check if this should be detected
        if i % 2 == 0 and i < 400:
            # Once detected, all subsequent visits should keep it True
            if checker.found_numerical:
                break


def test_large_scale_performance_no_detection():
    """Test performance with many visits when no library matches."""
    # Create a checker with specific libraries
    checker = NumericalUsageChecker(numerical_names={"numpy", "pandas", "scipy"})

    # Visit 300 non-matching attributes
    for i in range(300):
        node = ast.Attribute(value=ast.Name(id=f"module_{i}", ctx=ast.Load()), attr="function", ctx=ast.Load())
        checker.visit_Attribute(node)  # 582μs -> 134μs (335% faster)


def test_large_scale_complex_nested_with_many_branches():
    """Test deeply nested attributes with variety in depth."""
    # Create a checker with multiple numerical libraries
    checker = NumericalUsageChecker(numerical_names={"np", "tf", "torch"})

    # Create various nested chains with different depths
    test_cases = []

    # Case 1: Single level - np.array
    test_cases.append(ast.Attribute(value=ast.Name(id="np", ctx=ast.Load()), attr="array", ctx=ast.Load()))

    # Case 2: Two levels - np.random.rand
    inner = ast.Attribute(value=ast.Name(id="np", ctx=ast.Load()), attr="random", ctx=ast.Load())
    test_cases.append(ast.Attribute(value=inner, attr="rand", ctx=ast.Load()))

    # Case 3: Three levels
    inner = ast.Attribute(value=ast.Name(id="tf", ctx=ast.Load()), attr="nn", ctx=ast.Load())
    inner = ast.Attribute(value=inner, attr="relu", ctx=ast.Load())
    test_cases.append(inner)

    # Visit first test case
    checker.visit_Attribute(test_cases[0])  # 1.15μs -> 1.04μs (10.6% faster)


def test_large_scale_many_checkers_independent():
    """Test that multiple independent checkers work correctly."""
    # Create 100 different checkers with different configurations
    checkers = []
    for i in range(100):
        lib_name = f"lib_{i}"
        checker = NumericalUsageChecker(numerical_names={lib_name})
        checkers.append((checker, lib_name))

    # Test each checker
    for checker, lib_name in checkers:
        node = ast.Attribute(value=ast.Name(id=lib_name, ctx=ast.Load()), attr="method", ctx=ast.Load())
        checker.visit_Attribute(node)  # 45.2μs -> 43.7μs (3.30% faster)

    # Verify other checkers were not affected
    for checker, lib_name in checkers[:10]:
        pass


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-pr1051-2026-01-14T02.21.37 and push.

[codeflash-ai]

This PR has been automatically closed because the original PR #1051 by aseembits93 was closed.