⚡️ Speed up function `function_kind` by 6% by codeflash-ai[bot] · Pull Request #319 · codeflash-ai/codeflash

codeflash-ai · 2025-06-12T00:18:15Z

📄 6% (0.06x) speedup for `function_kind` in `codeflash/code_utils/static_analysis.py`

⏱️ Runtime : 494 microseconds → 468 microseconds (best of 69 runs)

📝 Explanation and details

Here’s an optimized rewrite focusing on.

Avoiding repeated list/dict construction: Replaces ["FunctionDef", "AsyncFunctionDef"] with a set and a local variable for quick lookup.
Decorator matching: Replaces the for-loop with a single pass using generator expression – thus returning immediately when a match is found, and avoiding extra checks.
Parent type check: Caches parents[0].type in a variable.
Reduces function calls: Pulls common checks out of loops.

Here's the optimized code (existing comments are unchanged, since logic is nearly the same).

Summary of improvements.

Eliminate duplicated string list creation per call.
Save parents[0].type only once
Early returns in the decorator loop; avoids further iteration if not necessary.

This should speed up your function, especially under heavy use!

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

from __future__ import annotations

import ast
from enum import Enum
from typing import List

# imports
import pytest  # used for our unit tests
from codeflash.code_utils.static_analysis import function_kind


class FunctionKind(Enum):
    FUNCTION = "function"
    INSTANCE_METHOD = "instance_method"
    CLASS_METHOD = "class_method"
    STATIC_METHOD = "static_method"

class FunctionParent:
    def __init__(self, type_: str):
        self.type = type_
from codeflash.code_utils.static_analysis import function_kind


# Helper to parse function source code and return ast.FunctionDef or ast.AsyncFunctionDef
def get_func_node(src: str):
    module = ast.parse(src)
    for node in ast.walk(module):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            return node
    raise ValueError("No function node found in source.")

# --------------------
# Unit Tests
# --------------------

# 1. Basic Test Cases

def test_function_at_module_level():
    # Basic: function at module level (no parents)
    node = get_func_node("def foo(): pass")
    parents = []
    codeflash_output = function_kind(node, parents) # 458ns -> 500ns

def test_function_with_function_parent():
    # Basic: function nested inside another function
    node = get_func_node("def bar():\n    def inner(): pass")
    # simulate being inside a function
    parents = [FunctionParent("FunctionDef")]
    codeflash_output = function_kind(node.body[0], parents) # 583ns -> 666ns

def test_function_with_asyncfunction_parent():
    # Basic: function nested inside async function
    node = get_func_node("async def bar():\n    def inner(): pass")
    parents = [FunctionParent("AsyncFunctionDef")]
    codeflash_output = function_kind(node.body[0], parents) # 583ns -> 583ns







def test_no_parents_none():
    # Edge: parents is None (should treat as no parents)
    node = get_func_node("def foo(): pass")
    codeflash_output = function_kind(node, None) # 542ns -> 833ns






def test_classdef_parent_not_first():
    # Edge: multiple parents, ClassDef not first
    node = get_func_node("def foo(): pass")
    parents = [FunctionParent("FunctionDef"), FunctionParent("ClassDef")]
    # Only first parent matters, so should be FUNCTION
    codeflash_output = function_kind(node, parents) # 708ns -> 1.25μs

def test_unknown_parent_type():
    # Edge: unknown parent type, should return None
    node = get_func_node("def foo(): pass")
    parents = [FunctionParent("Module")]
    codeflash_output = function_kind(node, parents) # 500ns -> 500ns



def test_async_function_at_module_level():
    # Edge: async function at module level
    node = get_func_node("async def foo(): pass")
    parents = []
    codeflash_output = function_kind(node, parents) # 500ns -> 667ns

def test_function_with_weird_parent():
    # Edge: parent type is an unexpected string
    node = get_func_node("def foo(): pass")
    parents = [FunctionParent("WeirdParent")]
    codeflash_output = function_kind(node, parents) # 500ns -> 625ns

def test_function_with_non_list_parents():
    # Edge: parents is not a list (e.g., tuple)
    node = get_func_node("def foo(): pass")
    parents = tuple()
    codeflash_output = function_kind(node, parents) # 500ns -> 458ns


