Merge pull request #1814 from codeflash-ai/fix/js-replacement-stale-line-numbers-after-global-declarations
fix: re-discover function position after add_global_declarations shifts line numbers
This commit is contained in:
mohammed ahmed
GitHub
commit
1e92f3d2ed
2 changed files with 232 additions and 4 deletions
|
|
@ -89,19 +89,24 @@ def replace_function_definitions_for_language(
|
|||
and function_to_optimize.ending_line
|
||||
and function_to_optimize.file_path == module_abspath
|
||||
):
|
||||
# Re-discover the function in the (possibly modified) source to get updated line numbers.
|
||||
# add_global_declarations may have inserted new declarations above the function,
|
||||
# shifting its line numbers from the original function_to_optimize values.
|
||||
current_func = _rediscover_function(lang_support, original_source_code, module_abspath, function_to_optimize)
|
||||
|
||||
# For Java, we need to pass the full optimized code so replace_function can
|
||||
# extract and add any new class members (static fields, helper methods).
|
||||
# For other languages, we extract just the target function.
|
||||
if language == Language.JAVA:
|
||||
new_code = lang_support.replace_function(original_source_code, function_to_optimize, code_to_apply)
|
||||
new_code = lang_support.replace_function(original_source_code, current_func, code_to_apply)
|
||||
else:
|
||||
optimized_func = _extract_function_from_code(
|
||||
lang_support, code_to_apply, function_to_optimize.function_name, module_abspath
|
||||
lang_support, code_to_apply, current_func.function_name, module_abspath
|
||||
)
|
||||
if optimized_func:
|
||||
new_code = lang_support.replace_function(original_source_code, function_to_optimize, optimized_func)
|
||||
new_code = lang_support.replace_function(original_source_code, current_func, optimized_func)
|
||||
else:
|
||||
new_code = lang_support.replace_function(original_source_code, function_to_optimize, code_to_apply)
|
||||
new_code = lang_support.replace_function(original_source_code, current_func, code_to_apply)
|
||||
else:
|
||||
new_code = original_source_code
|
||||
modified = False
|
||||
|
|
@ -144,6 +149,33 @@ def replace_function_definitions_for_language(
|
|||
return True
|
||||
|
||||
|
||||
def _rediscover_function(
|
||||
lang_support: LanguageSupport, source_code: str, file_path: Path, function_to_optimize: FunctionToOptimize
|
||||
) -> FunctionToOptimize:
|
||||
"""Re-discover a function in source code to get updated line numbers.
|
||||
|
||||
After add_global_declarations inserts new declarations, the target function's
|
||||
line numbers shift. This re-discovers the function by name to get current positions.
|
||||
|
||||
Returns the original function_to_optimize if rediscovery fails.
|
||||
"""
|
||||
try:
|
||||
current_functions = lang_support.discover_functions(source_code, file_path, _SOURCE_CRITERIA)
|
||||
target_name = function_to_optimize.function_name
|
||||
original_start = function_to_optimize.starting_line or 0
|
||||
candidates = [func for func in current_functions if func.function_name == target_name]
|
||||
if len(candidates) == 1:
|
||||
return candidates[0]
|
||||
if candidates:
|
||||
# For overloaded methods, pick the one closest to the original line position.
|
||||
# Global declarations shift lines down, so the new start >= original start.
|
||||
return min(candidates, key=lambda f: abs((f.starting_line or 0) - original_start))
|
||||
except Exception as e:
|
||||
logger.debug(f"Error rediscovering function {function_to_optimize.function_name}: {e}")
|
||||
|
||||
return function_to_optimize
|
||||
|
||||
|
||||
def _extract_function_from_code(
|
||||
lang_support: LanguageSupport, source_code: str, function_name: str, file_path: Path
|
||||
) -> str | None:
|
||||
|
|
|
|||
|
|
@ -2411,3 +2411,199 @@ export const fetchData = async (url) => {
|
|||
|
||||
assert "return (await fetch(url)).json();" in result
|
||||
assert js_support.validate_syntax(result) is True
|
||||
|
||||
|
||||
class TestNewGlobalDeclarationsShiftArrowFunctionLines:
|
||||
"""Tests for the bug where add_global_declarations shifts arrow function line numbers.
|
||||
|
||||
When optimized code introduces new global declarations (const, let, etc.) that get
|
||||
added to the original source via add_global_declarations, the target arrow function's
|
||||
line number shifts. But replace_function still uses the old starting_line, so
|
||||
_replace_function_body fails to find the function.
|
||||
|
||||
Reproduces: bytesToHumanReadable, isSelectable, and similar failures from Strapi logs.
|
||||
"""
|
||||
|
||||
def test_arrow_function_replacement_after_new_global_const(self, ts_support, temp_project):
|
||||
"""Test replacing an arrow function when optimized code adds new global constants.
|
||||
|
||||
The optimized code adds new module-level constants (SIZES, LOG_1000) that get
|
||||
inserted before the arrow function via add_global_declarations, shifting its line.
|
||||
The replacement must still find and update the function body.
|
||||
|
||||
Uses function_to_optimize (the if-branch in replace_function_definitions_for_language)
|
||||
to reproduce the real pipeline where stale starting_line causes failures.
|
||||
"""
|
||||
from codeflash.discovery.functions_to_optimize import FunctionToOptimize
|
||||
from codeflash.models.models import CodeString, CodeStringsMarkdown
|
||||
|
||||
original_source = """\
|
||||
import { Writable, WritableOptions } from 'node:stream';
|
||||
|
||||
const bytesToHumanReadable = (bytes: number) => {
|
||||
if (bytes === 0) return '0 Bytes';
|
||||
const sizes = ['Bytes', 'KB', 'MB', 'GB', 'TB', 'PB'];
|
||||
const i = Math.floor(Math.log(bytes) / Math.log(1000));
|
||||
return `${Math.round(bytes / Math.pow(1000, i))} ${sizes[i]}`;
|
||||
};
|
||||
|
||||
export { bytesToHumanReadable };
|
||||
"""
|
||||
file_path = temp_project / "file.ts"
|
||||
file_path.write_text(original_source, encoding="utf-8")
|
||||
|
||||
# Discover the function to get real starting_line/ending_line
|
||||
source = file_path.read_text(encoding="utf-8")
|
||||
functions = ts_support.discover_functions(
|
||||
source, file_path, FunctionFilterCriteria(require_return=False, require_export=False)
|
||||
)
|
||||
target = next(f for f in functions if f.function_name == "bytesToHumanReadable")
|
||||
|
||||
func_to_optimize = FunctionToOptimize(
|
||||
function_name=target.function_name,
|
||||
file_path=file_path,
|
||||
parents=target.parents,
|
||||
starting_line=target.starting_line,
|
||||
ending_line=target.ending_line,
|
||||
starting_col=target.starting_col,
|
||||
ending_col=target.ending_col,
|
||||
is_async=target.is_async,
|
||||
is_method=target.is_method,
|
||||
language=target.language,
|
||||
)
|
||||
|
||||
# Optimized code hoists constant arrays and precomputes LOG_1000
|
||||
optimized_code = """\
|
||||
const SIZES = ['Bytes', 'KB', 'MB', 'GB', 'TB', 'PB'];
|
||||
const LOG_1000 = Math.log(1000);
|
||||
|
||||
const bytesToHumanReadable = (bytes: number) => {
|
||||
if (bytes === 0) return '0 Bytes';
|
||||
const i = Math.floor(Math.log(bytes) / LOG_1000);
|
||||
return `${Math.round(bytes / 1000 ** i)} ${SIZES[i]}`;
|
||||
};
|
||||
"""
|
||||
|
||||
code_markdown = CodeStringsMarkdown(
|
||||
code_strings=[CodeString(code=optimized_code, file_path=Path("file.ts"), language="typescript")],
|
||||
language="typescript",
|
||||
)
|
||||
|
||||
replaced = replace_function_definitions_for_language(
|
||||
["bytesToHumanReadable"],
|
||||
code_markdown,
|
||||
file_path,
|
||||
temp_project,
|
||||
lang_support=ts_support,
|
||||
function_to_optimize=func_to_optimize,
|
||||
)
|
||||
|
||||
assert replaced, "Expected code to be replaced"
|
||||
result = file_path.read_text(encoding="utf-8")
|
||||
|
||||
expected = """\
|
||||
import { Writable, WritableOptions } from 'node:stream';
|
||||
|
||||
const LOG_1000 = Math.log(1000);
|
||||
|
||||
const SIZES = ['Bytes', 'KB', 'MB', 'GB', 'TB', 'PB'];
|
||||
|
||||
const bytesToHumanReadable = (bytes: number) => {
|
||||
if (bytes === 0) return '0 Bytes';
|
||||
const i = Math.floor(Math.log(bytes) / LOG_1000);
|
||||
return `${Math.round(bytes / 1000 ** i)} ${SIZES[i]}`;
|
||||
};
|
||||
|
||||
export { bytesToHumanReadable };
|
||||
"""
|
||||
assert result == expected, f"Result does not match expected output.\nExpected:\n{expected}\n\nGot:\n{result}"
|
||||
|
||||
def test_arrow_function_with_new_set_declaration(self, ts_support, temp_project):
|
||||
"""Test replacing arrow function when optimized code adds a Set declaration.
|
||||
|
||||
Reproduces the isSelectable failure pattern from Strapi logs where optimized code
|
||||
introduces a new Set (NON_FILE_TYPES/MEDIA_TYPES) above the arrow function.
|
||||
Uses function_to_optimize to trigger the stale-line-number bug.
|
||||
"""
|
||||
from codeflash.discovery.functions_to_optimize import FunctionToOptimize
|
||||
from codeflash.models.models import CodeString, CodeStringsMarkdown
|
||||
|
||||
original_source = """\
|
||||
export const isSelectable = (allowedTypes: string[], mime = '') => {
|
||||
if (!mime) return false;
|
||||
const fileType = mime.split('/')[0];
|
||||
if (allowedTypes.includes(fileType)) return true;
|
||||
if (allowedTypes.includes('file') && !['video', 'image', 'audio'].includes(fileType)) return true;
|
||||
return false;
|
||||
};
|
||||
"""
|
||||
file_path = temp_project / "isSelectable.ts"
|
||||
file_path.write_text(original_source, encoding="utf-8")
|
||||
|
||||
# Discover the function to get real starting_line/ending_line
|
||||
source = file_path.read_text(encoding="utf-8")
|
||||
functions = ts_support.discover_functions(source, file_path)
|
||||
target = next(f for f in functions if f.function_name == "isSelectable")
|
||||
|
||||
func_to_optimize = FunctionToOptimize(
|
||||
function_name=target.function_name,
|
||||
file_path=file_path,
|
||||
parents=target.parents,
|
||||
starting_line=target.starting_line,
|
||||
ending_line=target.ending_line,
|
||||
starting_col=target.starting_col,
|
||||
ending_col=target.ending_col,
|
||||
is_async=target.is_async,
|
||||
is_method=target.is_method,
|
||||
language=target.language,
|
||||
)
|
||||
|
||||
# Optimized code extracts the type set into a module-level constant
|
||||
optimized_code = """\
|
||||
const NON_FILE_TYPES = new Set(['video', 'image', 'audio']);
|
||||
|
||||
export const isSelectable = (allowedTypes: string[], mime = '') => {
|
||||
if (!mime) return false;
|
||||
const slashIndex = mime.indexOf('/');
|
||||
const fileType = slashIndex === -1 ? mime : mime.substring(0, slashIndex);
|
||||
for (let i = 0, len = allowedTypes.length; i < len; i++) {
|
||||
const t = allowedTypes[i];
|
||||
if (t === fileType) return true;
|
||||
if (t === 'file' && !NON_FILE_TYPES.has(fileType)) return true;
|
||||
}
|
||||
return false;
|
||||
};
|
||||
"""
|
||||
|
||||
code_markdown = CodeStringsMarkdown(
|
||||
code_strings=[CodeString(code=optimized_code, file_path=Path("isSelectable.ts"), language="typescript")],
|
||||
language="typescript",
|
||||
)
|
||||
|
||||
replaced = replace_function_definitions_for_language(
|
||||
["isSelectable"],
|
||||
code_markdown,
|
||||
file_path,
|
||||
temp_project,
|
||||
lang_support=ts_support,
|
||||
function_to_optimize=func_to_optimize,
|
||||
)
|
||||
|
||||
assert replaced, "Expected code to be replaced"
|
||||
result = file_path.read_text(encoding="utf-8")
|
||||
|
||||
expected = """\
|
||||
const NON_FILE_TYPES = new Set(['video', 'image', 'audio']);
|
||||
export const isSelectable = (allowedTypes: string[], mime = '') => {
|
||||
if (!mime) return false;
|
||||
const slashIndex = mime.indexOf('/');
|
||||
const fileType = slashIndex === -1 ? mime : mime.substring(0, slashIndex);
|
||||
for (let i = 0, len = allowedTypes.length; i < len; i++) {
|
||||
const t = allowedTypes[i];
|
||||
if (t === fileType) return true;
|
||||
if (t === 'file' && !NON_FILE_TYPES.has(fileType)) return true;
|
||||
}
|
||||
return false;
|
||||
};
|
||||
"""
|
||||
assert result == expected, f"Result does not match expected output.\nExpected:\n{expected}\n\nGot:\n{result}"
|
||||
|
|
|
|||
Loading…
Reference in a new issue