diff --git a/reai_toolkit/app/coordinators/ai_decomp_coordinator.py b/reai_toolkit/app/coordinators/ai_decomp_coordinator.py index 88a7278..733428e 100644 --- a/reai_toolkit/app/coordinators/ai_decomp_coordinator.py +++ b/reai_toolkit/app/coordinators/ai_decomp_coordinator.py @@ -17,6 +17,7 @@ RENAME_NOT_READY, RenderModel, display_rows_for_source_lines, + own_function_placeholder, render_progress, render_stream, render_view_with_map, @@ -339,19 +340,36 @@ def apply_predicted_name(self, name: str) -> None: return if self.ai_decomp_service.update_function_name(ea, name): - self.ai_decomp_service.tag_function_as_renamed(name) - self._hide_predicted_name() - self.refresh_disassembly_view() + self._after_function_renamed(ea, name) return final = self.ai_decomp_service.apply_deduped_name(ea, name) if final is None: self.show_info_dialog(msg=f"Could not rename this function to '{name}'.") return - self.ai_decomp_service.tag_function_as_renamed(final) + self._after_function_renamed(ea, final) + + def _after_function_renamed(self, ea: int, name: str) -> None: + self.ai_decomp_service.tag_function_as_renamed(name) self._hide_predicted_name() self.refresh_disassembly_view() + placeholder = ( + own_function_placeholder(self._current_tokenised) + if self._current_tokenised is not None + else None + ) + if placeholder is None: + self.refresh_current() + return + + self.ai_decomp_service.apply_overrides( + ea=ea, + overrides={placeholder: name}, + on_decomp=lambda response: self._on_decomp_complete(ea, response), + on_tokenised=lambda response: self._on_tokenised_complete(ea, response), + ) + def request_rename(self, display_line: int, word: str) -> None: ea = self._current_func_vaddr if ea is None: diff --git a/reai_toolkit/app/coordinators/ai_decomp_render.py b/reai_toolkit/app/coordinators/ai_decomp_render.py index d5fd48f..55fabc6 100644 --- a/reai_toolkit/app/coordinators/ai_decomp_render.py +++ b/reai_toolkit/app/coordinators/ai_decomp_render.py @@ -20,6 +20,8 @@ _IDENT_RE = re.compile(r"[A-Za-z_]\w*") +OWN_FUNCTION_KIND = "own_function" + RENAME_NOT_READY = "The AI decompilation is still loading — try again once it finishes." RENAME_NOT_DECOMP_LINE = "This line is not part of the AI decompilation." RENAME_NOT_CODE_LINE = ( @@ -251,6 +253,14 @@ def effective_values(tokens: "GetTokensResponse") -> dict[str, str]: return values +def own_function_placeholder(tokens: "GetTokensResponse") -> Optional[str]: + rendered = tokens.placeholder_to_rendered_token or {} + for placeholder, token in rendered.items(): + if token.kind == OWN_FUNCTION_KIND: + return placeholder + return None + + def names_token(ident: str, rendered_value: str) -> bool: return ident == rendered_value or ident in _IDENT_RE.findall(rendered_value) diff --git a/tests/unit/ai_decomp/test_render.py b/tests/unit/ai_decomp/test_render.py index 63cdede..f7fb90f 100644 --- a/tests/unit/ai_decomp/test_render.py +++ b/tests/unit/ai_decomp/test_render.py @@ -18,6 +18,7 @@ RENAME_NOT_ON_LINE, RENAME_UNRESOLVED, display_rows_for_source_lines, + own_function_placeholder, effective_values, find_token, index_of_identifier, @@ -420,3 +421,25 @@ def test_render_stream_of_a_finished_run_keeps_the_source(): def test_render_progress_without_steps_falls_back_to_status(): text = render_progress(_wp(step="", step_index=0, steps_total=0, status="PENDING")) assert "// PENDING" in text.split("\n") + + +def test_own_function_placeholder_finds_the_token_naming_the_function_itself(): + tokens = _tokens( + rendered={ + "@@V_v5@@": _rt("v5"), + "@@F_self@@": _rt("FUN_140004780", kind="own_function"), + "@@F_other@@": _rt("FUN_140001000", kind="function", function_id=9), + } + ) + + assert own_function_placeholder(tokens) == "@@F_self@@" + + +def test_own_function_placeholder_is_none_when_no_own_function_token_exists(): + tokens = _tokens(rendered={"@@V_v5@@": _rt("v5"), "@@P_a1@@": _rt("a1", kind="param")}) + + assert own_function_placeholder(tokens) is None + + +def test_own_function_placeholder_is_none_when_there_are_no_tokens_yet(): + assert own_function_placeholder(_tokens(rendered={})) is None