def test_many_methods_in_class():
    # Large scale: class with many methods, all instance methods
    src = "class C:\n" + "\n".join([f"    def foo{i}(self): pass" for i in range(500)])
    class_node = ast.parse(src).body[0]
    parents = [FunctionParent("ClassDef")]
    for func_node in class_node.body:
        codeflash_output = function_kind(func_node, parents) # 208ns -> 208ns


def test_large_nested_functions():
    # Large scale: deeply nested functions (up to 10 levels)
    src = ""
    indent = ""
    for i in range(10):
        src += f"{indent}def f{i}():\n"
        indent += "    "
    src += f"{indent}pass"
    node = ast.parse(src)
    # Go down to the innermost function
    func_node = node.body[0]
    parents = [FunctionParent("FunctionDef")] * 10
    # The innermost function node is at depth 10
    for _ in range(9):
        func_node = func_node.body[0]
    codeflash_output = function_kind(func_node, parents) # 750ns -> 875ns

def test_large_number_of_parents():
    # Large scale: parents list with many elements, first is ClassDef
    node = get_func_node("def foo(self): pass")
    parents = [FunctionParent("ClassDef")] + [FunctionParent("FunctionDef")] * 999
    codeflash_output = function_kind(node, parents) # 709ns -> 708ns

def test_large_number_of_module_level_functions():
    # Large scale: many module-level functions
    src = "\n".join([f"def foo{i}(): pass" for i in range(1000)])
    module = ast.parse(src)
    parents = []
    for func_node in module.body:
        codeflash_output = function_kind(func_node, parents) # 83ns -> 83ns
# 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
from enum import Enum
from typing import Any

# imports
import pytest  # used for our unit tests
from codeflash.code_utils.static_analysis import function_kind

# --- Mocking FunctionKind and FunctionParent for testability ---


class FunctionKind(Enum):
    FUNCTION = "function"
    CLASS_METHOD = "class_method"
    STATIC_METHOD = "static_method"
    INSTANCE_METHOD = "instance_method"

class FunctionParent:
    def __init__(self, type_: str):
        self.type = type_
from codeflash.code_utils.static_analysis import function_kind

# --- Unit tests ---

# Helper to create a FunctionDef or AsyncFunctionDef node with decorators
def make_func_node(
    name: str = "f",
    decorators: list[Any] = None,
    async_: bool = False
) -> ast.FunctionDef | ast.AsyncFunctionDef:
    if decorators is None:
        decorators = []
    args = ast.arguments(
        posonlyargs=[],
        args=[],
        kwonlyargs=[],
        kw_defaults=[],
        defaults=[]
    )
    if async_:
        return ast.AsyncFunctionDef(
            name=name,
            args=args,
            body=[],
            decorator_list=decorators,
            returns=None,
            type_comment=None
        )
    else:
        return ast.FunctionDef(
            name=name,
            args=args,
            body=[],
            decorator_list=decorators,
            returns=None,
            type_comment=None
        )

# ------------------------
# 1. Basic Test Cases
# ------------------------

def test_top_level_function_returns_functionkind_function():
    # Function at module level, no parents
    node = make_func_node()
    parents = []
    codeflash_output = function_kind(node, parents); result = codeflash_output # 417ns -> 458ns

