⚡️ Speed up method ScopedVisitor._is_defined by 4,947%

[codeflash-ai]

📄 4,947% (49.47x) speedup for ScopedVisitor._is_defined in marimo/_ast/visitor.py

⏱️ Runtime : 13.2 milliseconds → 262 microseconds (best of 64 runs)

📝 Explanation and details

The optimized code achieves a 4947% speedup through two key changes that eliminate inefficient generator expressions:

1. Block.is_defined() optimization:

• Original: any(name == defn for defn in self.defs) - O(n) linear search through all definitions
• Optimized: name in self.defs - O(1) set membership test
• This change alone reduces per-hit time from 110μs to 240ns (457x faster)

2. ScopedVisitor._is_defined() optimization:

• Original: any(block.is_defined(identifier) for block in self.block_stack) - Uses generator expression with overhead
• Optimized: Simple for-loop with early exit - Eliminates generator overhead and stops immediately when found
• Reduces per-hit time from 112μs to variable (231ns-2145ns depending on block position)

Why these optimizations work:

• Set membership (in) is fundamentally faster than iterating and comparing each element
• Early exit loops avoid unnecessary work when the identifier is found in early blocks
• Generator expressions have overhead for small datasets that direct loops avoid

Performance characteristics:

• Best case: Identifier found in first block (231ns vs 112μs)
• Worst case: Identifier not found anywhere (still much faster due to set lookups)
• The optimization is particularly effective for the test case with 1000 variable lookups (13.2ms → 249μs), showing it scales excellently with larger datasets

These changes are especially beneficial for AST visitor patterns where variable lookup is performed frequently during code analysis.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import ast
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Callable, Literal, Optional, Union
from uuid import uuid4

# imports
import pytest  # used for our unit tests
from marimo._ast.visitor import ScopedVisitor


# --- Minimal dependencies from marimo/_ast/variables.py ---
def is_local(name: str) -> bool:
    # Dummy implementation for test, not actually used in _is_defined
    return name == "__" or (name.startswith("_") and not name.startswith("__"))

# --- Classes and function under test (from marimo/_ast/visitor.py) ---
Name = str
Language = Literal["python", "sql"]
from marimo._ast.visitor import ScopedVisitor

# --- Unit Tests for _is_defined ---

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

# 1. Basic Test Cases


def test_basic_multiple_names(visitor):
    # Add several names
    visitor.block_stack[0].defs.update({"x", "y", "z"})
    codeflash_output = visitor._is_defined("x") # 2.32μs -> 778ns (198% faster)
    codeflash_output = visitor._is_defined("y") # 943ns -> 295ns (220% faster)
    codeflash_output = visitor._is_defined("z") # 676ns -> 257ns (163% faster)
    codeflash_output = visitor._is_defined("a") # 874ns -> 302ns (189% faster)

def test_basic_multiple_blocks(visitor):
    # Add a new block (scope) and define a name in it
    visitor.block_stack.append(Block(defs={"inner"}))
    codeflash_output = visitor._is_defined("inner")
    # Outer block does not have 'inner'
    visitor.block_stack.pop()
    codeflash_output = visitor._is_defined("inner")

def test_basic_shadowing(visitor):
    # Name defined in both outer and inner block
    visitor.block_stack[0].defs.add("dup")
    visitor.block_stack.append(Block(defs={"dup"}))
    codeflash_output = visitor._is_defined("dup")
    # Remove inner block, still defined in outer
    visitor.block_stack.pop()
    codeflash_output = visitor._is_defined("dup")

# 2. Edge Test Cases


def test_edge_empty_defs(visitor):
    # No names defined in any block
    visitor.block_stack = [Block(), Block()]
    codeflash_output = visitor._is_defined("nothing")

def test_edge_name_case_sensitivity(visitor):
    # Python is case-sensitive
    visitor.block_stack[0].defs.add("FOO")
    codeflash_output = visitor._is_defined("foo") # 2.60μs -> 941ns (176% faster)
    codeflash_output = visitor._is_defined("FOO") # 1.37μs -> 391ns (250% faster)

def test_edge_special_characters(visitor):
    # Names with special characters
    visitor.block_stack[0].defs.update({"$", "_foo", "bar_123", "αβγ"})
    codeflash_output = visitor._is_defined("$") # 2.11μs -> 770ns (174% faster)
    codeflash_output = visitor._is_defined("_foo") # 1.00μs -> 359ns (179% faster)
    codeflash_output = visitor._is_defined("bar_123") # 800ns -> 243ns (229% faster)
    codeflash_output = visitor._is_defined("αβγ") # 652ns -> 221ns (195% faster)
    codeflash_output = visitor._is_defined("not_present") # 918ns -> 355ns (159% faster)

