⚡️ Speed up function build_experiment_def by 10%

[codeflash-ai]

📄 10% (0.10x) speedup for build_experiment_def in higgsfield/internal/experiment/ast_parser.py

⏱️ Runtime : 1.09 milliseconds → 994 microseconds (best of 86 runs)

📝 Explanation and details

The optimized code achieves a 10% speedup through several key performance optimizations:

Primary Optimizations:

1. Local variable caching for frequently accessed globals: The code caches builder.get, type_dict.get, and AST class references (ast.Call, ast.Name, etc.) as local variables. This eliminates repeated global lookups during the main loop, which is critical since the function processes multiple decorators and keywords.

2. Reduced attribute access overhead: Instead of accessing decorator.args multiple times, it's cached as args = decorator.args. Similarly, kw.value is cached as value = kw.value to avoid repeated attribute lookups.

3. Optimized isinstance checks: The code combines two separate isinstance calls into a single tuple check: isinstance(value, (ast_Tuple, ast_List)) instead of isinstance(kw.value, ast.Tuple) or isinstance(kw.value, ast.List).

4. List comprehension over manual loop: Replaces the manual loop with .append() calls with a more efficient list comprehension: vals = [elt.value for elt in value.elts].

5. Type comparison optimization: Uses is comparison (dec_type is type_Expdec) instead of equality (type(dec) == Expdec), which is faster for type checks.

Performance Impact by Test Case:

• Large-scale tests show the biggest gains: Tests with 100-200 parameters see 10-13% speedups, demonstrating that the optimizations scale well with the number of decorators processed.
• Basic cases see modest improvements: Simple cases with few decorators show 2-4% improvements.
• Edge cases may be slightly slower: Some error-path tests show minor slowdowns due to the overhead of local variable setup, but this is negligible compared to the gains in normal execution paths.

The optimizations are most effective for functions with many decorators and complex parameter configurations, which matches the typical use case for experiment definition parsing.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import ast

# imports
import pytest
from higgsfield.internal.experiment.ast_parser import build_experiment_def


# Dummy classes to mimic Expdec and Paramdec for testing
class Expdec:
    def __init__(self):
        self.arg_pairs = {}

    def add_arg_pair(self, k, v):
        self.arg_pairs[k] = v

class Paramdec:
    def __init__(self):
        self.arg_pairs = {}

    def add_arg_pair(self, k, v):
        self.arg_pairs[k] = v

# function to test
builder = {"experiment": Expdec, "param": Paramdec}

type_dict = {
    "str": str,
    "int": int,
    "float": float,
    "bool": bool,
}
from higgsfield.internal.experiment.ast_parser import build_experiment_def

# unit tests

# Helper function to construct ast.Call for decorators
def make_decorator(name, args=None, keywords=None):
    args = args or []
    keywords = keywords or []
    return ast.Call(
        func=ast.Name(id=name, ctx=ast.Load()),
        args=args,
        keywords=keywords
    )

def make_constant(val):
    return ast.Constant(value=val)

def make_name(val):
    return ast.Name(id=val, ctx=ast.Load())

def make_tuple(vals):
    return ast.Tuple(elts=[make_constant(v) for v in vals], ctx=ast.Load())

def make_list(vals):
    return ast.List(elts=[make_constant(v) for v in vals], ctx=ast.Load())

def make_keyword(arg, value):
    return ast.keyword(arg=arg, value=value)

def make_function_def(decorators):
    return ast.FunctionDef(
        name="test_func",
        args=ast.arguments(posonlyargs=[], args=[], kwonlyargs=[], kw_defaults=[], defaults=[]),
        body=[],
        decorator_list=decorators
    )

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




def test_basic_param_bool_type():
    # Param with bool type keyword
    exp_dec = make_decorator(
        "experiment",
        args=[make_constant("exp4")]
    )
    param_dec = make_decorator(
        "param",
        args=[make_constant("flag")],
        keywords=[
            make_keyword("type", make_name("bool")),
            make_keyword("default", make_constant(True))
        ]
    )
    node = make_function_def([exp_dec, param_dec])
    codeflash_output = build_experiment_def(node); result = codeflash_output # 11.2μs -> 10.9μs (2.37% faster)
    experiment, params = result

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

