Skip to content

Conversation

@codeflash-ai
Copy link
Contributor

@codeflash-ai codeflash-ai bot commented Jun 12, 2025

📄 52% (0.52x) speedup for function_kind in codeflash/code_utils/static_analysis.py

⏱️ Runtime : 1.33 milliseconds 871 microseconds (best of 222 runs)

📝 Explanation and details

Here are the main issues and opportunities for speed from your code and profile.

  1. Pointless Reverse For Loop:

    for _i in range(len(parents) - 1, -1, -1).
        continue

    This loop does absolutely nothing and should be deleted.

  2. Repeated List/Attr Lookups:

    • parents[0].type is accessed up to three times.
    • You do two membership tests for function types with a literal in ["FunctionDef", "AsyncFunctionDef"] (slow linear scan).
  3. Decorator Search:

    • Instead of scanning all decorators, you can hoist checks and return ASAP.
    • Use a set for faster comparison ({"classmethod", "staticmethod"}).

Here is the rewritten, faster code.

Summary of speedups:

  • Removed the pointless loop (for _i in ...: continue)
  • Reduced repeated attribute/membership lookups
  • Used early returns and a set for quick decorator name tests
  • No change to function signature or result; all comments preserved since only the removed code was nonfunctional

This version will run markedly faster.

Correctness verification report:

Test Status
⚙️ Existing Unit Tests 🔘 None Found
🌀 Generated Regression Tests 3551 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, Optional

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


# Mocked FunctionParent and FunctionKind for testing
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 function to create a function node with decorators
def make_func_node(name, decorators=None, async_def=False):
    if decorators is None:
        decorators = []
    if async_def:
        return ast.AsyncFunctionDef(
            name=name,
            args=ast.arguments(posonlyargs=[], args=[], kwonlyargs=[], kw_defaults=[], defaults=[]),
            body=[],
            decorator_list=decorators,
            returns=None,
            type_comment=None,
        )
    else:
        return ast.FunctionDef(
            name=name,
            args=ast.arguments(posonlyargs=[], args=[], kwonlyargs=[], kw_defaults=[], defaults=[]),
            body=[],
            decorator_list=decorators,
            returns=None,
            type_comment=None,
        )

# ========================
# Basic Test Cases
# ========================

def test_top_level_function():
    # Function with no parents (top-level)
    node = make_func_node("foo")
    parents = []
    codeflash_output = function_kind(node, parents) # 333ns -> 458ns

def test_top_level_async_function():
    # Async function with no parents (top-level)
    node = make_func_node("foo", async_def=True)
    parents = []
    codeflash_output = function_kind(node, parents) # 334ns -> 375ns

def test_function_nested_in_function():
    # Function inside another function
    node = make_func_node("bar")
    parents = [FunctionParent("FunctionDef")]
    codeflash_output = function_kind(node, parents) # 458ns -> 500ns

def test_function_nested_in_async_function():
    # Function inside an async function
    node = make_func_node("bar")
    parents = [FunctionParent("AsyncFunctionDef")]
    codeflash_output = function_kind(node, parents) # 417ns -> 500ns

def test_instance_method_in_class():
    # Method inside a class, no decorators
    node = make_func_node("baz")
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.04μs -> 583ns

