fix: guard Java replacement against wrong-method-name candidates and anonymous-class method hoisting
Two bugs in _parse_optimization_source (replacement.py) caused Maven compilation
failures when codeflash optimised aerospike-client-java:
Bug 1 – standalone method with wrong name replaces target
When the LLM generated a standalone method whose name did not match the
optimisation target (e.g. generated `unpackMap` for target `unpackObjectMap`,
or generated `sizeTxn` for target `estimateKeySize`), the function fell back to
using the entire generated snippet as `target_method_source`. This silently
replaced the target with the wrong method, producing:
• a duplicate definition of the wrong method
• removal of the target method (breaking all callers)
Fix: after parsing standalone (class-free) code, verify that at least one
discovered method matches the target name. If no match is found, set
`target_method_source` to the empty string and log a warning. A corresponding
guard in `replace_function` returns the original source unchanged when
`target_method_source` is empty.
The same guard is applied to the full-class path: if the generated class does
not contain the target method, the candidate is also rejected.
Bug 2 – anonymous inner-class methods hoisted as top-level helpers
When an optimised method returned an anonymous class (e.g. `keySetIterator`
returning `new Iterator<LuaValue>() { … }`), tree-sitter's recursive walk
found the anonymous class's `hasNext`, `next`, and `remove` method_declaration
nodes and classified them as helpers to be inserted at the outer-class level.
The inserted methods carried `@Override` annotations that matched nothing in the
outer class and referenced local variables (`it`) that were only in scope inside
the optimised method body, producing compilation errors.
Fix: when extracting helpers from the optimised class, skip any method whose
line range is entirely contained within the target method's line range. Such
methods belong to anonymous/nested classes inside the method body and must not
be hoisted out as standalone class members.
Tests added for both bugs in TestWrongMethodNameGeneration and
TestAnonymousInnerClassMethods.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
misrasaurabh1
parent
24e26c3944
commit
2371540386
2 changed files with 298 additions and 3 deletions
|
|
@ -51,10 +51,11 @@ def _parse_optimization_source(new_source: str, target_method_name: str, analyze
|
|||
|
||||
Returns:
|
||||
ParsedOptimization with the method and any additional members.
|
||||
If the generated code contains no method matching target_method_name,
|
||||
target_method_source will be empty to signal that the candidate is invalid.
|
||||
|
||||
"""
|
||||
new_fields: list[str] = []
|
||||
new_helper_methods: list[str] = []
|
||||
target_method_source = new_source # Default to the whole source
|
||||
|
||||
# Check if this is a full class or just a method
|
||||
|
|
@ -83,11 +84,28 @@ def _parse_optimization_source(new_source: str, target_method_name: str, analyze
|
|||
start = (target_method.javadoc_start_line or target_method.start_line) - 1
|
||||
end = target_method.end_line
|
||||
target_method_source = "".join(lines[start:end])
|
||||
else:
|
||||
# Target method not found in the generated class — the LLM generated
|
||||
# a different method. Signal invalid candidate with empty source.
|
||||
logger.warning(
|
||||
"Generated class does not contain target method '%s'. Skipping candidate.",
|
||||
target_method_name,
|
||||
)
|
||||
target_method_source = ""
|
||||
|
||||
# Extract helper methods, categorised by position relative to the target
|
||||
# Extract helper methods, categorised by position relative to the target.
|
||||
# Skip methods whose line range falls entirely inside the target method's
|
||||
# range, as these belong to anonymous/inner classes inside the target body
|
||||
# and must not be hoisted out as top-level class members.
|
||||
lines = new_source.splitlines(keepends=True)
|
||||
for i, method in enumerate(methods):
|
||||
if method.name != target_method_name:
|
||||
lines = new_source.splitlines(keepends=True)
|
||||
# Skip methods nested inside the target (e.g. anonymous class methods)
|
||||
if target_method and (
|
||||
method.start_line >= target_method.start_line
|
||||
and method.end_line <= target_method.end_line
|
||||
):
|
||||
continue
|
||||
start = (method.javadoc_start_line or method.start_line) - 1
|
||||
end = method.end_line
|
||||
helper_source = "".join(lines[start:end])
|
||||
|
|
@ -101,6 +119,22 @@ def _parse_optimization_source(new_source: str, target_method_name: str, analyze
|
|||
if f.source_text:
|
||||
new_fields.append(f.source_text)
|
||||
|
||||
else:
|
||||
# No class found — generated code is a standalone method (or snippet).
|
||||
# Validate that it actually defines the target method; if it defines a
|
||||
# *different* method, applying it would corrupt the original source.
|
||||
standalone_methods = analyzer.find_methods(new_source)
|
||||
if standalone_methods:
|
||||
matching = [m for m in standalone_methods if m.name == target_method_name]
|
||||
if not matching:
|
||||
logger.warning(
|
||||
"Generated standalone method '%s' does not match target method '%s'. "
|
||||
"Skipping candidate to avoid corrupting the source.",
|
||||
standalone_methods[0].name,
|
||||
target_method_name,
|
||||
)
|
||||
target_method_source = ""
|
||||
|
||||
return ParsedOptimization(
|
||||
target_method_source=target_method_source,
|
||||
new_fields=new_fields,
|
||||
|
|
@ -276,6 +310,12 @@ def replace_function(
|
|||
# Parse the optimization to extract components
|
||||
parsed = _parse_optimization_source(new_source, func_name, analyzer)
|
||||
|
||||
# If the parsed optimization has no valid target source (e.g., the LLM generated
|
||||
# a method with a different name), skip this candidate entirely.
|
||||
if not parsed.target_method_source.strip():
|
||||
logger.warning("No valid replacement found for method '%s'. Returning original source.", func_name)
|
||||
return source
|
||||
|
||||
# Find the method in the original source
|
||||
methods = analyzer.find_methods(source)
|
||||
target_method = None
|
||||
|
|
|
|||
|
|
@ -1654,3 +1654,258 @@ public final class Buffer {{
|
|||
}
|
||||
"""
|
||||
assert new_code == expected
|
||||
|
||||
|
||||
class TestWrongMethodNameGeneration:
|
||||
"""Tests that guard against the LLM generating a different method name than the target.
|
||||
|
||||
When the optimizer generates code for method X but the LLM produces method Y instead,
|
||||
applying that replacement would:
|
||||
- Replace method X with the body of method Y (creating a duplicate of Y).
|
||||
- Remove method X from the source.
|
||||
|
||||
These tests verify that codeflash detects this mismatch and leaves the original
|
||||
source file unchanged.
|
||||
"""
|
||||
|
||||
def test_standalone_wrong_method_name_leaves_source_unchanged(self, tmp_path):
|
||||
"""Standalone generated method with wrong name must not replace the target.
|
||||
|
||||
Reproduces the Unpacker.unpackObjectMap bug: the LLM was asked to optimise
|
||||
``unpackObjectMap`` but generated ``unpackMap`` as a standalone method.
|
||||
Applying that would create a duplicate ``unpackMap`` and delete
|
||||
``unpackObjectMap``, causing compilation failures.
|
||||
"""
|
||||
from codeflash.discovery.functions_to_optimize import FunctionToOptimize
|
||||
|
||||
java_file = tmp_path / "Unpacker.java"
|
||||
original_code = """\
|
||||
public abstract class Unpacker {
|
||||
public static Object unpackObjectMap(byte[] buffer, int offset, int length) {
|
||||
return new Object();
|
||||
}
|
||||
|
||||
public Object unpackMap() {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
"""
|
||||
java_file.write_text(original_code, encoding="utf-8")
|
||||
|
||||
# The LLM generated an optimised ``unpackMap`` when it should have
|
||||
# optimised ``unpackObjectMap``.
|
||||
optimized_markdown = f"""```java:{java_file.relative_to(tmp_path)}
|
||||
public final Object unpackMap() {{
|
||||
return new Object();
|
||||
}}
|
||||
```"""
|
||||
|
||||
optimized_code = CodeStringsMarkdown.parse_markdown_code(optimized_markdown, expected_language="java")
|
||||
|
||||
function_to_optimize = FunctionToOptimize(
|
||||
function_name="unpackObjectMap",
|
||||
file_path=java_file,
|
||||
starting_line=2,
|
||||
ending_line=4,
|
||||
parents=[FunctionParent(name="Unpacker", type="ClassDef")],
|
||||
qualified_name="Unpacker.unpackObjectMap",
|
||||
is_method=True,
|
||||
)
|
||||
|
||||
result = replace_function_definitions_for_language(
|
||||
function_names=["unpackObjectMap"],
|
||||
optimized_code=optimized_code,
|
||||
module_abspath=java_file,
|
||||
project_root_path=tmp_path,
|
||||
function_to_optimize=function_to_optimize,
|
||||
)
|
||||
|
||||
# No modification should occur — wrong method name in generated code.
|
||||
assert result is False
|
||||
assert java_file.read_text(encoding="utf-8") == original_code
|
||||
|
||||
def test_class_wrapper_with_wrong_target_method_leaves_source_unchanged(self, tmp_path):
|
||||
"""Class-wrapped generated code missing the target method must not modify source.
|
||||
|
||||
Reproduces the Command.estimateKeySize bug: the LLM generated a class that
|
||||
contained only ``sizeTxn`` (a helper) and did not include ``estimateKeySize``
|
||||
(the target). Applying it would duplicate ``sizeTxn`` in the source.
|
||||
"""
|
||||
from codeflash.discovery.functions_to_optimize import FunctionToOptimize
|
||||
|
||||
java_file = tmp_path / "Command.java"
|
||||
original_code = """\
|
||||
public class Command {
|
||||
public int estimateKeySize(String key) {
|
||||
return key.length() + 4;
|
||||
}
|
||||
|
||||
private int sizeTxn(String key, Object txn, boolean hasWrite) {
|
||||
return key.length();
|
||||
}
|
||||
}
|
||||
"""
|
||||
java_file.write_text(original_code, encoding="utf-8")
|
||||
|
||||
# The LLM generated a class containing only ``sizeTxn`` instead of
|
||||
# the target ``estimateKeySize``.
|
||||
optimized_markdown = f"""```java:{java_file.relative_to(tmp_path)}
|
||||
public class Command {{
|
||||
private int sizeTxn(String key, Object txn, boolean hasWrite) {{
|
||||
return key.length() + 1;
|
||||
}}
|
||||
}}
|
||||
```"""
|
||||
|
||||
optimized_code = CodeStringsMarkdown.parse_markdown_code(optimized_markdown, expected_language="java")
|
||||
|
||||
function_to_optimize = FunctionToOptimize(
|
||||
function_name="estimateKeySize",
|
||||
file_path=java_file,
|
||||
starting_line=2,
|
||||
ending_line=4,
|
||||
parents=[FunctionParent(name="Command", type="ClassDef")],
|
||||
qualified_name="Command.estimateKeySize",
|
||||
is_method=True,
|
||||
)
|
||||
|
||||
result = replace_function_definitions_for_language(
|
||||
function_names=["estimateKeySize"],
|
||||
optimized_code=optimized_code,
|
||||
module_abspath=java_file,
|
||||
project_root_path=tmp_path,
|
||||
function_to_optimize=function_to_optimize,
|
||||
)
|
||||
|
||||
# No modification should occur — target method absent from generated class.
|
||||
assert result is False
|
||||
assert java_file.read_text(encoding="utf-8") == original_code
|
||||
|
||||
|
||||
class TestAnonymousInnerClassMethods:
|
||||
"""Tests that methods inside anonymous inner classes are not hoisted as helpers.
|
||||
|
||||
When an optimised method uses an anonymous class (e.g. an inline Iterator),
|
||||
the anonymous class's own methods (hasNext, next, remove ...) must NOT be
|
||||
extracted and inserted as top-level class members. Doing so would create
|
||||
broken methods: they would carry @Override annotations that do not correspond
|
||||
to any supertype method, and would reference variables only available in the
|
||||
enclosing method scope.
|
||||
"""
|
||||
|
||||
def test_anonymous_iterator_methods_not_hoisted_to_class(self, tmp_path):
|
||||
"""Reproduces the LuaMap.keySetIterator bug.
|
||||
|
||||
The LLM optimised ``keySetIterator`` by returning an anonymous
|
||||
``Iterator`` with ``hasNext``, ``next``, and ``remove`` methods.
|
||||
Those three methods must remain inside the anonymous class body and
|
||||
must NOT be added as top-level members of the outer class.
|
||||
"""
|
||||
from codeflash.discovery.functions_to_optimize import FunctionToOptimize
|
||||
|
||||
java_file = tmp_path / "LuaMap.java"
|
||||
original_code = """\
|
||||
import java.util.Iterator;
|
||||
import java.util.Map;
|
||||
|
||||
public final class LuaMap {
|
||||
private final Map<String, String> map;
|
||||
|
||||
public LuaMap(Map<String, String> map) {
|
||||
this.map = map;
|
||||
}
|
||||
|
||||
public Iterator<String> keySetIterator() {
|
||||
return map.keySet().iterator();
|
||||
}
|
||||
|
||||
public int size() {
|
||||
return map.size();
|
||||
}
|
||||
}
|
||||
"""
|
||||
java_file.write_text(original_code, encoding="utf-8")
|
||||
|
||||
# Optimised version returns a custom anonymous Iterator that avoids
|
||||
# creating a keySet view for empty maps.
|
||||
optimized_markdown = f"""```java:{java_file.relative_to(tmp_path)}
|
||||
import java.util.Iterator;
|
||||
import java.util.Map;
|
||||
|
||||
public final class LuaMap {{
|
||||
private final Map<String, String> map;
|
||||
|
||||
public LuaMap(Map<String, String> map) {{
|
||||
this.map = map;
|
||||
}}
|
||||
|
||||
public Iterator<String> keySetIterator() {{
|
||||
if (map.isEmpty()) {{
|
||||
return java.util.Collections.emptyIterator();
|
||||
}}
|
||||
final Iterator<Map.Entry<String, String>> it = map.entrySet().iterator();
|
||||
return new Iterator<String>() {{
|
||||
@Override
|
||||
public boolean hasNext() {{
|
||||
return it.hasNext();
|
||||
}}
|
||||
@Override
|
||||
public String next() {{
|
||||
return it.next().getKey();
|
||||
}}
|
||||
@Override
|
||||
public void remove() {{
|
||||
it.remove();
|
||||
}}
|
||||
}};
|
||||
}}
|
||||
|
||||
public int size() {{
|
||||
return map.size();
|
||||
}}
|
||||
}}
|
||||
```"""
|
||||
|
||||
optimized_code = CodeStringsMarkdown.parse_markdown_code(optimized_markdown, expected_language="java")
|
||||
|
||||
function_to_optimize = FunctionToOptimize(
|
||||
function_name="keySetIterator",
|
||||
file_path=java_file,
|
||||
starting_line=11,
|
||||
ending_line=13,
|
||||
parents=[FunctionParent(name="LuaMap", type="ClassDef")],
|
||||
qualified_name="LuaMap.keySetIterator",
|
||||
is_method=True,
|
||||
)
|
||||
|
||||
result = replace_function_definitions_for_language(
|
||||
function_names=["keySetIterator"],
|
||||
optimized_code=optimized_code,
|
||||
module_abspath=java_file,
|
||||
project_root_path=tmp_path,
|
||||
function_to_optimize=function_to_optimize,
|
||||
)
|
||||
|
||||
assert result is True
|
||||
new_code = java_file.read_text(encoding="utf-8")
|
||||
|
||||
# The replaced method should contain the anonymous iterator.
|
||||
assert "emptyIterator()" in new_code
|
||||
assert "it.hasNext()" in new_code
|
||||
|
||||
# Crucially, hasNext/next/remove must NOT appear as standalone top-level
|
||||
# class members outside the keySetIterator method body.
|
||||
lines = new_code.splitlines()
|
||||
size_idx = next(i for i, ln in enumerate(lines) if "public int size()" in ln)
|
||||
|
||||
# No anonymous-class methods should appear after size() definition
|
||||
for i in range(size_idx, len(lines)):
|
||||
assert "public boolean hasNext()" not in lines[i], (
|
||||
f"hasNext() was hoisted as a top-level class member at line {i + 1}"
|
||||
)
|
||||
assert "public String next()" not in lines[i], (
|
||||
f"next() was hoisted as a top-level class member at line {i + 1}"
|
||||
)
|
||||
assert "public void remove()" not in lines[i], (
|
||||
f"remove() was hoisted as a top-level class member at line {i + 1}"
|
||||
)
|
||||
|
|
|
|||
Loading…
Reference in a new issue