⚡️ Speed up function _replace_keys by 13%

[codeflash-ai]

📄 13% (0.13x) speedup for _replace_keys in framework/py/flwr/simulation/run_simulation.py

⏱️ Runtime : 4.62 milliseconds → 4.09 milliseconds (best of 68 runs)

📝 Explanation and details

The optimization achieves a 13% speedup by reducing unnecessary function calls and improving type checking efficiency:

Key optimizations:

1. Reduced recursive calls: The optimized version uses inline type checking (isinstance(v, (dict, list))) within the comprehensions to avoid calling _replace_keys on primitive values (strings, numbers, etc.). This eliminates thousands of unnecessary function calls for leaf nodes.

2. Early exit for primitives: Instead of always recursively calling _replace_keys and then checking types inside the function, the optimized version checks types first and only recurses when needed. This is particularly effective for structures with many primitive values.

3. Improved control flow: Changed the second if to elif to avoid checking both dict and list conditions for the same object.

4. Dictionary items caching: Stores d.items() in a local variable to avoid repeated method lookups during iteration.

Performance impact by test type:

• Large flat dictionaries: Most beneficial for structures with many primitive values where recursive calls are avoided
• Mixed nested structures: Effective when there are many leaf nodes (non-dict, non-list values)
• Deep nesting: Still provides benefits by reducing function call overhead at every level

The line profiler shows the optimization reduces total execution time from 27.8ms to 24.2ms, with the most significant improvement coming from avoiding redundant isinstance checks and recursive calls for primitive values.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from typing import Any

# imports
import pytest  # used for our unit tests
from simulation.run_simulation import _replace_keys

# unit tests

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

def test_simple_dict_replacement():
    # Basic dict, one key to replace
    d = {"foo": 1, "bar": 2}
    codeflash_output = _replace_keys(d, "foo", "baz"); result = codeflash_output

def test_no_match_key():
    # No keys match the string to replace
    d = {"a": 1, "b": 2}
    codeflash_output = _replace_keys(d, "z", "y"); result = codeflash_output

def test_multiple_keys_match():
    # Multiple keys match the string to replace
    d = {"foo": 1, "foobar": 2, "barfoo": 3}
    codeflash_output = _replace_keys(d, "foo", "baz"); result = codeflash_output

def test_nested_dicts():
    # Nested dictionaries
    d = {"foo": {"foo": 2, "bar": 3}, "bar": 1}
    codeflash_output = _replace_keys(d, "foo", "baz"); result = codeflash_output

def test_list_of_dicts():
    # List containing dictionaries
    d = [{"foo": 1}, {"bar": 2}]
    codeflash_output = _replace_keys(d, "foo", "baz"); result = codeflash_output

def test_dict_with_list_value():
    # Dict with a list as a value
    d = {"foo": [ {"foo": 1}, {"bar": 2} ]}
    codeflash_output = _replace_keys(d, "foo", "baz"); result = codeflash_output

def test_empty_dict():
    # Empty dict should return empty dict
    d = {}
    codeflash_output = _replace_keys(d, "foo", "baz"); result = codeflash_output

def test_empty_list():
    # Empty list should return empty list
    d = []
    codeflash_output = _replace_keys(d, "foo", "baz"); result = codeflash_output

def test_non_dict_non_list():
    # Non-dict, non-list input should return as is
    d = 42
    codeflash_output = _replace_keys(d, "foo", "baz"); result = codeflash_output

def test_key_with_multiple_occurrences():
    # Key contains the match string multiple times
    d = {"foofoo": 1, "bar": 2}
    codeflash_output = _replace_keys(d, "foo", "baz"); result = codeflash_output

def test_replace_with_empty_string():
    # Replace match with empty string
    d = {"foofoo": 1, "barfoo": 2}
    codeflash_output = _replace_keys(d, "foo", ""); result = codeflash_output

def test_replace_empty_string_in_key():
    # Replace empty string in key (should insert target between every character and at ends)
    d = {"abc": 1}
    codeflash_output = _replace_keys(d, "", "X"); result = codeflash_output

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


def test_list_of_lists():
    # List of lists, should recurse into all
    d = [[{"foo": 1}], [{"bar": 2}]]
    codeflash_output = _replace_keys(d, "foo", "baz"); result = codeflash_output

def test_key_is_empty_string():
    # Key is empty string
    d = {"": 1, "foo": 2}
    codeflash_output = _replace_keys(d, "foo", "bar"); result = codeflash_output

def test_replace_to_existing_key():
    # Replacement causes key collision
    d = {"foo": 1, "bar": 2}
    codeflash_output = _replace_keys(d, "foo", "bar"); result = codeflash_output

def test_nested_empty_dicts_and_lists():
    # Deeply nested empty dicts and lists
    d = {"foo": [], "bar": {}, "baz": [{"qux": {}}]}
    codeflash_output = _replace_keys(d, "foo", "zap"); result = codeflash_output


def test_value_is_none():
    # Value is None, should be preserved
    d = {"foo": None}
    codeflash_output = _replace_keys(d, "foo", "bar"); result = codeflash_output

def test_deeply_nested_mixed_types():
    # Deeply nested structure with mixed types
    d = {
        "foo": [
            {"foo": {"foo": 1}},
            {"bar": 2}
        ],
        "baz": {"foo": 3}
    }
    codeflash_output = _replace_keys(d, "foo", "qux"); result = codeflash_output

def test_key_is_int_string():
    # Key is a string that looks like an int
    d = {"1": "foo", "foo": "bar"}
    codeflash_output = _replace_keys(d, "foo", "baz"); result = codeflash_output

def test_key_is_unicode():
    # Key contains unicode characters
    d = {"föö": 1, "bar": 2}
    codeflash_output = _replace_keys(d, "ö", "o"); result = codeflash_output




def test_large_flat_dict():
    # Large flat dict with 1000 keys, half containing the match string
    d = {f"foo_{i}": i for i in range(500)}
    d.update({f"bar_{i}": i for i in range(500)})
    codeflash_output = _replace_keys(d, "foo", "baz"); result = codeflash_output
    expected = {f"baz_{i}": i for i in range(500)}
    expected.update({f"bar_{i}": i for i in range(500)})

def test_large_nested_dicts():
    # Large nested dict: 10 levels, each with 10 keys
    d = current = {}
    for i in range(10):
        new = {f"foo_{i}_{j}": j for j in range(10)}
        current["foo_level_" + str(i)] = new
        current = new
    codeflash_output = _replace_keys(d, "foo", "bar"); result = codeflash_output
    # Check that all keys at every level have 'foo' replaced by 'bar'
    current = result
    for i in range(10):
        current = current["bar_level_" + str(i)]
        for j in range(10):
            pass

def test_large_list_of_dicts():
    # List of 1000 dicts
    d = [{"foo": i, "bar": i*2} for i in range(1000)]
    codeflash_output = _replace_keys(d, "foo", "baz"); result = codeflash_output
    for i in range(1000):
        pass


def test_large_dict_with_no_matches():
    # Large dict, no keys contain match string
    d = {f"key_{i}": i for i in range(1000)}
    codeflash_output = _replace_keys(d, "foo", "baz"); result = codeflash_output

def test_large_dict_replace_with_empty_string():
    # Large dict, replace with empty string
    d = {f"foo_{i}": i for i in range(1000)}
    codeflash_output = _replace_keys(d, "foo_", ""); result = codeflash_output
    expected = {f"{i}": i for i in range(1000)}
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import copy
# function to test
# (as provided in the prompt)
from typing import Any

# imports
import pytest  # used for our unit tests
from simulation.run_simulation import _replace_keys

# unit tests

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

def test_empty_dict():
    # Test that an empty dict returns an empty dict
    codeflash_output = _replace_keys({}, "a", "b")

def test_empty_list():
    # Test that an empty list returns an empty list
    codeflash_output = _replace_keys([], "x", "y")

def test_no_match_key():
    # Test that a dict with no matching key substrings is unchanged
    d = {"foo": 1, "bar": 2}
    codeflash_output = _replace_keys(d, "baz", "qux")

def test_simple_replace():
    # Test that a single key is replaced
    d = {"foo": 1, "bar": 2}
    expected = {"fxx": 1, "bar": 2}
    codeflash_output = _replace_keys(d, "oo", "xx")

def test_multiple_keys_replace():
    # Test that multiple keys are replaced
    d = {"apple": 1, "apricot": 2, "banana": 3}
    expected = {"xpple": 1, "xpricot": 2, "banana": 3}
    codeflash_output = _replace_keys(d, "a", "x")

def test_nested_dict():
    # Test that nested dict keys are replaced recursively
    d = {"foo": {"bar": 1, "baz": 2}, "qux": 3}
    expected = {"fxx": {"bxxr": 1, "bxxz": 2}, "qux": 3}
    codeflash_output = _replace_keys(d, "oo", "xx")
    # Now with 'a' to 'x'
    expected = {"foo": {"bxr": 1, "bxz": 2}, "qux": 3}
    codeflash_output = _replace_keys(d, "a", "x")

def test_list_of_dicts():
    # Test that a list of dicts has keys replaced in each dict
    d = [{"foo": 1}, {"bar": 2}]
    expected = [{"fxx": 1}, {"bar": 2}]
    codeflash_output = _replace_keys(d, "oo", "xx")

def test_dict_with_list_values():
    # Test that dicts with list values are handled recursively
    d = {"foo": [ {"bar": 1}, {"baz": 2} ]}
    expected = {"fxx": [ {"bar": 1}, {"baz": 2} ]}
    codeflash_output = _replace_keys(d, "oo", "xx")

def test_list_with_nested_dicts():
    # Test that lists with nested dicts are handled recursively
    d = [ {"foo": {"bar": 1}}, {"baz": 2} ]
    expected = [ {"fxx": {"bar": 1}}, {"baz": 2} ]
    codeflash_output = _replace_keys(d, "oo", "xx")

def test_no_mutation_of_input():
    # Test that the input is not mutated (deepcopy safety)
    d = {"foo": {"bar": 1}}
    d_copy = copy.deepcopy(d)
    codeflash_output = _replace_keys(d, "foo", "baz"); _ = codeflash_output

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

def test_key_is_empty_string():
    # Test that empty string as a key is handled
    d = {"": 1, "foo": 2}
    expected = {"": 1, "foo": 2}
    codeflash_output = _replace_keys(d, "bar", "baz")

def test_match_is_empty_string():
    # If match is empty string, every position is a match, so every key gets target inserted between every character and at start/end
    d = {"ab": 1}
    # For match="", target="x", "ab" becomes "xaxb"
    codeflash_output = _replace_keys(d, "", "x")

def test_target_is_empty_string():
    # If target is empty, match is removed from key
    d = {"foo": 1, "food": 2}
    expected = {"fo": 1, "fod": 2}
    codeflash_output = _replace_keys(d, "o", "")

def test_key_is_only_match():
    # If the key is exactly the match string, it should be replaced with target
    d = {"foo": 1, "bar": 2}
    expected = {"baz": 1, "bar": 2}
    codeflash_output = _replace_keys(d, "foo", "baz")

def test_key_is_multiple_match():
    # If the key contains multiple matches, all should be replaced
    d = {"foofoo": 1}
    expected = {"barbar": 1}
    codeflash_output = _replace_keys(d, "foo", "bar")


def test_value_is_none():
    # Test that None values are preserved
    d = {"foo": None}
    expected = {"fxx": None}
    codeflash_output = _replace_keys(d, "oo", "xx")

def test_value_is_tuple():
    # Test that tuple values are preserved (not recursed into)
    d = {"foo": (1, 2)}
    expected = {"fxx": (1, 2)}
    codeflash_output = _replace_keys(d, "oo", "xx")

def test_nested_empty_structures():
    # Test that nested empty dicts/lists are handled
    d = {"foo": [], "bar": {}}
    expected = {"fxx": [], "bar": {}}
    codeflash_output = _replace_keys(d, "oo", "xx")

def test_deeply_nested():
    # Test that deep nesting works as expected
    d = {"a": [{"b": {"c": "d"}}]}
    expected = {"x": [{"b": {"c": "d"}}]}
    codeflash_output = _replace_keys(d, "a", "x")

def test_non_dict_list_input():
    # Test that non-dict, non-list input is returned as-is
    codeflash_output = _replace_keys(123, "x", "y")
    codeflash_output = _replace_keys("foo", "o", "a")  # not recursed into


def test_list_with_non_dict_elements():
    # Test that lists with non-dict elements are handled
    d = [1, 2, {"foo": 3}]
    expected = [1, 2, {"fxx": 3}]
    codeflash_output = _replace_keys(d, "oo", "xx")

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

def test_large_flat_dict():
    # Test with a large flat dictionary (1000 keys)
    d = {f"key{i}": i for i in range(1000)}
    expected = {f"val{i}": i for i in range(1000)}
    codeflash_output = _replace_keys(d, "key", "val")

def test_large_nested_dict():
    # Test with a large nested dict (10 levels deep, 10 keys at each level)
    def make_deep_dict(depth, width):
        if depth == 0:
            return "leaf"
        return {f"level{depth}_{i}": make_deep_dict(depth-1, width) for i in range(width)}
    d = make_deep_dict(5, 5)
    # Replace "level" with "stage"
    def check_keys_replaced(d):
        if isinstance(d, dict):
            for k in d:
                check_keys_replaced(d[k])
    codeflash_output = _replace_keys(d, "level", "stage"); result = codeflash_output
    check_keys_replaced(result)

def test_large_list_of_dicts():
    # Test with a large list of dicts (1000 elements)
    d = [{"foo": i, "bar": i*2} for i in range(1000)]
    expected = [{"fxx": i, "bar": i*2} for i in range(1000)]
    codeflash_output = _replace_keys(d, "oo", "xx")


def test_performance_large_structure():
    # Test that performance is reasonable on a large nested structure
    # (not a strict performance test, but ensures no stack overflow or excessive slowness)
    d = []
    for i in range(500):
        d.append({"foo": [{"bar": j} for j in range(2)]})
    codeflash_output = _replace_keys(d, "foo", "baz"); result = codeflash_output
    for item in result:
        for subitem in item["baz"]:
            pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from ray.runtime_env.runtime_env import RuntimeEnv
from simulation.run_simulation import _replace_keys

def test__replace_keys():
    _replace_keys(RuntimeEnv(py_modules=[], working_dir=None, pip=None, conda={}, container=None, env_vars=None, worker_process_setup_hook='', config=None, _validate=False, ={'': {}}), '', '')

To edit these changes git checkout codeflash/optimize-_replace_keys-mh9yps0j and push.