Fix new LibBS decompiler interface (#12)

* Fix new decomp interface

* bump

* bump libbs version for copy fix

Author: mahaloz

setup.cfg

@@ -18,7 +18,7 @@ install_requires =
     transformers
     tqdm
     dailalib
-    libbs
+    libbs>=1.18.1
 
 python_requires = >= 3.8
 packages = find:

varbert/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "2.2.0"
+__version__ = "2.2.1"
 
 import importlib.resources
 import tarfile

varbert/api.py

@@ -20,8 +20,8 @@ def __init__(self, **kwargs):
         ) if not self._delay_init else None
 
     def predict_variable_names(
-            self, function: Function = None, decompilation_text: Optional[str] = None, use_decompiler=True,
-            remove_bad_names=True
+        self, function: Function = None, decompilation_text: Optional[str] = None, use_decompiler=True,
+        remove_bad_names=True, **kwargs
     ) -> Tuple[Dict[str, str], str]:
         """
         Predict variable names for a function or decompilation text. You can use this function in two ways:
@@ -61,7 +61,10 @@ def predict_variable_names(
             if not use_decompiler:
                 raise ValueError("Must provide decompilation text if not using decompiler")
 
-            decompilation_text = self._dec_interface.decompile(function.addr)
+            decomp = self._dec_interface.decompile(function.addr)
+            if decomp is None or not decomp.text:
+                raise ValueError(f"Unable to decompile function {function}")
+            decompilation_text = decomp.text
 
         # preprocess text for training
         preprocessor = DecompilationTextProcessor(

varbert/text_processor.py

@@ -89,7 +89,12 @@ def _process_code_with_decompiler(self):
         # replace all original names with tmp tokenized names
         self._decompiler.rename_local_variables_by_names(self._func, og_name_to_tokenized_name)
         # get the decomp, fix the tmp tokens
-        tokenized_dec_text = self._decompiler.decompile(self._func.addr)
+        tokenized_decomp = self._decompiler.decompile(self._func.addr)
+        if tokenized_decomp is None:
+            _l.error("Decompiler failed to decompile function.")
+            return
+
+        tokenized_dec_text = tokenized_decomp.text
         tokenized_dec_text = tokenized_dec_text.replace(tmp_token, "@@")
 
         # revert to the original names in the decomp