def test_no_decorators_returns_none():
    # No decorators at all
    node = make_function_def([])
    codeflash_output = build_experiment_def(node); result = codeflash_output # 815ns -> 1.41μs (42.2% slower)

def test_only_param_decorator_returns_none():
    # Only param decorator, no experiment
    param_dec = make_decorator(
        "param",
        args=[make_constant("p3")],
        keywords=[make_keyword("type", make_name("str"))]
    )
    node = make_function_def([param_dec])
    codeflash_output = build_experiment_def(node); result = codeflash_output # 6.59μs -> 6.41μs (2.81% faster)

def test_non_call_decorator_ignored():
    # Non-call decorator in the list
    exp_dec = make_decorator("experiment", args=[make_constant("exp5")])
    node = make_function_def([ast.Name(id="not_a_call", ctx=ast.Load()), exp_dec])
    codeflash_output = build_experiment_def(node); result = codeflash_output # 4.86μs -> 5.06μs (3.97% slower)
    experiment, params = result

def test_decorator_with_no_func_stops():
    # ast.Call with no func attribute
    broken_dec = ast.Call(func=None, args=[], keywords=[])
    node = make_function_def([broken_dec])
    codeflash_output = build_experiment_def(node); result = codeflash_output # 1.18μs -> 1.63μs (27.5% slower)

def test_unknown_decorator_stops():
    # Decorator not in builder
    unknown_dec = make_decorator("unknown", args=[make_constant("x")])
    node = make_function_def([unknown_dec])
    codeflash_output = build_experiment_def(node); result = codeflash_output # 1.45μs -> 1.69μs (14.1% slower)

def test_experiment_with_multiple_args_raises():
    # Experiment decorator with more than one arg
    exp_dec = make_decorator(
        "experiment",
        args=[make_constant("exp6"), make_constant("extra")]
    )
    node = make_function_def([exp_dec])
    with pytest.raises(ValueError):
        build_experiment_def(node) # 3.76μs -> 4.16μs (9.48% slower)

def test_param_with_multiple_args_raises():
    # Param decorator with more than one arg
    exp_dec = make_decorator("experiment", args=[make_constant("exp7")])
    param_dec = make_decorator(
        "param",
        args=[make_constant("p4"), make_constant("extra")]
    )
    node = make_function_def([exp_dec, param_dec])
    with pytest.raises(ValueError):
        build_experiment_def(node) # 7.19μs -> 7.25μs (0.787% slower)

def test_duplicate_param_names_raises():
    # Two param decorators with same name
    exp_dec = make_decorator("experiment", args=[make_constant("exp8")])
    param_dec1 = make_decorator("param", args=[make_constant("dup")])
    param_dec2 = make_decorator("param", args=[make_constant("dup")])
    node = make_function_def([exp_dec, param_dec1, param_dec2])
    with pytest.raises(ValueError):
        build_experiment_def(node) # 9.20μs -> 9.24μs (0.541% slower)

def test_multiple_experiment_decorators_raises():
    # Two experiment decorators
    exp_dec1 = make_decorator("experiment", args=[make_constant("exp9")])
    exp_dec2 = make_decorator("experiment", args=[make_constant("exp10")])
    node = make_function_def([exp_dec1, exp_dec2])
    with pytest.raises(ValueError):
        build_experiment_def(node) # 6.41μs -> 6.75μs (5.02% slower)




def test_large_number_of_params():
    # Experiment with many param decorators
    exp_dec = make_decorator("experiment", args=[make_constant("exp_large")])
    param_decs = [
        make_decorator(
            "param",
            args=[make_constant(f"p{i}")],
            keywords=[make_keyword("type", make_name("int")), make_keyword("default", make_constant(i))]
        )
        for i in range(100)
    ]
    node = make_function_def([exp_dec] + param_decs)
    codeflash_output = build_experiment_def(node); result = codeflash_output # 206μs -> 186μs (10.7% faster)
    experiment, params = result
    for i in range(100):
        pass


