⚡️ Speed up method CfoVisitor._get_called_func_name by 18% in PR #501 (runtime-fixes-2)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #501

If you approve this dependent PR, these changes will be merged into the original PR branch runtime-fixes-2.

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

---

📄 18% (0.18x) speedup for CfoVisitor._get_called_func_name in codeflash/code_utils/edit_generated_tests.py

⏱️ Runtime : 1.48 milliseconds → 1.26 milliseconds (best of 142 runs)

📝 Explanation and details

Key optimizations:

• Avoid multiple isinstance calls by caching type(node) in a variable and using identity checks against the expected types (_ast_Name, _ast_Attribute).
• This speeds up the type checks, since identity comparison is faster than repeated isinstance checks, and in this hot path, there are only two accepted types.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

import ast

# imports
import pytest  # used for our unit tests
from codeflash.code_utils.edit_generated_tests import CfoVisitor

# unit tests

# Helper function to parse an expression and return its AST node
def parse_expr(expr: str):
    """Parse a Python expression and return its AST node."""
    return ast.parse(expr, mode='eval').body

# Fixture for a CfoVisitor instance (source code and function_name are irrelevant for _get_called_func_name)
@pytest.fixture
def visitor():
    return CfoVisitor("dummy_func", "")

# ---------------- BASIC TEST CASES ----------------

def test_get_called_func_name_simple_name(visitor):
    # Test with a simple function call: foo()
    node = parse_expr("foo")
    codeflash_output = visitor._get_called_func_name(node) # 552ns -> 501ns (10.2% faster)

def test_get_called_func_name_simple_attribute(visitor):
    # Test with an attribute access: obj.method
    node = parse_expr("obj.method")
    codeflash_output = visitor._get_called_func_name(node) # 652ns -> 541ns (20.5% faster)

def test_get_called_func_name_nested_attribute(visitor):
    # Test with nested attributes: obj.foo.bar
    node = parse_expr("obj.foo.bar")
    # Should return the last attribute's name
    codeflash_output = visitor._get_called_func_name(node) # 651ns -> 521ns (25.0% faster)

def test_get_called_func_name_name_in_call(visitor):
    # Test with a function call node: foo()
    node = parse_expr("foo()").func
    codeflash_output = visitor._get_called_func_name(node) # 511ns -> 501ns (2.00% faster)

def test_get_called_func_name_attribute_in_call(visitor):
    # Test with a method call: obj.method()
    node = parse_expr("obj.method()").func
    codeflash_output = visitor._get_called_func_name(node) # 682ns -> 530ns (28.7% faster)

# ---------------- EDGE TEST CASES ----------------

def test_get_called_func_name_none(visitor):
    # Test with None as input
    codeflash_output = visitor._get_called_func_name(None) # 631ns -> 461ns (36.9% faster)

def test_get_called_func_name_constant(visitor):
    # Test with a constant (should not have a function name)
    node = parse_expr("42")
    codeflash_output = visitor._get_called_func_name(node) # 691ns -> 471ns (46.7% faster)

def test_get_called_func_name_tuple(visitor):
    # Test with a tuple node
    node = parse_expr("(foo, bar)")
    codeflash_output = visitor._get_called_func_name(node) # 962ns -> 481ns (100% faster)

def test_get_called_func_name_lambda(visitor):
    # Test with a lambda node
    node = parse_expr("lambda x: x + 1")
    codeflash_output = visitor._get_called_func_name(node) # 641ns -> 491ns (30.5% faster)

def test_get_called_func_name_subscript(visitor):
    # Test with a subscript node: arr[0]
    node = parse_expr("arr[0]")
    codeflash_output = visitor._get_called_func_name(node) # 601ns -> 490ns (22.7% faster)

def test_get_called_func_name_call_of_attribute(visitor):
    # Test with a call to a method: obj.method().attr
    node = parse_expr("obj.method().attr")
    # Should return 'attr'
    codeflash_output = visitor._get_called_func_name(node) # 682ns -> 571ns (19.4% faster)

