⚡️ Speed up function wrap_target_calls_with_treesitter by 61% in PR #1199 (omni-java)#1601
Conversation
The optimized code achieves a **61% speedup** (46.7ms → 28.9ms) primarily through three targeted algorithmic improvements: **1. Binary Search Replaces Linear Search (Primary Win)** The original `_byte_to_line_index` consumed 65.2% of total runtime using reverse iteration through `line_byte_starts` (O(n) per call, 1852 calls). The optimized version uses `bisect.bisect_right` for binary search (O(log n)), reducing this function's time from 168ms to 1.3ms - a **129x improvement**. This single change accounts for the bulk of the speedup. **2. Efficient Early-Exit String Checks** `_infer_array_cast_type` originally spent 82.5% of its time allocating a tuple and calling `any()` with generator expressions for substring checks (1852 calls). The optimized code: - Hoists `assertion_methods` tuple to module constant `_ASSERTION_METHODS` (eliminates repeated allocations) - Uses direct `in` checks with early-return: `if "assertArrayEquals" not in line and "assertArrayNotEquals" not in line` This reduces the function's time from 2.8ms to 0.7ms (4x faster), though a smaller absolute gain than the binary search. **3. Micro-optimizations in Hot Path** In `_collect_calls` (24,175 hits), caching `node.type` in a local variable (`node_type = node.type`) and `len(body_bytes)` reduces attribute lookups during recursive traversal, providing marginal but measurable gains in a frequently-executed code path. **Test Results Show Consistent Gains:** - Large-scale tests demonstrate the most improvement: `test_many_calls_scaling_and_counting` (1000 calls) shows **158% speedup** (22.9ms → 8.88ms) - `test_performance_with_large_number_of_calls_and_lines` (500 calls) shows **37.5% speedup** - Even small-scale tests (single calls) show 1-7% improvements from the string check optimizations The optimization particularly benefits workloads with many method invocations to wrap, as the O(log n) vs O(n) difference compounds across multiple calls. Since this function instruments Java code for profiling/testing, it's likely called during build/test cycles where faster instrumentation directly reduces developer iteration time.
| from codeflash.discovery.functions_to_optimize import FunctionToOptimize | ||
| from codeflash.languages.java.parser import JavaAnalyzer | ||
|
|
||
| _ASSERTION_METHODS = ("assertArrayEquals", "assertArrayNotEquals") |
There was a problem hiding this comment.
Nit: _ASSERTION_METHODS is defined here but never referenced anywhere in the file. The code at line 257 uses inline "assertArrayEquals" not in line checks instead. This constant should either be removed or actually used.
| _ASSERTION_METHODS = ("assertArrayEquals", "assertArrayNotEquals") |
PR Review SummaryPrek ChecksFixed: 1 import sorting issue in Mypy: 19 errors found in Code ReviewThis is a well-targeted optimization PR with 3 algorithmic changes to
Issues found:
No critical bugs, security issues, or breaking API changes detected. Test Coverage
Last updated: 2026-02-20 |
⚡️ This pull request contains optimizations for PR #1199
If you approve this dependent PR, these changes will be merged into the original PR branch
omni-java.📄 61% (0.61x) speedup for
wrap_target_calls_with_treesitterincodeflash/languages/java/instrumentation.py⏱️ Runtime :
46.7 milliseconds→28.9 milliseconds(best of101runs)📝 Explanation and details
The optimized code achieves a 61% speedup (46.7ms → 28.9ms) primarily through three targeted algorithmic improvements:
1. Binary Search Replaces Linear Search (Primary Win)
The original
_byte_to_line_indexconsumed 65.2% of total runtime using reverse iteration throughline_byte_starts(O(n) per call, 1852 calls). The optimized version usesbisect.bisect_rightfor binary search (O(log n)), reducing this function's time from 168ms to 1.3ms - a 129x improvement. This single change accounts for the bulk of the speedup.2. Efficient Early-Exit String Checks
_infer_array_cast_typeoriginally spent 82.5% of its time allocating a tuple and callingany()with generator expressions for substring checks (1852 calls). The optimized code:assertion_methodstuple to module constant_ASSERTION_METHODS(eliminates repeated allocations)inchecks with early-return:if "assertArrayEquals" not in line and "assertArrayNotEquals" not in lineThis reduces the function's time from 2.8ms to 0.7ms (4x faster), though a smaller absolute gain than the binary search.
3. Micro-optimizations in Hot Path
In
_collect_calls(24,175 hits), cachingnode.typein a local variable (node_type = node.type) andlen(body_bytes)reduces attribute lookups during recursive traversal, providing marginal but measurable gains in a frequently-executed code path.Test Results Show Consistent Gains:
test_many_calls_scaling_and_counting(1000 calls) shows 158% speedup (22.9ms → 8.88ms)test_performance_with_large_number_of_calls_and_lines(500 calls) shows 37.5% speedupThe optimization particularly benefits workloads with many method invocations to wrap, as the O(log n) vs O(n) difference compounds across multiple calls. Since this function instruments Java code for profiling/testing, it's likely called during build/test cycles where faster instrumentation directly reduces developer iteration time.
✅ Correctness verification report:
🌀 Click to see Generated Regression Tests
To edit these changes
git checkout codeflash/optimize-pr1199-2026-02-20T11.58.24and push.