def test_function_inside_function_returns_functionkind_function():
    # Function nested inside another function
    node = make_func_node()
    parents = [FunctionParent("FunctionDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 625ns -> 709ns

def test_async_function_top_level():
    node = make_func_node(async_=True)
    parents = []
    codeflash_output = function_kind(node, parents); result = codeflash_output # 417ns -> 417ns

def test_class_instance_method():
    # Function inside a class, no decorators
    node = make_func_node()
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 708ns -> 750ns

def test_class_classmethod():
    # Function inside a class, with @classmethod
    decorator = ast.Name(id="classmethod", ctx=ast.Load())
    node = make_func_node(decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 875ns -> 1.04μs

def test_class_staticmethod():
    # Function inside a class, with @staticmethod
    decorator = ast.Name(id="staticmethod", ctx=ast.Load())
    node = make_func_node(decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 834ns -> 875ns

def test_class_multiple_decorators_staticmethod_last():
    # Function inside a class, multiple decorators, staticmethod last
    decorator1 = ast.Name(id="something", ctx=ast.Load())
    decorator2 = ast.Name(id="staticmethod", ctx=ast.Load())
    node = make_func_node(decorators=[decorator1, decorator2])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 1.04μs -> 1.08μs

def test_class_multiple_decorators_classmethod_first():
    # Function inside a class, multiple decorators, classmethod first
    decorator1 = ast.Name(id="classmethod", ctx=ast.Load())
    decorator2 = ast.Name(id="other", ctx=ast.Load())
    node = make_func_node(decorators=[decorator1, decorator2])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 791ns -> 833ns

# ------------------------
# 2. Edge Test Cases
# ------------------------

def test_empty_parents_returns_function():
    # No parents, should be FUNCTION
    node = make_func_node()
    codeflash_output = function_kind(node, []); result = codeflash_output # 417ns -> 459ns

def test_parent_is_unknown_type_returns_none():
    # Parent is not ClassDef or FunctionDef/AsyncFunctionDef
    node = make_func_node()
    parents = [FunctionParent("Module")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 458ns -> 542ns

def test_decorator_is_not_ast_name():
    # Decorator is not ast.Name (e.g., ast.Call for @decorator())
    decorator = ast.Call(
        func=ast.Name(id="classmethod", ctx=ast.Load()),
        args=[],
        keywords=[]
    )
    node = make_func_node(decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 875ns -> 958ns

def test_decorator_is_ast_attribute():
    # Decorator is ast.Attribute (e.g., @abc.abstractmethod)
    decorator = ast.Attribute(
        value=ast.Name(id="abc", ctx=ast.Load()),
        attr="abstractmethod",
        ctx=ast.Load()
    )
    node = make_func_node(decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 750ns -> 792ns

def test_class_with_both_class_and_static_decorators():
    # Both classmethod and staticmethod present, classmethod first
    decorator1 = ast.Name(id="classmethod", ctx=ast.Load())
    decorator2 = ast.Name(id="staticmethod", ctx=ast.Load())
    node = make_func_node(decorators=[decorator1, decorator2])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 792ns -> 833ns

def test_class_with_both_static_and_class_decorators():
    # Both staticmethod and classmethod, staticmethod first
    decorator1 = ast.Name(id="staticmethod", ctx=ast.Load())
    decorator2 = ast.Name(id="classmethod", ctx=ast.Load())
    node = make_func_node(decorators=[decorator1, decorator2])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 875ns -> 834ns

def test_class_with_nonstring_decorator_id():
    # Decorator with id that is not a recognized string
    decorator = ast.Name(id="otherdecorator", ctx=ast.Load())
    node = make_func_node(decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 833ns -> 875ns

def test_async_function_inside_class():
    # Async function inside class, no decorators
    node = make_func_node(async_=True)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 667ns -> 666ns

def test_function_with_multiple_parents():
    # Multiple parents, only first is checked
    node = make_func_node()
    parents = [FunctionParent("ClassDef"), FunctionParent("Module")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 625ns -> 667ns

def test_function_with_parent_asyncfunctiondef():
    # Parent is AsyncFunctionDef
    node = make_func_node()
    parents = [FunctionParent("AsyncFunctionDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 625ns -> 583ns

def test_function_with_parent_none_type():
    # Parent with type None
    class DummyParent:
        def __init__(self):
            self.type = None
    node = make_func_node()
    parents = [DummyParent()]
    codeflash_output = function_kind(node, parents); result = codeflash_output

# ------------------------
# 3. Large Scale Test Cases
# ------------------------

def test_many_functions_in_module():
    # Test many top-level functions
    for i in range(100):
        node = make_func_node(name=f"func_{i}")
        parents = []
        codeflash_output = function_kind(node, parents); result = codeflash_output # 125ns -> 83ns

def test_many_class_methods():
    # Test many @classmethods in a class
    decorator = ast.Name(id="classmethod", ctx=ast.Load())
    for i in range(100):
        node = make_func_node(name=f"class_func_{i}", decorators=[decorator])
        parents = [FunctionParent("ClassDef")]
        codeflash_output = function_kind(node, parents); result = codeflash_output # 291ns -> 250ns

def test_many_static_methods():
    # Test many @staticmethods in a class
    decorator = ast.Name(id="staticmethod", ctx=ast.Load())
    for i in range(100):
        node = make_func_node(name=f"static_func_{i}", decorators=[decorator])
        parents = [FunctionParent("ClassDef")]
        codeflash_output = function_kind(node, parents); result = codeflash_output # 291ns -> 250ns

def test_many_instance_methods_with_varied_decorators():
    # Test many instance methods with non-matching decorators
    for i in range(100):
        decorator = ast.Name(id=f"decorator_{i}", ctx=ast.Load())
        node = make_func_node(name=f"instance_func_{i}", decorators=[decorator])
        parents = [FunctionParent("ClassDef")]
        codeflash_output = function_kind(node, parents); result = codeflash_output # 291ns -> 250ns

def test_large_decorator_list_with_classmethod_at_end():
    # Function with 999 unrelated decorators and classmethod at the end
    decorators = [ast.Name(id=f"decorator_{i}", ctx=ast.Load()) for i in range(999)]
    decorators.append(ast.Name(id="classmethod", ctx=ast.Load()))
    node = make_func_node(decorators=decorators)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 59.8μs -> 52.8μs

def test_large_decorator_list_with_staticmethod_at_start():
    # Function with staticmethod as first of 1000 decorators
    decorators = [ast.Name(id="staticmethod", ctx=ast.Load())] + [
        ast.Name(id=f"decorator_{i}", ctx=ast.Load()) for i in range(999)
    ]
    node = make_func_node(decorators=decorators)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 875ns -> 875ns

def test_large_number_of_parents():
    # Only first parent should matter, even if there are many
    node = make_func_node()
    parents = [FunctionParent("ClassDef")] + [FunctionParent("Module") for _ in range(999)]
    codeflash_output = function_kind(node, parents); result = codeflash_output # 709ns -> 708ns

def test_large_number_of_functions_varied_contexts():
    # Mix of top-level, nested, and class methods
    for i in range(50):
        node = make_func_node(name=f"func_{i}")
        codeflash_output = function_kind(node, []) # 125ns -> 125ns
        codeflash_output = function_kind(node, [FunctionParent("FunctionDef")]) # 125ns -> 125ns
        codeflash_output = function_kind(node, [FunctionParent("ClassDef")]) # 125ns -> 125ns
        decorator = ast.Name(id="classmethod", ctx=ast.Load())
        node2 = make_func_node(name=f"class_func_{i}", decorators=[decorator])
        codeflash_output = function_kind(node2, [FunctionParent("ClassDef")]) # 125ns -> 125ns
# 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-function_kind-mbsmr7uy and push.

Here’s an optimized rewrite focusing on. - **Avoiding repeated list/dict construction:** Replaces `["FunctionDef", "AsyncFunctionDef"]` with a `set` and a local variable for quick lookup. - **Decorator matching:** Replaces the for-loop with a single pass using generator expression – thus returning immediately when a match is found, and avoiding extra checks. - **Parent type check:** Caches `parents[0].type` in a variable. - **Reduces function calls:** Pulls common checks out of loops. Here's the optimized code (existing comments are unchanged, since logic is nearly the same). ### Summary of improvements. - Eliminate duplicated string list creation per call. - Save parents[0].type only once - Early returns in the decorator loop; avoids further iteration if not necessary. This should speed up your function, especially under heavy use!

codeflash-ai Bot added the ⚡️ codeflash Optimization PR opened by Codeflash AI label Jun 12, 2025

codeflash-ai Bot requested a review from KRRT7 June 12, 2025 00:18

KRRT7 closed this Jun 12, 2025

codeflash-ai Bot deleted the codeflash/optimize-function_kind-mbsmr7uy branch June 12, 2025 00:18

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

⚡️ Speed up function `function_kind` by 6%#319

⚡️ Speed up function `function_kind` by 6%#319
codeflash-ai[bot] wants to merge 1 commit into
updated-vsc-extensionfrom
codeflash/optimize-function_kind-mbsmr7uy

codeflash-ai Bot commented Jun 12, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

codeflash-ai Bot commented Jun 12, 2025

📄 6% (0.06x) speedup for function_kind in codeflash/code_utils/static_analysis.py

📝 Explanation and details

Summary of improvements.

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

📄 6% (0.06x) speedup for `function_kind` in `codeflash/code_utils/static_analysis.py`