def test_get_called_func_name_call_of_call(visitor):
    # Test with a call to a function returning a function: foo()()
    node = parse_expr("foo()()").func
    # This is a Call node, not Name or Attribute, so should return None
    codeflash_output = visitor._get_called_func_name(node) # 611ns -> 420ns (45.5% faster)

def test_get_called_func_name_attribute_of_name(visitor):
    # Test with attribute of a name: foo.bar
    node = parse_expr("foo.bar")
    codeflash_output = visitor._get_called_func_name(node) # 651ns -> 511ns (27.4% faster)

def test_get_called_func_name_attribute_of_attribute(visitor):
    # Test with attribute of attribute: foo.bar.baz
    node = parse_expr("foo.bar.baz")
    codeflash_output = visitor._get_called_func_name(node) # 661ns -> 481ns (37.4% faster)

def test_get_called_func_name_invalid_node(visitor):
    # Test with an invalid AST node type (e.g., ast.Module)
    node = ast.Module(body=[])
    codeflash_output = visitor._get_called_func_name(node) # 581ns -> 441ns (31.7% faster)

# ---------------- LARGE SCALE TEST CASES ----------------

def test_get_called_func_name_many_names(visitor):
    # Test with a large number of Name nodes
    for i in range(1000):
        node = ast.Name(id=f"func_{i}", ctx=ast.Load())
        codeflash_output = visitor._get_called_func_name(node)

def test_get_called_func_name_many_attributes(visitor):
    # Test with a large number of Attribute nodes
    for i in range(1000):
        node = ast.Attribute(value=ast.Name(id="obj", ctx=ast.Load()), attr=f"method_{i}", ctx=ast.Load())
        codeflash_output = visitor._get_called_func_name(node)

def test_get_called_func_name_mixed_large(visitor):
    # Test a mix of Name, Attribute, and other node types in a large list
    for i in range(333):
        node = ast.Name(id=f"func_{i}", ctx=ast.Load())
        codeflash_output = visitor._get_called_func_name(node) # 230ns -> 180ns (27.8% faster)
    for i in range(333):
        node = ast.Attribute(value=ast.Name(id="obj", ctx=ast.Load()), attr=f"method_{i}", ctx=ast.Load())
        codeflash_output = visitor._get_called_func_name(node)
    for i in range(334):
        node = ast.Constant(value=i)
        codeflash_output = visitor._get_called_func_name(node)

def test_get_called_func_name_performance(visitor):
    # Performance test: ensure function is fast for 1000 Name nodes
    import time
    nodes = [ast.Name(id=f"func_{i}", ctx=ast.Load()) for i in range(1000)]
    start = time.time()
    for node in nodes:
        codeflash_output = visitor._get_called_func_name(node)
    elapsed = time.time() - start

# ---------------- ADDITIONAL EDGE CASES ----------------

def test_get_called_func_name_keyword(visitor):
    # Test with a keyword node (should return None)
    node = ast.keyword(arg="foo", value=ast.Constant(value=1))
    codeflash_output = visitor._get_called_func_name(node) # 661ns -> 561ns (17.8% faster)

def test_get_called_func_name_starred(visitor):
    # Test with a starred node (should return None)
    node = ast.Starred(value=ast.Name(id="foo", ctx=ast.Load()), ctx=ast.Load())
    codeflash_output = visitor._get_called_func_name(node) # 622ns -> 450ns (38.2% faster)

def test_get_called_func_name_comprehension(visitor):
    # Test with a comprehension node (should return None)
    node = parse_expr("[x for x in y]")
    codeflash_output = visitor._get_called_func_name(node) # 922ns -> 452ns (104% faster)

