Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix new LibBS decompiler interface #12

Merged
merged 3 commits into from
Aug 7, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion setup.cfg
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ install_requires =
transformers
tqdm
dailalib
libbs
libbs>=1.18.1

python_requires = >= 3.8
packages = find:
Expand Down
2 changes: 1 addition & 1 deletion varbert/__init__.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
__version__ = "2.2.0"
__version__ = "2.2.1"

import importlib.resources
import tarfile
Expand Down
9 changes: 6 additions & 3 deletions varbert/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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(
Expand Down
7 changes: 6 additions & 1 deletion varbert/text_processor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading