Implement Extended Masking

.gitignore

@@ -20,4 +20,6 @@ mise.toml
 dist/
 decompiled_*/
 decompiled_*.py
+ROADMAP.md
 
+.opencode-cache/

pylingual/decompiler.py

@@ -44,6 +44,7 @@
 from pylingual.models import CacheTranslator, load_models
 from pylingual.utils.generate_bytecode import CompileError, PyenvError, compile_version
 from pylingual.masking.model_disasm import create_global_masker, restore_masked_source_text
+from pylingual.preprocessor import Preprocessor
 from pylingual.editable_bytecode import PYCFile
 from pylingual.segmentation.segmentation_search_strategies import get_top_k_predictions, m_deep_top_k, naive_confidence_priority, filter_subwords
 from pylingual.segmentation.sliding_window import merge, sliding_window
@@ -122,6 +123,8 @@ def __call__(self):
         with tempfile.TemporaryDirectory() as tmp:
             self.tmp = Path(tmp)
 
+            self.preprocess_bytecode()
+            self.check_timeout()
             self.mask_bytecode()
             self.check_timeout()
             self.run_segmentation()
@@ -208,6 +211,13 @@ def purge_comp_errors(self):
             self.source_context.purged_cfts = []
             return []
 
+    def preprocess_bytecode(self):
+        logger.info(f"Preprocessing bytecode for {self.name}...")
+        try:
+            Preprocessor().preprocess(self.pyc)
+        except Exception as e:
+            logger.warning(f"Preprocessing failed for {self.name}, continuing without container recovery: {e}")
+
     def mask_bytecode(self):
         logger.info(f"Masking bytecode for {self.name}...")
         try:

pylingual/editable_bytecode/EditableBytecode.py

@@ -389,6 +389,7 @@ def _bake_jumps(self, add_extended=False):
                                 instruction.argval = instruction.target.offset
                             else:
                                 instruction.arg = instruction.argval = instruction.target.offset
+                        instruction.refresh_argrepr()
                     if add_extended:
                         # count the number of necessary preceeding EXTENDED_ARG instructions
                         n_extendeds_needed = (instruction.arg > 0xFF) + (instruction.arg > 0xFFFF) + (instruction.arg > 0xFFFFFF)

pylingual/editable_bytecode/Instruction.py

@@ -133,6 +133,12 @@ def target(self, value):
             raise AttributeError("Only jump instructions have target attributes, not " + repr(self))
         self._target = value
 
+    def refresh_argrepr(self):
+        """Refresh argrepr to reflect the current state of the instruction, e.g. after offset changes."""
+        if self.is_jump and self._target is not None:
+            self.argval = self._target.offset
+            self.argrepr = "to " + repr(self._target.offset)
+
     def add_reqs(self, *reqs):
         for r in reqs:
             if isinstance(r, tuple):

pylingual/masking/global_masker.py

@@ -1,3 +1,4 @@
+import ast
 from collections.abc import MutableMapping
 from copy import deepcopy
 from xdis import iscode
@@ -45,10 +46,22 @@ def _key_transform(self, key):
             return (key.value.decode("utf-8"), str)
         if type(key) == LongTypeForPython3:
             return (key.value, int)
+        if isinstance(key, (list, set, dict, tuple, frozenset)):
+            return (repr(key), type(key))
+        try:
+            hash(key)
+        except TypeError:
+            return (repr(key), type(key))
         return (key, type(key))
 
     def _key_restore(self, key):
-        return key[0]
+        value, typ = key
+        if type(value) is str and typ is not str:
+            try:
+                return ast.literal_eval(value)
+            except (ValueError, SyntaxError):
+                return value
+        return value
 
 
 #### Main Masker
@@ -183,16 +196,20 @@ def get_model_view(self, inst: Inst) -> str:
             elif isinstance(inst.argval, str) and inst.argval in self.global_tab:  # have to do this check incase string is varname of type annotation
                 view = f"{inst.opname} , {repr(self.mask(inst.argval))}"
 
+            elif type(inst.argval) in (list, tuple, frozenset, set, dict) and getattr(inst, "preprocessed_container", False):
+                # Container constants recovered by the preprocessor are masked atomically.
+                # TODO: Providing type information via type(inst.argval).__name__(<mask_N>) would
+                # give the model useful context, but the current model was not trained on that format
+                # and produces garbled output (e.g. splitting "dict" across a mask boundary). Reverting
+                # to a plain mask token matches the training data. Revisit this after model retraining.
+                view = f"{inst.opname} , {self.mask(inst.argval)}"
             elif type(inst.argval) in (list, tuple, frozenset, set):
-                # do recursive in-place replacement of list elems if they are strs or bytestrs
-
                 def replace_list(consts):
-                    """recursive replacement of elements in arbitrary list-like objects"""
                     for idx, const in enumerate(consts):
                         if isinstance(const, str):
                             consts[idx] = repr(self.mask(inst.bytecode.resolve_namespace(const)))
                         elif const is None:
-                            continue  # don't mask None
+                            continue
                         elif type(const) in (list, tuple, frozenset):
                             consts[idx] = type(const)(replace_list(list(const)))
                         elif type(const) is slice:
@@ -201,19 +218,12 @@ def replace_list(consts):
                             consts[idx] = self.mask(const)
                     return consts
 
-                consts = list(deepcopy(inst.argval))
-                consts = replace_list(consts)
-
+                consts = replace_list(list(deepcopy(inst.argval)))
                 if inst.bytecode.version < (3, 11):
-                    # Format keyword argument list
-                    # We left pad the list of kwargs so the model doesnt have to "look ahead"
                     next_insts = inst.next_instructions
-                    next_inst = next_insts[0] if next_insts != [] else None
+                    next_inst = next_insts[0] if next_insts else None
                     if next_inst is not None and inst.opname == "LOAD_CONST" and next_inst.opname == "CALL_FUNCTION_KW":
                         consts = ["<KWARG_PAD>"] * (next_inst.argval - len(consts)) + consts
-
-                # cast back to original type and print repr
-                # demote quotes one layer
                 arg_repr = repr(type(inst.argval)(consts)).replace("'", "").replace('"', "'")
                 view = f"{inst.opname} , {arg_repr}"
             elif type(inst.argval) is slice:

pylingual/masking/model_disasm.py

@@ -36,6 +36,7 @@ def create_global_masker(bytecode: EditableBytecode) -> Masker:
 
     for bc in bytecode.iter_bytecodes():
         bc_co = bc.to_code(no_lnotab=True)
+        preprocessed = [inst.argval for inst in bc.instructions if getattr(inst, "preprocessed_container", False)]
 
         # create consts
         consts = list(deepcopy(bc_co.co_consts))
@@ -54,13 +55,13 @@ def create_global_masker(bytecode: EditableBytecode) -> Masker:
             # Don't needlessly increment the global_idx
             if const in global_tab:
                 continue
-            if type(const) in (list, tuple, frozenset, set):
+            if type(const) in (list, tuple, frozenset, set) and not any(type(const) is type(value) and const == value for value in preprocessed):
                 consts.extend(const)
             elif type(const) is slice:
                 # decompose slice constants for 3.14+
                 consts.extend([const.start, const.stop, const.step])
             else:
-                global_tab.update({bc.resolve_namespace(const): f"<mask_{global_idx}>"})
+                global_tab[bc.resolve_namespace(const)] = f"<mask_{global_idx}>"
                 global_idx += 1
 
         # create names
@@ -69,41 +70,41 @@ def create_global_masker(bytecode: EditableBytecode) -> Masker:
                 for n in name:
                     if n in global_tab:
                         continue
-                    global_tab.update({bc.resolve_namespace(n): f"<mask_{global_idx}>"})
+                    global_tab[bc.resolve_namespace(n)] = f"<mask_{global_idx}>"
                     global_idx += 1
             else:
                 if name in global_tab:
                     continue
-                global_tab.update({bc.resolve_namespace(name): f"<mask_{global_idx}>"})
+                global_tab[bc.resolve_namespace(name)] = f"<mask_{global_idx}>"
                 global_idx += 1
 
         for free in bc_co.co_freevars:
             if free in global_tab:
                 continue
-            global_tab.update({free: f"<mask_{global_idx}>"})
+            global_tab[free] = f"<mask_{global_idx}>"
             global_idx += 1
 
         if bc.version >= (3, 11):
             for cell in bc_co.co_cellvars:
                 if cell in global_tab:
                     continue
-                global_tab.update({cell: f"<mask_{global_idx}>"})
+                global_tab[cell] = f"<mask_{global_idx}>"
                 global_idx += 1
 
         for local in bc_co.co_varnames:
             if isinstance(local, tuple):
                 for local_item in local:
                     if local_item in global_tab:
                         continue
-                    global_tab.update({bc.resolve_namespace(local_item): f"<mask_{global_idx}>"})
+                    global_tab[bc.resolve_namespace(local_item)] = f"<mask_{global_idx}>"
                     global_idx += 1
             else:
                 if local in global_tab:
                     continue
-                global_tab.update({bc.resolve_namespace(local): f"<mask_{global_idx}>"})
+                global_tab[bc.resolve_namespace(local)] = f"<mask_{global_idx}>"
                 global_idx += 1
 
-        global_tab.update({bc_co.co_name: f"<mask_{global_idx}>"})
+        global_tab[bc_co.co_name] = f"<mask_{global_idx}>"
         global_idx += 1
 
     return global_masker

pylingual/preprocessor/__init__.py

@@ -0,0 +1 @@
+from .preprocessor import Preprocessor

pylingual/preprocessor/container_recovery/recovery.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from collections import defaultdict
+
+from .segment import Recovery, Segment
+
+
+_strategies: dict[int, list] = defaultdict(list)
+
+
+def register_recovery_strategy(priority: int):
+    def decorator(func):
+        _strategies[priority].append(func)
+        return func
+    return decorator
+
+
+def get_strategies() -> list:
+    result = []
+    for priority in sorted(_strategies.keys()):
+        result.extend(_strategies[priority])
+    return result
+
+
+def recover(seg: Segment, indent: int = 0) -> Recovery:
+    for strategy in get_strategies():
+        result = strategy(seg, indent)
+        if result is not None:
+            return result
+    return Recovery(None, False)
+
+
+from . import strategies  # noqa: E402,F401  — trigger strategy registration

pylingual/preprocessor/container_recovery/segment.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+INDENT = "  "
+SIDE_OFFSET = " " * 14
+
+
+@dataclass
+class Recovery:
+    value: object
+    complete: bool
+
+
+@dataclass
+class Segment:
+    tag: str
+    ordered_children: list[Segment | tuple[int, object]]
+    base_stack_depth: int | None = None
+    start_offset: int | None = None
+    end_offset: int | None = None
+
+    def compute_offsets(self) -> None:
+        offsets = []
+        for child in self.ordered_children:
+            if isinstance(child, Segment):
+                child.compute_offsets()
+                if child.start_offset is not None:
+                    offsets.append(child.start_offset)
+                if child.end_offset is not None:
+                    offsets.append(child.end_offset)
+            else:
+                instr = child[1]
+                if hasattr(instr, "offset"):
+                    offsets.append(instr.offset)
+        if offsets:
+            self.start_offset = min(offsets)
+            self.end_offset = max(offsets)
+
+    def print(self, indent=0):
+        print(INDENT * indent, self.tag, sep="")
+        for child in self.ordered_children:
+            if isinstance(child, type(self)):
+                child.print(indent + 1)
+            else:
+                colpad = len(f"{SIDE_OFFSET}{child[0]} {child[1].opname}")
+                print(f"{SIDE_OFFSET}{child[0]} {child[1].opname}" + (" " * (20 - colpad) + f"{self.base_stack_depth}" if self.base_stack_depth else ""))
+        print()