def test_get_called_func_name_binop(visitor):
    # Test with a binary operation node (should return None)
    node = parse_expr("a + b")
    codeflash_output = visitor._get_called_func_name(node) # 611ns -> 471ns (29.7% faster)

def test_get_called_func_name_call(visitor):
    # Test with a Call node itself (should return None)
    node = parse_expr("foo()")
    codeflash_output = visitor._get_called_func_name(node) # 641ns -> 460ns (39.3% faster)
# 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

import ast

# imports
import pytest  # used for our unit tests
from codeflash.code_utils.edit_generated_tests import CfoVisitor

# unit tests

# Helper to parse code and get the node for _get_called_func_name
def get_call_node(expr_code):
    """Helper to parse an expression and return the AST node representing the function being called."""
    # Parse as an expression
    expr_ast = ast.parse(expr_code, mode='eval')
    # The function being called is expr_ast.body.func for a Call node
    if isinstance(expr_ast.body, ast.Call):
        return expr_ast.body.func
    # If not a call, just return the body itself
    return expr_ast.body

# Helper to create a visitor instance (source_code is irrelevant for this test)
def get_visitor():
    return CfoVisitor(function_name="dummy", source_code="")

# 1. Basic Test Cases

def test_get_called_func_name_simple_name():
    # Test with a simple function name: foo()
    node = get_call_node("foo()")
    visitor = get_visitor()
    codeflash_output = visitor._get_called_func_name(node); result = codeflash_output # 551ns -> 541ns (1.85% faster)

def test_get_called_func_name_simple_attribute():
    # Test with a simple attribute call: obj.bar()
    node = get_call_node("obj.bar()")
    visitor = get_visitor()
    codeflash_output = visitor._get_called_func_name(node); result = codeflash_output # 651ns -> 561ns (16.0% faster)

def test_get_called_func_name_nested_attribute():
    # Test with nested attributes: obj.foo.bar()
    node = get_call_node("obj.foo.bar()")
    visitor = get_visitor()
    codeflash_output = visitor._get_called_func_name(node); result = codeflash_output # 611ns -> 521ns (17.3% faster)

def test_get_called_func_name_builtin():
    # Test with a builtin function: len()
    node = get_call_node("len()")
    visitor = get_visitor()
    codeflash_output = visitor._get_called_func_name(node); result = codeflash_output # 490ns -> 492ns (0.407% slower)

def test_get_called_func_name_module_function():
    # Test with a module function: math.sqrt()
    node = get_call_node("math.sqrt()")
    visitor = get_visitor()
    codeflash_output = visitor._get_called_func_name(node); result = codeflash_output # 681ns -> 521ns (30.7% faster)

# 2. Edge Test Cases

def test_get_called_func_name_lambda():
    # Test with a lambda being called: (lambda x: x+1)(5)
    node = get_call_node("(lambda x: x+1)(5)")
    visitor = get_visitor()
    codeflash_output = visitor._get_called_func_name(node); result = codeflash_output # 631ns -> 461ns (36.9% faster)

def test_get_called_func_name_call_of_call():
    # Test with calling the result of another call: foo()()
    node = get_call_node("foo()()")
    visitor = get_visitor()
    codeflash_output = visitor._get_called_func_name(node); result = codeflash_output # 591ns -> 451ns (31.0% faster)

def test_get_called_func_name_subscript():
    # Test with calling a subscript: arr[0]()
    node = get_call_node("arr[0]()")
    visitor = get_visitor()
    codeflash_output = visitor._get_called_func_name(node); result = codeflash_output # 581ns -> 451ns (28.8% faster)

def test_get_called_func_name_attribute_of_call():
    # Test with calling an attribute of a call: foo().bar()
    node = get_call_node("foo().bar()")
    visitor = get_visitor()
    codeflash_output = visitor._get_called_func_name(node); result = codeflash_output # 601ns -> 501ns (20.0% faster)