def test_edge_empty_string_name(visitor):
    # Empty string as a name
    visitor.block_stack[0].defs.add("")
    codeflash_output = visitor._is_defined("") # 1.86μs -> 639ns (192% faster)
    codeflash_output = visitor._is_defined(" ") # 920ns -> 335ns (175% faster)

def test_edge_none_and_non_str(visitor):
    # Non-string identifiers should not be found
    visitor.block_stack[0].defs.add("None")
    codeflash_output = visitor._is_defined(None) # 1.79μs -> 669ns (167% faster)
    codeflash_output = visitor._is_defined(123) # 828ns -> 333ns (149% faster)
    codeflash_output = visitor._is_defined("None") # 950ns -> 338ns (181% faster)

def test_edge_duplicate_names_in_blocks(visitor):
    # Add the same name to multiple blocks
    visitor.block_stack[0].defs.add("dup")
    visitor.block_stack.append(Block(defs={"dup"}))
    codeflash_output = visitor._is_defined("dup")

def test_edge_block_with_no_defs(visitor):
    # Block exists but has no defs
    visitor.block_stack.append(Block(defs=set()))
    codeflash_output = visitor._is_defined("foo")

def test_edge_block_with_only_global_names(visitor):
    # Only global_names, not defs
    visitor.block_stack[0].global_names.add("global_var")
    codeflash_output = visitor._is_defined("global_var") # 2.37μs -> 869ns (172% faster)

def test_edge_defs_with_falsy_names(visitor):
    # Falsy names such as "0", "False"
    visitor.block_stack[0].defs.update({"0", "False", "None"})
    codeflash_output = visitor._is_defined("0") # 2.33μs -> 782ns (197% faster)
    codeflash_output = visitor._is_defined("False") # 1.04μs -> 287ns (262% faster)
    codeflash_output = visitor._is_defined("None") # 774ns -> 230ns (237% faster)

# 3. Large Scale Test Cases

def test_large_many_blocks_and_defs():
    # 100 blocks, each with 10 unique names
    v = ScopedVisitor()
    v.block_stack = []
    all_names = set()
    for i in range(100):
        names = {f"name_{i}_{j}" for j in range(10)}
        all_names.update(names)
        v.block_stack.append(Block(defs=names))
    # Should find all names
    for name in all_names:
        codeflash_output = v._is_defined(name)
    # Should not find a name not present
    codeflash_output = v._is_defined("not_defined_anywhere")

def test_large_name_at_last_block():
    # Only the last block has the name
    v = ScopedVisitor()
    v.block_stack = [Block(defs=set()) for _ in range(99)]
    v.block_stack.append(Block(defs={"needle"}))
    codeflash_output = v._is_defined("needle")
    # Remove last block, should not be found
    v.block_stack.pop()
    codeflash_output = v._is_defined("needle")

def test_large_all_blocks_empty():
    # Many blocks, none define anything
    v = ScopedVisitor()
    v.block_stack = [Block(defs=set()) for _ in range(100)]
    codeflash_output = v._is_defined("nothing")

def test_large_performance_many_lookups():
    # 1000 names in one block, check all
    v = ScopedVisitor()
    names = {f"var_{i}" for i in range(1000)}
    v.block_stack[0].defs = names
    for i in range(1000):
        codeflash_output = v._is_defined(f"var_{i}") # 13.2ms -> 249μs (5184% faster)
    # Check a name not present
    codeflash_output = v._is_defined("not_here") # 31.0μs -> 354ns (8650% faster)

def test_large_shadowing_across_blocks():
    # Name appears in several blocks, should always return True
    v = ScopedVisitor()
    for i in range(10):
        v.block_stack.append(Block(defs={f"shadowed", f"unique_{i}"}))
    codeflash_output = v._is_defined("shadowed")
    for i in range(10):
        codeflash_output = v._is_defined(f"unique_{i}")
    codeflash_output = v._is_defined("not_present")
# 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 marimo._ast.visitor import ScopedVisitor

# --- Unit tests for ScopedVisitor._is_defined ---

# Basic Test Cases




















#------------------------------------------------
from marimo._ast.visitor import ScopedVisitor

def test_ScopedVisitor__is_defined():
    ScopedVisitor._is_defined(ScopedVisitor(), '')

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_o_lbxivc/tmpmjy6qswu/test_concolic_coverage.py::test_ScopedVisitor__is_defined	3.17μs	1.44μs	120%✅

To edit these changes git checkout codeflash/optimize-ScopedVisitor._is_defined-mhcw6q6h and push.