fix(PLU-315): slow data types sync

pyproject.toml

@@ -66,4 +66,5 @@ testpaths = ["tests"]
 addopts = "--import-mode=importlib -s"
 markers = [
     "idalib: tests that require an activated idalib (IDA Pro headless). Run with `pytest -m idalib`. In CI these run on the `ida` self-hosted runner image (IDA + idalib preinstalled, licensing provided by the runner environment).",
+    "ida_ui: tests that launch the real IDA GUI to observe widget behavior. Opt-in: set REAI_UI_TESTS=1 and ensure no interactive IDA instance is running (needs the license seat and a display).",
 ]

reai_toolkit/app/components/dialogs/matching_dialog.py

@@ -1240,7 +1240,7 @@ def import_data_types(self):
                 remote_function_id: int = self.local_func_id_to_remote_func_id[local_function_id]
                 selected_matches[remote_function_id] = self.matched_func_to_original_ea[remote_function_id]
 
-        self.data_types_service.import_data_types(selected_matches)
+        self.data_types_service.import_data_types_async(selected_matches)
 
     # =====================================================================
     # (Optional) page-switch helpers

reai_toolkit/app/services/data_types/data_types_service.py

@@ -1,3 +1,4 @@
+import threading
 from dataclasses import dataclass, field
 
 from revengai import ApiException, Configuration
@@ -30,6 +31,14 @@ class ImportDataTypesService(IThreadService):
     def __init__(self, netstore_service: SimpleNetStore, sdk_config: Configuration) -> None:
         super().__init__(netstore_service=netstore_service, sdk_config=sdk_config)
 
+    def import_data_types_async(self, matches: dict[int, int]) -> None:
+        self.start_worker(target=self._import_worker, args=(matches,))
+
+    def _import_worker(self, _: threading.Event, matches: dict[int, int]) -> None:
+        result: DataTypesImportResult = self.import_data_types(matches)
+        if result.error:
+            logger.error(f"RevEng.AI: {result.error}")
+
     def import_data_types(self, matches: dict[int, int]) -> DataTypesImportResult:
         if len(matches) == 0:
             return DataTypesImportResult()

reai_toolkit/app/services/variable_sync/variable_sync_service.py

@@ -4,7 +4,8 @@
 from typing import Optional, Tuple
 
 from libbs.artifacts import Enum, FunctionHeader, StackVariable, Struct, Typedef
-from libbs.decompilers.ida.compat import execute_read
+from libbs.decompilers.ida.compat import DummyIDACodeView, execute_read
+from libbs.decompilers.ida.compat import function as read_ida_function
 from loguru import logger
 
 from revengai import (
@@ -29,7 +30,20 @@
 
 @execute_read
 def _read_decompiler_function(deci, func_addr: int):
-    return deci.functions.get(func_addr)
+    lowered: int = deci.art_lifter.lower_addr(func_addr)
+    try:
+        if deci.decompiler_available:
+            code_view = DummyIDACodeView(lowered)
+            if code_view.cfunc is None:
+                return None
+            func = read_ida_function(lowered, ida_code_view=code_view)
+        else:
+            func = read_ida_function(lowered, decompiler_available=False)
+    except Exception as e:
+        logger.debug(f"RevEng.AI: could not read function at 0x{lowered:x}: {e}")
+        return None
+
+    return deci.art_lifter.lift(func) if func is not None else None
 
 
 @execute_read

reai_toolkit/app/transformations/import_data_types.py

@@ -1,21 +1,25 @@
 from typing import cast
 
+import ida_funcs
+import ida_typeinf
+import idaapi
+
 import libbs.artifacts
 from libbs.api import DecompilerInterface
-from libbs.decompilers.ida.compat import execute_ui, convert_type_str_to_ida_type
+from libbs.decompilers.ida.compat import execute_write, convert_type_str_to_ida_type
 from loguru import logger
 from revengai import (
-    FunctionArgument,
     Enumeration,
+    FunctionArgument,
     FunctionDataTypesList,
-    FunctionHeader,
-    V2FunctionInfoFuncDepsInner,
-    FunctionInfo,
+    FunctionDataTypesListItem,
     FunctionType,
     Structure,
     TypeDefinition,
 )
 
+APPLY_CHUNK_SIZE = 50
+
 
 class TaggedDependency:
     def __init__(self, dependency: Structure | Enumeration | TypeDefinition) -> None:
@@ -31,69 +35,144 @@ class ImportDataTypes:
     def __init__(self) -> None:
         self.deci: DecompilerInterface
 
-    @execute_ui
     def execute(self, functions: FunctionDataTypesList, matched_function_mapping: dict[int, int] = {}) -> set[int]:
-        self.deci = DecompilerInterface.discover(force_decompiler="ida") # type: ignore
-        lookup: dict[str, TaggedDependency] = {}
+        items: list[FunctionDataTypesListItem] = [
+            item for item in functions.items if item.data_types is not None
+        ]
+        if not items:
+            return set()
+
+        # Track processed dependencies to prevent duplicate imports.
+        # Without this:
+        # - Shared dependencies get re-processed, breaking references (shows as invalid ordinals in IDA)
+        # - Cannot resolve subdependencies (e.g. struct fields that reference other imported types)
+        lookup: dict[str, TaggedDependency] = self._build_lookup(items)
+        if lookup:
+            self._apply_dependencies(lookup)
+
         failed: set[int] = set()
+        total: int = len(items)
+        for start in range(0, total, APPLY_CHUNK_SIZE):
+            chunk: list[FunctionDataTypesListItem] = items[start:start + APPLY_CHUNK_SIZE]
+            failed |= self._apply_chunk(chunk, matched_function_mapping)
+            logger.info(
+                f"RevEng.AI: applied data types to {min(start + APPLY_CHUNK_SIZE, total)}/{total} functions"
+            )
 
-        for function in functions.items:
-            data_types: FunctionInfo | None = function.data_types
+        return failed
 
-            if data_types is None:
-                continue
-
-            # Track processed dependencies to prevent duplicate imports.
-            # Without this:
-            # - Shared dependencies get re-processed, breaking references (shows as invalid ordinals in IDA)
-            # - Cannot resolve subdependencies (e.g. struct fields that reference other imported types)
-            for dep in data_types.func_deps:
+    def _build_lookup(self, items: list[FunctionDataTypesListItem]) -> dict[str, TaggedDependency]:
+        lookup: dict[str, TaggedDependency] = {}
+        for item in items:
+            for dep in item.data_types.func_deps or []:
                 if dep.actual_instance is None:
                     continue
 
                 if dep.actual_instance.name not in lookup:
-                    lookup.update({dep.actual_instance.name: TaggedDependency(dep.actual_instance)}) # type: ignore
+                    lookup[dep.actual_instance.name] = TaggedDependency(dep.actual_instance) # type: ignore
 
-            dependency: V2FunctionInfoFuncDepsInner
-            for dependency in data_types.func_deps:
-                if dependency.actual_instance is None:
-                    continue
+        return lookup
+
+    @execute_write
+    def _apply_dependencies(self, lookup: dict[str, TaggedDependency]) -> None:
+        self.deci = DecompilerInterface.discover(force_decompiler="ida") # type: ignore
+        for tagged_dependency in lookup.values():
+            try:
+                self.process_dependency(tagged_dependency, lookup)
+            except Exception as e:
+                logger.warning(
+                    f"RevEng.AI: skipped dependency {tagged_dependency.name!r}: {e!r}"
+                )
+                tagged_dependency.processed = True
+
+    @execute_write
+    def _apply_chunk(
+        self, chunk: list[FunctionDataTypesListItem], matched_function_mapping: dict[int, int]
+    ) -> set[int]:
+        failed: set[int] = set()
+        for item in chunk:
+            func: FunctionType | None = item.data_types.func_types
+            if func is None:
+                continue
 
-                tagged_dependency: TaggedDependency | None = lookup.get(dependency.actual_instance.name)
-                if tagged_dependency:
-                    try:
-                        self.process_dependency(tagged_dependency, lookup)
-                    except Exception as e:
-                        logger.warning(
-                            f"RevEng.AI: skipped dependency {tagged_dependency.name!r}: {e!r}"
-                        )
-                        tagged_dependency.processed = True
-
-            func: FunctionType | None = data_types.func_types
-            if func:
-                try:
-                    if matched_function_mapping:
-                        ea: int = matched_function_mapping[function.function_id]
-                    else:
-                        ea: int = func.addr
-
-                    if not self.update_function(func, ea) or self._function_types_unparseable(func):
-                        failed.add(function.function_id)
-                except Exception as e:
-                    logger.warning(
-                        f"RevEng.AI: skipped data types for function {function.function_id}: {e!r}"
-                    )
-                    failed.add(function.function_id)
+            try:
+                if matched_function_mapping:
+                    ea: int = matched_function_mapping[item.function_id]
+                else:
+                    ea: int = func.addr
+
+                if not self.apply_function_type(func, ea):
+                    failed.add(item.function_id)
+            except Exception as e:
+                logger.warning(
+                    f"RevEng.AI: skipped data types for function {item.function_id}: {e!r}"
+                )
+                failed.add(item.function_id)
 
         return failed
 
+    def apply_function_type(self, func: FunctionType, ea: int) -> bool:
+        if ida_funcs.get_func(ea) is None:
+            logger.warning(f"failed to update function: {func.name} at 0x{ea:x}")
+            return False
+
+        args: list[FunctionArgument] = sorted(func.header.args.values(), key=lambda a: a.offset)
+        arg_types: list[tuple[str, ida_typeinf.tinfo_t]] = []
+        for arg in args:
+            arg_tif = convert_type_str_to_ida_type(self.normalise_type(arg.type)) if arg.type else None
+            if arg_tif is None:
+                return False
+            arg_types.append((arg.name, arg_tif))
+
+        details: ida_typeinf.func_type_data_t = self._current_func_details(ea)
+
+        ret_type: str = self.normalise_type(func.header.type) if func.header.type else ""
+        if ret_type:
+            ret_tif = convert_type_str_to_ida_type(ret_type)
+            if ret_tif is None:
+                return False
+            details.rettype = ret_tif
+        elif details.rettype.empty():
+            details.rettype = convert_type_str_to_ida_type("void")
+
+        if arg_types:
+            details.clear()
+            for name, arg_tif in arg_types:
+                funcarg = ida_typeinf.funcarg_t()
+                funcarg.name = name
+                funcarg.type = arg_tif
+                details.push_back(funcarg)
+
+        proto = ida_typeinf.tinfo_t()
+        if not proto.create_func(details):
+            return False
+
+        return bool(ida_typeinf.apply_tinfo(ea, proto, ida_typeinf.TINFO_DEFINITE))
+
+    @staticmethod
+    def _current_func_details(ea: int) -> "ida_typeinf.func_type_data_t":
+        details = ida_typeinf.func_type_data_t()
+        existing = ida_typeinf.tinfo_t()
+        if idaapi.get_tinfo(existing, ea) and existing.is_func() and existing.get_func_details(details):
+            return details
+
+        if (
+            ida_typeinf.guess_tinfo(existing, ea) != ida_typeinf.GUESS_FUNC_FAILED
+            and existing.is_func()
+            and existing.get_func_details(details)
+        ):
+            return details
+
+        details = ida_typeinf.func_type_data_t()
+        details.cc = ida_typeinf.CM_CC_UNKNOWN
+        return details
 
     def process_dependency(
         self, tagged_dependency: TaggedDependency, lookup: dict[str, TaggedDependency]
     ) -> None:
         if tagged_dependency.processed:
             return
-        
+
         dependency: Structure | Enumeration | TypeDefinition = tagged_dependency.dependency
         match dependency:
             case Structure():
@@ -110,7 +189,7 @@ def process_dependency(
     def update_struct(self, imported_struct: Structure, lookup: dict[str, TaggedDependency]) -> None:
         if imported_struct.size is None:
             return
-        
+
         for member in imported_struct.members.values():
             subdependency: TaggedDependency | None = lookup.get(member.type)
             if subdependency:
@@ -134,56 +213,6 @@ def update_typedef(self, imported_typedef: TypeDefinition, lookup: dict[str, Tag
             name=imported_typedef.name, type_=normalized_type
         )
 
-    def update_function(self, func: FunctionType, ea: int) -> bool:
-        base_address: int = self.deci.binary_base_addr
-        rva: int = ea - base_address
-
-        target_func: libbs.artifacts.Function | None = self.deci.functions.get(rva) # type: ignore
-        if target_func is None:
-            logger.warning(f"failed to update function: {func.name} at rva: 0x{rva:0x}")
-            return False
-
-        target_func.name = func.name
-        target_func.size = func.size
-        target_func.type = func.type
-
-        # Check the target function has a header.
-        if target_func.header:
-            self.update_header(func.header, target_func)
-
-        self.deci.functions[rva] = target_func
-        return True
-
-    def _function_types_unparseable(self, func: FunctionType) -> bool:
-        if not (func.header and func.header.args):
-            return False
-        for arg in func.header.args.values():
-            type_str: str = self.normalise_type(arg.type) if arg.type else ""
-            if type_str and convert_type_str_to_ida_type(type_str) is None:
-                return True
-        return False
-
-    def update_header(
-        self, imported_header: FunctionHeader, target_function: libbs.artifacts.Function
-    ) -> None:
-        if target_function.header is None:
-            return
-
-        target_function.header.name = imported_header.name
-        target_function.header.type = self.normalise_type(imported_header.type)
-        self.update_function_arguments(imported_header.args, target_function)
-
-    def update_function_arguments(
-        self, imported_args: dict[str, FunctionArgument], target_function: libbs.artifacts.Function
-    ) -> None:
-        if target_function.header is None:
-            return
-
-        for arg in imported_args.values():
-            arg.type = self.normalise_type(arg.type)
-
-        target_function.header.args = {v.offset: v for v in imported_args.values()}
-
     @staticmethod
     def normalise_type(data_type: str) -> str:
         # When we obtain a type from DWARF information, it often looks something like `DWARF/stdint-uintn.h::uint32_t`