def test_large_number_of_params_and_keywords():
    # Experiment with many params, each with multiple keywords
    exp_dec = make_decorator("experiment", args=[make_constant("exp_multi_kw")])
    param_decs = [
        make_decorator(
            "param",
            args=[make_constant(f"p{i}")],
            keywords=[
                make_keyword("type", make_name("float")),
                make_keyword("default", make_constant(float(i))),
                make_keyword("description", make_constant(f"param {i}"))
            ]
        )
        for i in range(50)
    ]
    node = make_function_def([exp_dec] + param_decs)
    codeflash_output = build_experiment_def(node); result = codeflash_output # 128μs -> 113μs (13.4% faster)
    experiment, params = result
    for i in range(50):
        pass

def test_large_scale_with_mixed_types():
    # Params with mixed types and large scale
    exp_dec = make_decorator("experiment", args=[make_constant("exp_mixed")])
    param_decs = []
    types = ["int", "str", "float", "bool"]
    for i in range(200):
        t = types[i % 4]
        param_decs.append(
            make_decorator(
                "param",
                args=[make_constant(f"p{i}")],
                keywords=[make_keyword("type", make_name(t)), make_keyword("default", make_constant(i))]
            )
        )
    node = make_function_def([exp_dec] + param_decs)
    codeflash_output = build_experiment_def(node); result = codeflash_output # 401μs -> 356μs (12.7% faster)
    experiment, params = result
    for i in range(200):
        t = types[i % 4]
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import ast

# imports
import pytest
from higgsfield.internal.experiment.ast_parser import build_experiment_def


# Dummy classes to simulate Expdec and Paramdec behavior for testing
class Expdec:
    def __init__(self):
        self.arg_pairs = {}

    def add_arg_pair(self, key, value):
        # Simulate expected behavior: store key-value pairs
        self.arg_pairs[key] = value

class Paramdec:
    def __init__(self):
        self.arg_pairs = {}

    def add_arg_pair(self, key, value):
        # Simulate expected behavior: store key-value pairs
        self.arg_pairs[key] = value

# function to test (copied from above, with builder using our dummy classes)
builder = {"experiment": Expdec, "param": Paramdec}

type_dict = {
    "str": str,
    "int": int,
    "float": float,
    "bool": bool,
}
from higgsfield.internal.experiment.ast_parser import build_experiment_def


# Helper function to create decorator AST nodes
def make_decorator(name, args=None, keywords=None):
    args = args or []
    keywords = keywords or []
    return ast.Call(
        func=ast.Name(id=name, ctx=ast.Load()),
        args=args,
        keywords=keywords
    )

def make_functiondef(decorators):
    return ast.FunctionDef(
        name="test_func",
        args=ast.arguments(posonlyargs=[], args=[], kwonlyargs=[], kw_defaults=[], defaults=[]),
        body=[],
        decorator_list=decorators
    )

# -------------------- UNIT TESTS --------------------

# 1. Basic Test Cases



def test_basic_only_param_no_experiment():
    # Test a function with only a param decorator (should return None)
    param_dec = make_decorator("param", args=[ast.Constant(value="p2")], keywords=[ast.keyword(arg="type", value=ast.Name(id="float", ctx=ast.Load()))])
    node = make_functiondef([param_dec])
    codeflash_output = build_experiment_def(node); result = codeflash_output # 8.36μs -> 8.08μs (3.44% faster)



def test_no_decorators():
    # Function with no decorators should return None
    node = make_functiondef([])
    codeflash_output = build_experiment_def(node); result = codeflash_output # 1.04μs -> 1.85μs (43.8% slower)

def test_non_call_decorator_ignored():
    # Non-call decorator should be ignored
    node = make_functiondef([ast.Name(id="experiment", ctx=ast.Load())])
    codeflash_output = build_experiment_def(node); result = codeflash_output # 1.26μs -> 1.68μs (24.8% slower)

def test_decorator_with_missing_func():
    # Decorator with missing func attribute (simulate by passing ast.Call with func=None)
    dec = ast.Call(func=None, args=[], keywords=[])
    node = make_functiondef([dec])
    codeflash_output = build_experiment_def(node); result = codeflash_output # 1.29μs -> 1.66μs (22.4% slower)

