Merge pull request #1597 from codeflash-ai/codeflash/optimize-pr1199-2026-02-20T10.16.02

⚡️ Speed up method `JavaAssertTransformer._build_target_call` by 1,947% in PR #1199 (`omni-java`)

Author: misrasaurabh1

codeflash/languages/java/parser.py

@@ -7,6 +7,7 @@
 from __future__ import annotations
 
 import logging
+from bisect import bisect_right
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
@@ -111,6 +112,13 @@ def __init__(self) -> None:
         """Initialize the Java analyzer."""
         self._parser: Parser | None = None
 
+        # Caches for the last decoded source to avoid repeated decodes.
+        self._cached_source_bytes: bytes | None = None
+        self._cached_source_str: str | None = None
+        # cumulative byte counts per character: cum_bytes[i] == total bytes for first i characters
+        # length is number_of_chars + 1, cum_bytes[0] == 0
+        self._cached_cum_bytes: list[int] | None = None
+
     @property
     def parser(self) -> Parser:
         """Get the parser, creating it lazily."""
@@ -678,6 +686,41 @@ def get_package_name(self, source: str) -> str | None:
 
         return None
 
+    def _ensure_decoded(self, source: bytes) -> None:
+        """Ensure the provided source bytes are decoded and cumulative byte mapping is built.
+
+        Caches the decoded string and cumulative byte-lengths for the last-seen `source` bytes
+        to make slicing by node byte offsets into string slices much cheaper.
+        """
+        if source is self._cached_source_bytes:
+            return
+
+        decoded = source.decode("utf8")
+        # Build cumulative bytes per character. cum[0] = 0, cum[i] = bytes for first i chars.
+        cum: list[int] = [0]
+        # Building the cumulative mapping is done once per distinct source and is faster than
+        # repeatedly decoding prefixes for many nodes.
+        # A local variable for append and encode reduces attribute lookups.
+        append = cum.append
+        for ch in decoded:
+            append(cum[-1] + len(ch.encode("utf8")))
+
+        self._cached_source_bytes = source
+        self._cached_source_str = decoded
+        self._cached_cum_bytes = cum
+
+    def byte_to_char_index(self, byte_offset: int, source: bytes) -> int:
+        """Convert a byte offset into a character index for the given source bytes.
+
+        This uses a cached cumulative byte-length mapping so repeated conversions are O(log n)
+        (binary search) instead of re-decoding prefixes O(n).
+        """
+        self._ensure_decoded(source)
+        # cum is a non-decreasing list: find largest k where cum[k] <= byte_offset
+        cum = self._cached_cum_bytes  # type: ignore[assignment]
+        # bisect_right returns insertion point; subtract 1 to get character count
+        return bisect_right(cum, byte_offset) - 1
+
 
 def get_java_analyzer() -> JavaAnalyzer:
     """Get a JavaAnalyzer instance.

codeflash/languages/java/remove_asserts.py

@@ -23,6 +23,8 @@
 from codeflash.languages.java.parser import get_java_analyzer
 
 if TYPE_CHECKING:
+    from tree_sitter import Node
+
     from codeflash.discovery.functions_to_optimize import FunctionToOptimize
     from codeflash.languages.java.parser import JavaAnalyzer
 
@@ -548,7 +550,7 @@ def _extract_target_calls(self, content: str, base_offset: int) -> list[TargetCa
 
     def _collect_target_invocations(
         self,
-        node,
+        node: Node,
         wrapper_bytes: bytes,
         content_bytes: bytes,
         base_offset: int,
@@ -601,32 +603,40 @@ def _collect_target_invocations(
             self._collect_target_invocations(child, wrapper_bytes, content_bytes, base_offset, out, seen_top_level)
 
     def _build_target_call(
-        self, node, wrapper_bytes: bytes, content_bytes: bytes, start_byte: int, end_byte: int, base_offset: int
+        self, node: Node, wrapper_bytes: bytes, content_bytes: bytes, start_byte: int, end_byte: int, base_offset: int
     ) -> TargetCall:
         """Build a TargetCall from a tree-sitter method_invocation node."""
-        get_text = self.analyzer.get_node_text
-
         object_node = node.child_by_field_name("object")
         args_node = node.child_by_field_name("arguments")
-        args_text = get_text(args_node, wrapper_bytes) if args_node else ""
+
+        if args_node:
+            args_text = wrapper_bytes[args_node.start_byte : args_node.end_byte].decode("utf8")
+        else:
+            args_text = ""
         # argument_list node includes parens, strip them
-        if args_text.startswith("(") and args_text.endswith(")"):
+        if args_text and args_text[0] == "(" and args_text[-1] == ")":
             args_text = args_text[1:-1]
 
-        # Byte offsets -> char offsets for correct Python string indexing
-        start_char = len(content_bytes[:start_byte].decode("utf8"))
-        end_char = len(content_bytes[:end_byte].decode("utf8"))
+        # Byte offsets -> char offsets for correct Python string indexing using analyzer mapping
+        start_char = self.analyzer.byte_to_char_index(start_byte, content_bytes)
+        end_char = self.analyzer.byte_to_char_index(end_byte, content_bytes)
+
+        # Extract receiver and full call text from the wrapper bytes directly (fast for small wrappers)
+        receiver_text = (
+            wrapper_bytes[object_node.start_byte : object_node.end_byte].decode("utf8") if object_node else None
+        )
+        full_call_text = wrapper_bytes[node.start_byte : node.end_byte].decode("utf8")
 
         return TargetCall(
-            receiver=get_text(object_node, wrapper_bytes) if object_node else None,
+            receiver=receiver_text,
             method_name=self.func_name,
             arguments=args_text,
-            full_call=get_text(node, wrapper_bytes),
+            full_call=full_call_text,
             start_pos=base_offset + start_char,
             end_pos=base_offset + end_char,
         )
 
-    def _find_top_level_arg_node(self, target_node, wrapper_bytes: bytes):
+    def _find_top_level_arg_node(self, target_node: Node, wrapper_bytes: bytes) -> Node | None:
         """Find the top-level argument expression containing a nested target call.
 
         Walks up the AST from target_node to the wrapper _d() call's argument_list.