From f75b9f038e6078659781babe0f34da26294b32d6 Mon Sep 17 00:00:00 2001
From: Jennifer Taylor <dragonminded@dragonminded.com>
Date: Wed, 5 May 2021 00:55:02 +0000
Subject: [PATCH] Implement a few more opcodes to get a few more Bishi files
 decompiling.

---
 bemani/format/afp/decompile.py     | 52 ++++++++++++++++++++++++++++--
 bemani/tests/test_afp_decompile.py |  3 --
 2 files changed, 49 insertions(+), 6 deletions(-)

diff --git a/bemani/format/afp/decompile.py b/bemani/format/afp/decompile.py
index 3678039..2e280cf 100644
--- a/bemani/format/afp/decompile.py
+++ b/bemani/format/afp/decompile.py
@@ -426,6 +426,22 @@ class Array(Expression):
         return f"[{', '.join(params)}]"
 
 
+class Object(Expression):
+    # Call a method on an object.
+    def __init__(self, params: Dict[Any, Any]) -> None:
+        self.params = params
+
+    def __repr__(self) -> str:
+        return self.render("")
+
+    def render(self, parent_prefix: str, nested: bool = False) -> str:
+        params = [f"{value_ref(key, parent_prefix)}: {value_ref(val, parent_prefix)}" for (key, val) in self.params.items()]
+        lpar = "{"
+        rpar = "}"
+
+        return f"{lpar}{', '.join(params)}{rpar}"
+
+
 class FunctionCall(Expression):
     # Call a method on an object.
     def __init__(self, name: Union[str, StringConstant], params: List[Any]) -> None:
@@ -2653,10 +2669,24 @@ class ByteCodeDecompiler(VerboseOutput):
                     num_entries = stack.pop()
                     if not isinstance(num_entries, int):
                         raise Exception("Logic error!")
-                    params = []
+                    arrparams = []
                     for _ in range(num_entries):
-                        params.append(stack.pop())
-                    stack.append(Array(params))
+                        arrparams.append(stack.pop())
+                    stack.append(Array(arrparams))
+
+                    chunk.actions[i] = NopStatement()
+                    continue
+
+                if action.opcode == AP2Action.INIT_OBJECT:
+                    num_entries = stack.pop()
+                    if not isinstance(num_entries, int):
+                        raise Exception("Logic error!")
+                    objparams: Dict[Any, Any] = {}
+                    for _ in range(num_entries):
+                        val = stack.pop()
+                        key = stack.pop()
+                        objparams[key] = val
+                    stack.append(Object(objparams))
 
                     chunk.actions[i] = NopStatement()
                     continue
@@ -2730,6 +2760,22 @@ class ByteCodeDecompiler(VerboseOutput):
                     chunk.actions[i] = NopStatement()
                     continue
 
+                if action.opcode == AP2Action.BIT_L_SHIFT:
+                    shift_amt = stack.pop()
+                    shift_val = stack.pop()
+                    stack.append(ArithmeticExpression(shift_val, "<<", shift_amt))
+
+                    chunk.actions[i] = NopStatement()
+                    continue
+
+                if action.opcode in {AP2Action.BIT_R_SHIFT, AP2Action.BIT_U_R_SHIFT}:
+                    shift_amt = stack.pop()
+                    shift_val = stack.pop()
+                    stack.append(ArithmeticExpression(shift_val, ">>", shift_amt))
+
+                    chunk.actions[i] = NopStatement()
+                    continue
+
                 if action.opcode == AP2Action.EQUALS2:
                     expr2 = stack.pop()
                     expr1 = stack.pop()
diff --git a/bemani/tests/test_afp_decompile.py b/bemani/tests/test_afp_decompile.py
index 1b2534e..5e7a6d7 100644
--- a/bemani/tests/test_afp_decompile.py
+++ b/bemani/tests/test_afp_decompile.py
@@ -628,9 +628,6 @@ class TestAFPDecompile(ExtendedTestCase):
             # "returning" early from a function.
         ])
         statements = self.__call_decompile(bytecode)
-
-        # TODO: The output should be optimized to remove the early return and move the
-        # start playing section inside the if.
         self.assertEqual(self.__equiv(statements), ["if (not True) {\n  return\n}", "builtin_StartPlaying()"])
 
     def test_if_handling_diamond(self) -> None: