⚡️ Speed up method InMemoryNodeState.verify_token by 12%

[codeflash-ai]

📄 12% (0.12x) speedup for InMemoryNodeState.verify_token in framework/py/flwr/supernode/nodestate/in_memory_nodestate.py

⏱️ Runtime : 1.08 milliseconds → 965 microseconds (best of 201 runs)

📝 Explanation and details

The optimization replaces the with statement lock context manager with explicit acquire()/release() calls and moves the comparison operation outside the critical section.

Key changes:

1. Explicit lock management: Uses acquire() and release() instead of with statement to avoid the overhead of context manager protocol
2. Reduced critical section: The dictionary lookup (token_store.get(run_id)) happens inside the lock, but the comparison (value == token) is performed outside the lock

Why it's faster:

• Context manager overhead: The with statement invokes __enter__ and __exit__ methods, adding function call overhead. Direct acquire()/release() eliminates this
• Shorter lock duration: By storing the lookup result and releasing the lock before comparison, other threads can access the lock sooner, improving concurrency
• Hot path optimization: The profiler shows this method is called frequently (3096+ times), making even small per-call improvements significant

Performance characteristics:
The optimization is most effective for workloads with high-frequency token verification calls, as demonstrated by the test cases with large token stores and many verification attempts. The 11% speedup is achieved through reduced per-call overhead in this critical path function.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import pytest  # used for our unit tests
from supernode.nodestate.in_memory_nodestate import InMemoryNodeState

# unit tests

@pytest.fixture
def node_state():
    """Fixture to provide a fresh InMemoryNodeState for each test."""
    return InMemoryNodeState()

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

def test_verify_token_valid(node_state):
    """Test valid token for run_id returns True."""
    node_state.token_store[1] = "abc123"
    codeflash_output = node_state.verify_token(1, "abc123")

def test_verify_token_invalid_token(node_state):
    """Test invalid token for run_id returns False."""
    node_state.token_store[1] = "abc123"
    codeflash_output = node_state.verify_token(1, "wrongtoken")

def test_verify_token_invalid_run_id(node_state):
    """Test valid token but invalid run_id returns False."""
    node_state.token_store[1] = "abc123"
    codeflash_output = node_state.verify_token(2, "abc123")

def test_verify_token_empty_token(node_state):
    """Test empty token string for valid run_id."""
    node_state.token_store[1] = ""
    codeflash_output = node_state.verify_token(1, "")
    codeflash_output = node_state.verify_token(1, "notempty")

def test_verify_token_empty_store(node_state):
    """Test verify_token when token_store is empty."""
    codeflash_output = node_state.verify_token(1, "anytoken")

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

def test_verify_token_none_token(node_state):
    """Test None as token value."""
    node_state.token_store[1] = None
    codeflash_output = node_state.verify_token(1, None)
    codeflash_output = node_state.verify_token(1, "abc123")

def test_verify_token_none_run_id(node_state):
    """Test None as run_id."""
    node_state.token_store[None] = "abc123"
    codeflash_output = node_state.verify_token(None, "abc123")
    codeflash_output = node_state.verify_token(None, "wrongtoken")

def test_verify_token_nonexistent_run_id(node_state):
    """Test run_id not present in token_store."""
    node_state.token_store[1] = "abc123"
    node_state.token_store[2] = "def456"
    codeflash_output = node_state.verify_token(3, "abc123")
    codeflash_output = node_state.verify_token(3, "def456")

def test_verify_token_non_string_token(node_state):
    """Test token values that are not strings."""
    node_state.token_store[1] = 12345  # int token
    codeflash_output = node_state.verify_token(1, 12345)
    codeflash_output = node_state.verify_token(1, "12345")

def test_verify_token_negative_run_id(node_state):
    """Test negative run_id values."""
    node_state.token_store[-1] = "negtoken"
    codeflash_output = node_state.verify_token(-1, "negtoken")
    codeflash_output = node_state.verify_token(-1, "wrongtoken")

def test_verify_token_zero_run_id(node_state):
    """Test run_id of zero."""
    node_state.token_store[0] = "zerotoken"
    codeflash_output = node_state.verify_token(0, "zerotoken")
    codeflash_output = node_state.verify_token(0, "other")

def test_verify_token_special_characters(node_state):
    """Test tokens with special characters."""
    special_token = "!@#$%^&*()_+"
    node_state.token_store[42] = special_token
    codeflash_output = node_state.verify_token(42, "!@#$%^&*()_+")
    codeflash_output = node_state.verify_token(42, "wrong!")

def test_verify_token_unicode_token(node_state):
    """Test tokens with unicode characters."""
    unicode_token = "токен"
    node_state.token_store[7] = unicode_token
    codeflash_output = node_state.verify_token(7, "токен")
    codeflash_output = node_state.verify_token(7, "token")

def test_verify_token_large_integer_token(node_state):
    """Test large integer token values."""
    large_token = 10**18
    node_state.token_store[99] = large_token
    codeflash_output = node_state.verify_token(99, 10**18)
    codeflash_output = node_state.verify_token(99, 10**17)

def test_verify_token_float_token(node_state):
    """Test float token values."""
    node_state.token_store[3] = 1.2345
    codeflash_output = node_state.verify_token(3, 1.2345)
    codeflash_output = node_state.verify_token(3, 1.2346)

def test_verify_token_boolean_token(node_state):
    """Test boolean token values."""
    node_state.token_store[4] = True
    codeflash_output = node_state.verify_token(4, True)
    codeflash_output = node_state.verify_token(4, False)

def test_verify_token_mutation(node_state):
    """Test that changing token_store affects verify_token result."""
    node_state.token_store[5] = "first"
    codeflash_output = node_state.verify_token(5, "first")
    node_state.token_store[5] = "second"
    codeflash_output = node_state.verify_token(5, "first")
    codeflash_output = node_state.verify_token(5, "second")

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

def test_verify_token_many_entries(node_state):
    """Test verify_token with many entries in token_store."""
    # Add 1000 run_id/token pairs
    for i in range(1000):
        node_state.token_store[i] = f"token_{i}"
    # Check a few random valid pairs
    codeflash_output = node_state.verify_token(0, "token_0")
    codeflash_output = node_state.verify_token(999, "token_999")
    codeflash_output = node_state.verify_token(500, "token_500")
    # Check a few invalid tokens
    codeflash_output = node_state.verify_token(100, "wrong")
    codeflash_output = node_state.verify_token(1001, "token_1001")

def test_verify_token_many_collisions(node_state):
    """Test verify_token with many run_ids sharing the same token value."""
    for i in range(1000):
        node_state.token_store[i] = "sharedtoken"
    # All should verify with "sharedtoken"
    for i in range(0, 1000, 100):
        codeflash_output = node_state.verify_token(i, "sharedtoken")
        codeflash_output = node_state.verify_token(i, "wrongtoken")

def test_verify_token_large_token_values(node_state):
    """Test verify_token with very large token strings."""
    large_token = "a" * 1000
    node_state.token_store[123] = large_token
    codeflash_output = node_state.verify_token(123, "a" * 1000)
    codeflash_output = node_state.verify_token(123, "a" * 999)

def test_verify_token_performance(node_state):
    """Test verify_token performance with large token_store."""
    # Add 1000 entries
    for i in range(1000):
        node_state.token_store[i] = f"token_{i}"
    # Ensure lookup is correct for last entry
    codeflash_output = node_state.verify_token(999, "token_999")
    # Ensure lookup is correct for first entry
    codeflash_output = node_state.verify_token(0, "token_0")
    # Ensure lookup fails for missing entry
    codeflash_output = node_state.verify_token(1001, "token_1001")

def test_verify_token_varied_types(node_state):
    """Test verify_token with varied token types in large store."""
    # Mix string, int, float, None, bool
    node_state.token_store[1] = "str"
    node_state.token_store[2] = 123
    node_state.token_store[3] = 1.23
    node_state.token_store[4] = None
    node_state.token_store[5] = True
    codeflash_output = node_state.verify_token(1, "str")
    codeflash_output = node_state.verify_token(2, 123)
    codeflash_output = node_state.verify_token(3, 1.23)
    codeflash_output = node_state.verify_token(4, None)
    codeflash_output = node_state.verify_token(5, True)
    # Wrong types
    codeflash_output = node_state.verify_token(1, 123)
    codeflash_output = node_state.verify_token(2, "123")
    codeflash_output = node_state.verify_token(3, "1.23")
    codeflash_output = node_state.verify_token(4, "None")
    codeflash_output = node_state.verify_token(5, False)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import pytest  # used for our unit tests
from supernode.nodestate.in_memory_nodestate import InMemoryNodeState

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

@pytest.fixture
def node_state():
    """Fixture to provide a fresh InMemoryNodeState for each test."""
    return InMemoryNodeState()

# 1. Basic Test Cases

def test_verify_token_correct_token(node_state):
    """Test that verify_token returns True for correct token."""
    node_state.token_store[1] = "token123"
    codeflash_output = node_state.verify_token(1, "token123")

def test_verify_token_incorrect_token(node_state):
    """Test that verify_token returns False for incorrect token."""
    node_state.token_store[1] = "token123"
    codeflash_output = node_state.verify_token(1, "wrongtoken")

def test_verify_token_missing_run_id(node_state):
    """Test that verify_token returns False if run_id is not in token_store."""
    node_state.token_store[1] = "token123"
    codeflash_output = node_state.verify_token(2, "token123")

def test_verify_token_empty_token(node_state):
    """Test that verify_token returns False if token is empty string."""
    node_state.token_store[1] = "token123"
    codeflash_output = node_state.verify_token(1, "")

def test_verify_token_empty_token_store(node_state):
    """Test that verify_token returns False if token_store is empty."""
    codeflash_output = node_state.verify_token(1, "token123")

# 2. Edge Test Cases

def test_verify_token_none_token(node_state):
    """Test that verify_token returns False if token is None."""
    node_state.token_store[1] = "token123"
    codeflash_output = node_state.verify_token(1, None)

def test_verify_token_none_run_id(node_state):
    """Test that verify_token returns False if run_id is None."""
    node_state.token_store[1] = "token123"
    codeflash_output = node_state.verify_token(None, "token123")

def test_verify_token_token_is_integer(node_state):
    """Test that verify_token returns False if token is an integer."""
    node_state.token_store[1] = "token123"
    codeflash_output = node_state.verify_token(1, 123)

def test_verify_token_run_id_is_string(node_state):
    """Test that verify_token returns False if run_id is a string."""
    node_state.token_store[1] = "token123"
    codeflash_output = node_state.verify_token("1", "token123")

def test_verify_token_token_is_empty_string_in_store(node_state):
    """Test that verify_token returns True if token in store is empty string and matches."""
    node_state.token_store[1] = ""
    codeflash_output = node_state.verify_token(1, "")

def test_verify_token_run_id_negative(node_state):
    """Test that verify_token works for negative run_id."""
    node_state.token_store[-1] = "negtoken"
    codeflash_output = node_state.verify_token(-1, "negtoken")
    codeflash_output = node_state.verify_token(-1, "wrongtoken")

def test_verify_token_run_id_zero(node_state):
    """Test that verify_token works for run_id zero."""
    node_state.token_store[0] = "zerotoken"
    codeflash_output = node_state.verify_token(0, "zerotoken")
    codeflash_output = node_state.verify_token(0, "wrongtoken")

def test_verify_token_token_is_special_characters(node_state):
    """Test that verify_token works for tokens with special characters."""
    special_token = "!@#$%^&*()_+-=[]{}|;:',.<>/?"
    node_state.token_store[42] = special_token
    codeflash_output = node_state.verify_token(42, special_token)
    codeflash_output = node_state.verify_token(42, "wrongtoken")

def test_verify_token_run_id_not_integer_type(node_state):
    """Test that verify_token returns False for non-integer run_id types."""
    node_state.token_store[1] = "token123"
    codeflash_output = node_state.verify_token(1.0, "token123")
    codeflash_output = node_state.verify_token([1], "token123")
    codeflash_output = node_state.verify_token((1,), "token123")

def test_verify_token_token_is_bytes(node_state):
    """Test that verify_token returns False if token is bytes."""
    node_state.token_store[1] = "token123"
    codeflash_output = node_state.verify_token(1, b"token123")

# 3. Large Scale Test Cases

def test_verify_token_large_token_store(node_state):
    """Test verify_token performance and correctness with large token_store."""
    # Fill with 1000 run_id-token pairs
    for i in range(1000):
        node_state.token_store[i] = f"token{i}"
    # Test a few random correct tokens
    codeflash_output = node_state.verify_token(0, "token0")
    codeflash_output = node_state.verify_token(999, "token999")
    # Test incorrect tokens
    codeflash_output = node_state.verify_token(500, "wrongtoken")
    # Test missing run_id
    codeflash_output = node_state.verify_token(1001, "token1001")

def test_verify_token_large_token_values(node_state):
    """Test verify_token with large token string values."""
    large_token = "x" * 1000  # 1000-character token
    node_state.token_store[123] = large_token
    codeflash_output = node_state.verify_token(123, large_token)
    codeflash_output = node_state.verify_token(123, large_token + "y")

def test_verify_token_many_false_positives(node_state):
    """Test verify_token does not return True for similar but incorrect tokens."""
    # Fill with 1000 tokens, all similar
    for i in range(1000):
        node_state.token_store[i] = f"token{i}"
    # Try tokens that are off by one character
    for i in range(1000):
        wrong_token = f"token{i}x"
        codeflash_output = node_state.verify_token(i, wrong_token)

def test_verify_token_all_run_ids(node_state):
    """Test verify_token for all run_ids in a large store."""
    for i in range(1000):
        node_state.token_store[i] = f"token{i}"
    # Check that all correct tokens return True
    for i in range(1000):
        codeflash_output = node_state.verify_token(i, f"token{i}")

def test_verify_token_all_run_ids_incorrect(node_state):
    """Test verify_token for all run_ids with incorrect tokens in a large store."""
    for i in range(1000):
        node_state.token_store[i] = f"token{i}"
    # Check that all incorrect tokens return False
    for i in range(1000):
        codeflash_output = node_state.verify_token(i, f"token{i+1}")
# 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-InMemoryNodeState.verify_token-mh9zuts0 and push.