def test_get_called_func_name_name_node():
    # Test with just a Name node, not a call
    node = get_call_node("foo")
    visitor = get_visitor()
    codeflash_output = visitor._get_called_func_name(node); result = codeflash_output # 521ns -> 511ns (1.96% faster)

def test_get_called_func_name_attribute_node():
    # Test with just an Attribute node, not a call
    node = get_call_node("obj.method")
    visitor = get_visitor()
    codeflash_output = visitor._get_called_func_name(node); result = codeflash_output # 581ns -> 511ns (13.7% faster)

def test_get_called_func_name_constant():
    # Test with a constant value (should return None)
    node = get_call_node("42")
    visitor = get_visitor()
    codeflash_output = visitor._get_called_func_name(node); result = codeflash_output # 551ns -> 451ns (22.2% faster)

def test_get_called_func_name_none():
    # Test with None as input
    visitor = get_visitor()
    codeflash_output = visitor._get_called_func_name(None); result = codeflash_output # 601ns -> 420ns (43.1% faster)

def test_get_called_func_name_unusual_node_type():
    # Test with an AST node type not handled (e.g., BinOp)
    node = get_call_node("1 + 2")
    visitor = get_visitor()
    codeflash_output = visitor._get_called_func_name(node); result = codeflash_output # 541ns -> 471ns (14.9% faster)

# 3. Large Scale Test Cases

def test_get_called_func_name_many_names():
    # Test with many different Name nodes
    visitor = get_visitor()
    for i in range(1000):
        node = ast.Name(id=f"func_{i}", ctx=ast.Load())
        codeflash_output = visitor._get_called_func_name(node); result = codeflash_output

def test_get_called_func_name_many_attributes():
    # Test with many different Attribute nodes
    visitor = get_visitor()
    for i in range(1000):
        node = ast.Attribute(value=ast.Name(id="obj", ctx=ast.Load()), attr=f"method_{i}", ctx=ast.Load())
        codeflash_output = visitor._get_called_func_name(node); result = codeflash_output

def test_get_called_func_name_mixed_types():
    # Test with a mix of valid and invalid node types
    visitor = get_visitor()
    nodes = []
    # 500 Name nodes
    for i in range(500):
        nodes.append(ast.Name(id=f"n_{i}", ctx=ast.Load()))
    # 250 Attribute nodes
    for i in range(250):
        nodes.append(ast.Attribute(value=ast.Name(id="o", ctx=ast.Load()), attr=f"a_{i}", ctx=ast.Load()))
    # 125 Constant nodes
    for i in range(125):
        nodes.append(ast.Constant(value=i))
    # 125 BinOp nodes
    for i in range(125):
        nodes.append(ast.BinOp(left=ast.Constant(value=i), op=ast.Add(), right=ast.Constant(value=i+1)))
    # Now test
    for node in nodes[:500]:
        pass
    for node in nodes[500:750]:
        pass
    for node in nodes[750:875]:
        codeflash_output = visitor._get_called_func_name(node)
    for node in nodes[875:]:
        codeflash_output = visitor._get_called_func_name(node)

def test_get_called_func_name_long_attribute_chain():
    # Test with a very long attribute chain: a.b.c.d.e.f.g.h.i.j()
    chain = "a" + "".join(f".b{i}" for i in range(10)) + "()"
    node = get_call_node(chain)
    visitor = get_visitor()
    codeflash_output = visitor._get_called_func_name(node); result = codeflash_output # 672ns -> 590ns (13.9% faster)

def test_get_called_func_name_large_call_tree():
    # Test with a large number of nested calls: f(g(h(i(j(k(l(m(n(o(p()))))))))))
    expr = "p()"
    for fn in reversed("abcdefghijklmno"):
        expr = f"{fn}({expr})"
    node = get_call_node(expr)
    visitor = get_visitor()
    codeflash_output = visitor._get_called_func_name(node); result = codeflash_output
# 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-pr501-2025-07-03T16.17.03 and push.