def test_unknown_decorator():
    # Decorator with unknown name should cause stop
    unknown_dec = make_decorator("unknown", args=[ast.Constant(value="x")])
    node = make_functiondef([unknown_dec])
    codeflash_output = build_experiment_def(node); result = codeflash_output # 1.50μs -> 1.73μs (13.1% slower)

def test_experiment_with_multiple_args_raises():
    # Experiment decorator with more than one arg should raise
    experiment_dec = make_decorator("experiment", args=[ast.Constant(value="exp"), ast.Constant(value="extra")])
    node = make_functiondef([experiment_dec])
    with pytest.raises(ValueError):
        build_experiment_def(node) # 4.22μs -> 4.54μs (7.04% slower)

def test_param_with_multiple_args_raises():
    # Param decorator with more than one arg should raise
    param_dec = make_decorator("param", args=[ast.Constant(value="p"), ast.Constant(value="extra")])
    experiment_dec = make_decorator("experiment", args=[ast.Constant(value="exp")])
    node = make_functiondef([experiment_dec, param_dec])
    with pytest.raises(ValueError):
        build_experiment_def(node) # 8.01μs -> 7.69μs (4.04% faster)

def test_param_with_duplicate_name_raises():
    # Two param decorators with same name should raise
    param_dec1 = make_decorator("param", args=[ast.Constant(value="dup")])
    param_dec2 = make_decorator("param", args=[ast.Constant(value="dup")])
    experiment_dec = make_decorator("experiment", args=[ast.Constant(value="exp")])
    node = make_functiondef([experiment_dec, param_dec1, param_dec2])
    with pytest.raises(ValueError):
        build_experiment_def(node) # 9.48μs -> 9.46μs (0.254% faster)

def test_multiple_experiment_decorators_raises():
    # Two experiment decorators should raise
    experiment_dec1 = make_decorator("experiment", args=[ast.Constant(value="exp1")])
    experiment_dec2 = make_decorator("experiment", args=[ast.Constant(value="exp2")])
    node = make_functiondef([experiment_dec1, experiment_dec2])
    with pytest.raises(ValueError):
        build_experiment_def(node) # 6.75μs -> 6.51μs (3.67% faster)



def test_many_params():
    # Function with experiment and many param decorators
    experiment_dec = make_decorator("experiment", args=[ast.Constant(value="exp")])
    param_decs = []
    for i in range(100):  # 100 param decorators
        param_decs.append(make_decorator("param", args=[ast.Constant(value=f"p{i}")], keywords=[ast.keyword(arg="type", value=ast.Name(id="int", ctx=ast.Load()))]))
    node = make_functiondef([experiment_dec] + param_decs)
    codeflash_output = build_experiment_def(node); result = codeflash_output # 170μs -> 152μs (12.2% faster)
    exp, params = result
    for i in range(100):
        pass



def test_many_params_with_varied_types():
    # Many param decorators with varied types
    experiment_dec = make_decorator("experiment", args=[ast.Constant(value="exp_varied")])
    param_decs = []
    for i in range(50):
        t = ["int", "float", "str", "bool"][i % 4]
        param_decs.append(make_decorator("param", args=[ast.Constant(value=f"p{i}")], keywords=[ast.keyword(arg="type", value=ast.Name(id=t, ctx=ast.Load()))]))
    node = make_functiondef([experiment_dec] + param_decs)
    codeflash_output = build_experiment_def(node); result = codeflash_output # 90.2μs -> 84.7μs (6.43% faster)
    _, params = result
    for i in range(50):
        t = ["int", "float", "str", "bool"][i % 4]
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from ast import FunctionDef
from higgsfield.internal.experiment.ast_parser import build_experiment_def
import pytest

def test_build_experiment_def():
    with pytest.raises(AttributeError, match="'FunctionDef'\\ object\\ has\\ no\\ attribute\\ 'decorator_list'"):
        build_experiment_def(FunctionDef())

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_0kzwu12q/tmpikm_6za0/test_concolic_coverage.py::test_build_experiment_def	1.50μs	2.14μs	-29.8%⚠️

To edit these changes git checkout codeflash/optimize-build_experiment_def-mglp4fdp and push.