def test_class_method_in_class():
    # Method inside class, with @classmethod decorator
    decorator = ast.Name(id="classmethod", ctx=ast.Load())
    node = make_func_node("baz", decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.21μs -> 750ns

def test_static_method_in_class():
    # Method inside class, with @staticmethod decorator
    decorator = ast.Name(id="staticmethod", ctx=ast.Load())
    node = make_func_node("baz", decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.29μs -> 750ns

# ========================
# Edge Test Cases
# ========================

def test_class_method_with_multiple_decorators():
    # Method with multiple decorators, @classmethod not first
    decorator1 = ast.Name(id="otherdecorator", ctx=ast.Load())
    decorator2 = ast.Name(id="classmethod", ctx=ast.Load())
    node = make_func_node("baz", decorators=[decorator1, decorator2])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.38μs -> 916ns

def test_static_method_with_multiple_decorators():
    # Method with multiple decorators, @staticmethod not first
    decorator1 = ast.Name(id="otherdecorator", ctx=ast.Load())
    decorator2 = ast.Name(id="staticmethod", ctx=ast.Load())
    node = make_func_node("baz", decorators=[decorator1, decorator2])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.46μs -> 958ns

def test_instance_method_with_non_name_decorator():
    # Method with decorator that is not ast.Name (should not match)
    decorator = ast.Attribute(value=ast.Name(id="foo", ctx=ast.Load()), attr="bar", ctx=ast.Load())
    node = make_func_node("baz", decorators=[decorator])
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.25μs -> 750ns

def test_function_with_unknown_parent_type():
    # Function with a parent that's not function or class
    node = make_func_node("foo")
    parents = [FunctionParent("Module")]
    codeflash_output = function_kind(node, parents) # 834ns -> 375ns

def test_function_with_empty_parents_list():
    # Function with empty parents list (should be top-level function)
    node = make_func_node("foo")
    parents = []
    codeflash_output = function_kind(node, parents) # 375ns -> 375ns

def test_function_with_extra_parents():
    # Function with multiple parents, first is ClassDef
    node = make_func_node("foo")
    parents = [FunctionParent("ClassDef"), FunctionParent("Module")]
    codeflash_output = function_kind(node, parents) # 1.12μs -> 583ns

def test_async_class_method():
    # Async method with @classmethod
    decorator = ast.Name(id="classmethod", ctx=ast.Load())
    node = make_func_node("baz", decorators=[decorator], async_def=True)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.12μs -> 791ns

def test_async_static_method():
    # Async method with @staticmethod
    decorator = ast.Name(id="staticmethod", ctx=ast.Load())
    node = make_func_node("baz", decorators=[decorator], async_def=True)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.29μs -> 791ns

def test_async_instance_method():
    # Async method with no decorator
    node = make_func_node("baz", async_def=True)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.08μs -> 583ns

def test_function_with_non_name_decorator():
    # Function with decorator that's not ast.Name (should be ignored)
    decorator = ast.Attribute(value=ast.Name(id="foo", ctx=ast.Load()), attr="bar", ctx=ast.Load())
    node = make_func_node("foo", decorators=[decorator])
    parents = []
    codeflash_output = function_kind(node, parents) # 375ns -> 375ns

def test_function_with_multiple_parents_first_is_function():
    # Function with multiple parents, first is a function
    node = make_func_node("foo")
    parents = [FunctionParent("FunctionDef"), FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 458ns -> 459ns

# ========================
# Large Scale Test Cases
# ========================

def test_many_top_level_functions():
    # Test 1000 top-level functions
    for i in range(1000):
        node = make_func_node(f"func_{i}")
        parents = []
        codeflash_output = function_kind(node, parents) # 83ns -> 83ns

def test_many_methods_in_class():
    # Test 1000 instance methods in a class
    for i in range(1000):
        node = make_func_node(f"method_{i}")
        parents = [FunctionParent("ClassDef")]
        codeflash_output = function_kind(node, parents) # 333ns -> 167ns

def test_many_class_methods_and_static_methods():
    # Test 500 classmethods and 500 staticmethods in a class
    for i in range(500):
        node = make_func_node(f"class_method_{i}", decorators=[ast.Name(id="classmethod", ctx=ast.Load())])
        parents = [FunctionParent("ClassDef")]
        codeflash_output = function_kind(node, parents) # 375ns -> 209ns
    for i in range(500):
        node = make_func_node(f"static_method_{i}", decorators=[ast.Name(id="staticmethod", ctx=ast.Load())])
        parents = [FunctionParent("ClassDef")]
        codeflash_output = function_kind(node, parents) # 375ns -> 209ns

def test_large_mixed_decorator_list():
    # Test a method with 1000 decorators, only one is 'classmethod'
    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("baz", decorators=decorators)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 53.7μs -> 47.0μs

def test_large_mixed_decorator_list_staticmethod():
    # Test a method with 1000 decorators, only one is 'staticmethod'
    decorators = [ast.Name(id=f"decorator_{i}", ctx=ast.Load()) for i in range(999)]
    decorators.insert(500, ast.Name(id="staticmethod", ctx=ast.Load()))
    node = make_func_node("baz", decorators=decorators)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 27.8μs -> 24.3μs

def test_large_mixed_decorator_list_instance_method():
    # Test a method with 1000 decorators, none are classmethod or staticmethod
    decorators = [ast.Name(id=f"decorator_{i}", ctx=ast.Load()) for i in range(1000)]
    node = make_func_node("baz", decorators=decorators)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 53.9μs -> 46.8μs

# ========================
# Mutation Testing Guards
# ========================

def test_mutation_guard_classmethod_vs_staticmethod():
    # If both classmethod and staticmethod are present, classmethod wins (first found in list)
    decorators = [
        ast.Name(id="staticmethod", ctx=ast.Load()),
        ast.Name(id="classmethod", ctx=ast.Load())
    ]
    node = make_func_node("baz", decorators=decorators)
    parents = [FunctionParent("ClassDef")]
    # Should return classmethod because it's found first in the loop
    codeflash_output = function_kind(node, parents) # 1.21μs -> 792ns

def test_mutation_guard_staticmethod_vs_classmethod():
    # If staticmethod is before classmethod, staticmethod wins
    decorators = [
        ast.Name(id="staticmethod", ctx=ast.Load()),
        ast.Name(id="other", ctx=ast.Load()),
    ]
    node = make_func_node("baz", decorators=decorators)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.21μs -> 792ns

def test_mutation_guard_classmethod_vs_staticmethod_order():
    # If classmethod is before staticmethod, classmethod wins
    decorators = [
        ast.Name(id="classmethod", ctx=ast.Load()),
        ast.Name(id="staticmethod", ctx=ast.Load())
    ]
    node = make_func_node("baz", decorators=decorators)
    parents = [FunctionParent("ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.17μs -> 708ns
# 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 dataclasses import dataclass
from enum import Enum

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


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

# Mocked FunctionParent dataclass for testing
@dataclass
class FunctionParent:
    type: str  # e.g., "Module", "ClassDef", "FunctionDef", "AsyncFunctionDef"
from codeflash.code_utils.static_analysis import function_kind


# Helper function to create a decorated function node
def make_func_node(name="f", decorators=None, async_=False):
    if decorators is None:
        decorators = []
    if async_:
        return ast.AsyncFunctionDef(
            name=name,
            args=ast.arguments(posonlyargs=[], args=[], kwonlyargs=[], kw_defaults=[], defaults=[]),
            body=[],
            decorator_list=decorators,
            returns=None,
            type_comment=None
        )
    else:
        return ast.FunctionDef(
            name=name,
            args=ast.arguments(posonlyargs=[], args=[], kwonlyargs=[], kw_defaults=[], defaults=[]),
            body=[],
            decorator_list=decorators,
            returns=None,
            type_comment=None
        )

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

def test_module_level_function():
    # Function at module level (no parents)
    node = make_func_node()
    parents = []
    codeflash_output = function_kind(node, parents) # 333ns -> 375ns

def test_module_level_async_function():
    # Async function at module level (no parents)
    node = make_func_node(async_=True)
    parents = []
    codeflash_output = function_kind(node, parents) # 375ns -> 375ns

def test_function_inside_function():
    # Function nested inside another function
    node = make_func_node()
    parents = [FunctionParent(type="FunctionDef")]
    codeflash_output = function_kind(node, parents) # 458ns -> 417ns

def test_function_inside_async_function():
    # Function nested inside an async function
    node = make_func_node()
    parents = [FunctionParent(type="AsyncFunctionDef")]
    codeflash_output = function_kind(node, parents) # 458ns -> 500ns

def test_instance_method_in_class():
    # Regular method in a class (no decorators)
    node = make_func_node()
    parents = [FunctionParent(type="ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.04μs -> 583ns

def test_class_method_in_class():
    # Method decorated with @classmethod
    decorator = ast.Name(id="classmethod")
    node = make_func_node(decorators=[decorator])
    parents = [FunctionParent(type="ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.21μs -> 750ns

def test_static_method_in_class():
    # Method decorated with @staticmethod
    decorator = ast.Name(id="staticmethod")
    node = make_func_node(decorators=[decorator])
    parents = [FunctionParent(type="ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.29μs -> 750ns

def test_class_method_and_static_method_decorators_order():
    # Both @classmethod and @staticmethod, should prefer first found
    decorator1 = ast.Name(id="classmethod")
    decorator2 = ast.Name(id="staticmethod")
    node = make_func_node(decorators=[decorator1, decorator2])
    parents = [FunctionParent(type="ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.17μs -> 750ns

def test_static_method_and_class_method_decorators_order():
    # Both @staticmethod and @classmethod, should prefer first found
    decorator1 = ast.Name(id="staticmethod")
    decorator2 = ast.Name(id="classmethod")
    node = make_func_node(decorators=[decorator1, decorator2])
    parents = [FunctionParent(type="ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.29μs -> 709ns

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

def test_class_method_with_alias():
    # Decorator is not a simple ast.Name (e.g., ast.Attribute or alias)
    decorator = ast.Attribute(value=ast.Name(id="something"), attr="classmethod")
    node = make_func_node(decorators=[decorator])
    parents = [FunctionParent(type="ClassDef")]
    # Should NOT recognize as classmethod, so should be instance method
    codeflash_output = function_kind(node, parents) # 1.21μs -> 709ns

def test_static_method_with_alias():
    # Decorator is not a simple ast.Name (e.g., ast.Attribute or alias)
    decorator = ast.Attribute(value=ast.Name(id="something"), attr="staticmethod")
    node = make_func_node(decorators=[decorator])
    parents = [FunctionParent(type="ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.17μs -> 708ns

def test_function_with_other_decorators():
    # Method with unrelated decorator
    decorator = ast.Name(id="otherdecorator")
    node = make_func_node(decorators=[decorator])
    parents = [FunctionParent(type="ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.25μs -> 791ns

def test_function_with_multiple_non_method_decorators():
    # Multiple unrelated decorators
    decorators = [ast.Name(id="foo"), ast.Name(id="bar")]
    node = make_func_node(decorators=decorators)
    parents = [FunctionParent(type="ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.38μs -> 916ns

def test_function_with_no_decorators_in_class():
    # No decorators in class context
    node = make_func_node()
    parents = [FunctionParent(type="ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.04μs -> 583ns

def test_function_with_unknown_parent_type():
    # Parent type is not handled
    node = make_func_node()
    parents = [FunctionParent(type="UnknownType")]
    codeflash_output = function_kind(node, parents) # 834ns -> 375ns

def test_function_with_empty_parents_list():
    # No parents supplied
    node = make_func_node()
    parents = []
    codeflash_output = function_kind(node, parents) # 375ns -> 375ns

def test_function_with_multiple_parents():
    # Nested inside a class, which is inside a function
    node = make_func_node()
    parents = [FunctionParent(type="ClassDef"), FunctionParent(type="FunctionDef")]
    codeflash_output = function_kind(node, parents) # 1.12μs -> 583ns

def test_function_with_module_parent():
    # Parent is module
    node = make_func_node()
    parents = [FunctionParent(type="Module")]
    codeflash_output = function_kind(node, parents) # 833ns -> 375ns

def test_function_with_large_decorator_list():
    # Large number of unrelated decorators
    decorators = [ast.Name(id=f"decorator{i}") for i in range(100)]
    node = make_func_node(decorators=decorators)
    parents = [FunctionParent(type="ClassDef")]
    codeflash_output = function_kind(node, parents) # 7.21μs -> 5.88μs

def test_async_method_in_class():
    # Async method in class (should still be instance method)
    node = make_func_node(async_=True)
    parents = [FunctionParent(type="ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.00μs -> 583ns

def test_async_class_method():
    # Async method with @classmethod
    decorator = ast.Name(id="classmethod")
    node = make_func_node(decorators=[decorator], async_=True)
    parents = [FunctionParent(type="ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.12μs -> 791ns

def test_async_static_method():
    # Async method with @staticmethod
    decorator = ast.Name(id="staticmethod")
    node = make_func_node(decorators=[decorator], async_=True)
    parents = [FunctionParent(type="ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.29μs -> 791ns

def test_function_with_decorator_not_name():
    # Decorator is a call, not a name
    decorator = ast.Call(func=ast.Name(id="classmethod"), args=[], keywords=[])
    node = make_func_node(decorators=[decorator])
    parents = [FunctionParent(type="ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.17μs -> 750ns

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

def test_many_functions_in_class():
    # Test many functions in a class context, all should be instance methods
    for i in range(100):
        node = make_func_node(name=f"f{i}")
        parents = [FunctionParent(type="ClassDef")]
        codeflash_output = function_kind(node, parents) # 333ns -> 208ns

def test_many_class_methods_in_class():
    # Test many classmethods in a class context
    decorator = ast.Name(id="classmethod")
    for i in range(100):
        node = make_func_node(name=f"f{i}", decorators=[decorator])
        parents = [FunctionParent(type="ClassDef")]
        codeflash_output = function_kind(node, parents) # 416ns -> 250ns

def test_many_static_methods_in_class():
    # Test many staticmethods in a class context
    decorator = ast.Name(id="staticmethod")
    for i in range(100):
        node = make_func_node(name=f"f{i}", decorators=[decorator])
        parents = [FunctionParent(type="ClassDef")]
        codeflash_output = function_kind(node, parents) # 416ns -> 250ns

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

def test_many_nested_functions():
    # Test many functions nested inside functions
    for i in range(100):
        node = make_func_node(name=f"f{i}")
        parents = [FunctionParent(type="FunctionDef")]
        codeflash_output = function_kind(node, parents) # 125ns -> 166ns

def test_large_number_of_parents():
    # Test with a long parent chain, but topmost is ClassDef
    node = make_func_node()
    parents = [FunctionParent(type="ClassDef")] + [FunctionParent(type="FunctionDef")] * 999
    codeflash_output = function_kind(node, parents) # 10.9μs -> 584ns

def test_large_number_of_parents_not_classdef():
    # Test with a long parent chain, but topmost is not ClassDef
    node = make_func_node()
    parents = [FunctionParent(type="Module")] + [FunctionParent(type="FunctionDef")] * 999
    codeflash_output = function_kind(node, parents) # 10.4μs -> 375ns

def test_large_decorator_list_with_classmethod_last():
    # Large decorator list, classmethod at the end
    decorators = [ast.Name(id=f"decorator{i}") for i in range(99)] + [ast.Name(id="classmethod")]
    node = make_func_node(decorators=decorators)
    parents = [FunctionParent(type="ClassDef")]
    # Should still find the classmethod
    codeflash_output = function_kind(node, parents) # 7.25μs -> 6.00μs

def test_large_decorator_list_with_staticmethod_first():
    # Large decorator list, staticmethod at the start
    decorators = [ast.Name(id="staticmethod")] + [ast.Name(id=f"decorator{i}") for i in range(99)]
    node = make_func_node(decorators=decorators)
    parents = [FunctionParent(type="ClassDef")]
    codeflash_output = function_kind(node, parents) # 1.21μs -> 792ns
# 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-mbsnv8im and push.

Codeflash

Here are the main issues and opportunities for speed from your code and profile.

1. **Pointless Reverse For Loop:**  
   ```python
   for _i in range(len(parents) - 1, -1, -1).
       continue
   ```
   This loop does absolutely nothing and should be deleted.

2. **Repeated List/Attr Lookups:**  
   - `parents[0].type` is accessed up to three times.  
   - You do two membership tests for function types with a literal `in ["FunctionDef", "AsyncFunctionDef"]` (slow linear scan).  

3. **Decorator Search:**  
   - Instead of scanning all decorators, you can hoist checks and return ASAP.
   - Use a set for faster comparison (`{"classmethod", "staticmethod"}`).

Here is the rewritten, faster code.



**Summary of speedups:**
- Removed the pointless loop (`for _i in ...: continue`)
- Reduced repeated attribute/membership lookups
- Used early returns and a set for quick decorator name tests
- No change to function signature or result; all comments preserved since only the removed code was nonfunctional

**This version will run markedly faster.**
@codeflash-ai codeflash-ai bot added the ⚡️ codeflash Optimization PR opened by Codeflash AI label Jun 12, 2025
@codeflash-ai codeflash-ai bot requested a review from KRRT7 June 12, 2025 00:49
@KRRT7 KRRT7 closed this Jun 12, 2025
@codeflash-ai codeflash-ai bot deleted the codeflash/optimize-function_kind-mbsnv8im branch June 12, 2025 01:11
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

⚡️ codeflash Optimization PR opened by Codeflash AI

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants