From 880795d2395732f991ad771cc592054e619342d6 Mon Sep 17 00:00:00 2001
From: Mc-muffin <8714476+Mc-muffin@users.noreply.github.com>
Date: Mon, 25 Nov 2024 22:23:54 -0500
Subject: [PATCH] organization

---
 pythonlib/formats/fps4.py                     |  4 +-
 pythonlib/formats/pak.py                      |  3 +
 pythonlib/formats/rebirth/__init__.py         |  0
 pythonlib/formats/{ => rebirth}/pak2.py       |  0
 pythonlib/formats/{ => rebirth}/scpk.py       |  6 +-
 pythonlib/formats/{ => rebirth}/theirsce.py   | 63 ++++++++++++++++++-
 .../formats/{ => rebirth}/theirsce_funcs.py   |  0
 .../{ => rebirth}/theirsce_instructions.py    |  2 +
 pythonlib/games/ToolsTOR.py                   |  8 +--
 9 files changed, 75 insertions(+), 11 deletions(-)
 create mode 100644 pythonlib/formats/rebirth/__init__.py
 rename pythonlib/formats/{ => rebirth}/pak2.py (100%)
 rename pythonlib/formats/{ => rebirth}/scpk.py (97%)
 rename pythonlib/formats/{ => rebirth}/theirsce.py (83%)
 rename pythonlib/formats/{ => rebirth}/theirsce_funcs.py (100%)
 rename pythonlib/formats/{ => rebirth}/theirsce_instructions.py (99%)

diff --git a/pythonlib/formats/fps4.py b/pythonlib/formats/fps4.py
index 1737ef2..63554e2 100644
--- a/pythonlib/formats/fps4.py
+++ b/pythonlib/formats/fps4.py
@@ -1,4 +1,6 @@
-import os, struct, sys
+import os
+import struct
+import sys
 
 def dump_fps4(name, name2, base_path):
 
diff --git a/pythonlib/formats/pak.py b/pythonlib/formats/pak.py
index 44ae608..1347522 100644
--- a/pythonlib/formats/pak.py
+++ b/pythonlib/formats/pak.py
@@ -186,6 +186,9 @@ class Pak:
 
     def __getitem__(self, item):
         return self.files[item]
+    
+    def __setitem__(self, item, data):
+        self.files[item] = data
 
     def __len__(self):
         return len(self.files)
diff --git a/pythonlib/formats/rebirth/__init__.py b/pythonlib/formats/rebirth/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/pythonlib/formats/pak2.py b/pythonlib/formats/rebirth/pak2.py
similarity index 100%
rename from pythonlib/formats/pak2.py
rename to pythonlib/formats/rebirth/pak2.py
diff --git a/pythonlib/formats/scpk.py b/pythonlib/formats/rebirth/scpk.py
similarity index 97%
rename from pythonlib/formats/scpk.py
rename to pythonlib/formats/rebirth/scpk.py
index 154c74f..cac773d 100644
--- a/pythonlib/formats/scpk.py
+++ b/pythonlib/formats/rebirth/scpk.py
@@ -2,9 +2,9 @@ import struct
 from dataclasses import dataclass
 from pathlib import Path
 
-from ..formats.FileIO import FileIO
-from ..formats.pak import Pak
-from ..utils import comptolib
+from ..FileIO import FileIO
+from ..pak import Pak
+from ...utils import comptolib
 
 MAGIC = b"SCPK"
 
diff --git a/pythonlib/formats/theirsce.py b/pythonlib/formats/rebirth/theirsce.py
similarity index 83%
rename from pythonlib/formats/theirsce.py
rename to pythonlib/formats/rebirth/theirsce.py
index e360d18..e1b1f66 100644
--- a/pythonlib/formats/theirsce.py
+++ b/pythonlib/formats/rebirth/theirsce.py
@@ -1,10 +1,29 @@
+from dataclasses import dataclass
 from io import BytesIO
 from pathlib import Path
 from typing import Generator, Union
 
-from .FileIO import FileIO
-from .theirsce_funcs import *
-from .theirsce_instructions import *
+from ..FileIO import FileIO
+from .theirsce_funcs import SYSCALL_NAMES
+from .theirsce_instructions import (
+    AluOperation,
+    BranchType,
+    InstructionType,
+    ReferenceScope,
+    TheirsceAcquireInstruction,
+    TheirsceAluInstruction,
+    TheirsceBaseInstruction,
+    TheirsceBranchInstruction,
+    TheirsceBreakInstruction,
+    TheirsceLocalCallInstruction,
+    TheirscePushInstruction,
+    TheirsceReferenceInstruction,
+    TheirsceReturnInstruction,
+    TheirsceSpecialReferenceInstruction,
+    TheirsceStringInstruction,
+    TheirsceSyscallInstruction,
+    VariableType,
+)
 
 # I followed other project when making this class
 # not sure how if it's the best approach
@@ -17,6 +36,12 @@ class subsection:
     unk2: int
     off:  int
 
+@dataclass
+class function:
+    arguments: list
+    body:  list
+    visited: bool = False
+
 
 class Theirsce(FileIO):
     def __init__(self, path: Union[Path, str, BytesIO, bytes] ="") -> None:
@@ -48,6 +73,8 @@ class Theirsce(FileIO):
                 subsections.append(sub)
             self.sections.append(subsections)
             self.seek(pos)
+
+        self.functions = self.get_functions()
     
     def __enter__(self):
         return super().__enter__()
@@ -122,6 +149,36 @@ class Theirsce(FileIO):
 
         return data
 
+    def get_functions(self):
+        funs: dict[int, function] = {}
+        ops = list(self.walk_code())
+
+        bd = []
+        fn_st = ops[0].position; ops.pop(0)
+        for val in ops:
+            if val.type == InstructionType.BREAK:
+                if bd[0].type == InstructionType.ACQUIRE:
+                    funs[fn_st] = function(bd[0].params, bd)
+                    bd.pop(0)
+                else:
+                    funs[fn_st] = function([], bd)
+                fn_st = val.position
+                bd = []
+            else:
+                bd.append(val)
+
+        if bd[0].type == InstructionType.ACQUIRE:
+            funs[fn_st] = function(bd[0].params, bd)
+            bd.pop(0)
+        else:
+            funs[fn_st] = function([], bd)
+
+        return funs
+    
+    def reset_seen_functions(self):
+        for fun in self.functions.values():
+            fun.visited = False
+
     def read_opcode(self) -> TheirsceBaseInstruction:
         pos = self.tell()
         opcode = self.read_uint8()
diff --git a/pythonlib/formats/theirsce_funcs.py b/pythonlib/formats/rebirth/theirsce_funcs.py
similarity index 100%
rename from pythonlib/formats/theirsce_funcs.py
rename to pythonlib/formats/rebirth/theirsce_funcs.py
diff --git a/pythonlib/formats/theirsce_instructions.py b/pythonlib/formats/rebirth/theirsce_instructions.py
similarity index 99%
rename from pythonlib/formats/theirsce_instructions.py
rename to pythonlib/formats/rebirth/theirsce_instructions.py
index 3e8e0fd..dc52276 100644
--- a/pythonlib/formats/theirsce_instructions.py
+++ b/pythonlib/formats/rebirth/theirsce_instructions.py
@@ -83,6 +83,7 @@ class AluOperation(Enum):
 
 @dataclass
 class TheirsceBaseInstruction:
+    position: int
     mnemonic: str = field(default=False, init=False) # type: ignore
     type: InstructionType = field(default=False, init=False) # type: ignore
     #size: int
@@ -182,6 +183,7 @@ class TheirsceReferenceInstruction(TheirsceBaseInstruction):
     offset: int
     shift:  int
     position: int
+    is_array: bool = False
     mnemonic = "REF"
     type = InstructionType.REFERENCE
 
diff --git a/pythonlib/games/ToolsTOR.py b/pythonlib/games/ToolsTOR.py
index bbec446..114d750 100644
--- a/pythonlib/games/ToolsTOR.py
+++ b/pythonlib/games/ToolsTOR.py
@@ -16,13 +16,13 @@ import pyjson5 as json
 from tqdm import tqdm
 import io
 
-import pythonlib.formats.pak2 as pak2lib
+import pythonlib.formats.rebirth.pak2 as pak2lib
 import pythonlib.utils.comptolib as comptolib
 from pythonlib.formats.FileIO import FileIO
 from pythonlib.formats.pak import Pak
-from pythonlib.formats.scpk import Scpk
-from pythonlib.formats.theirsce import Theirsce
-from pythonlib.formats.theirsce_instructions import (AluOperation,
+from pythonlib.formats.rebirth.scpk import Scpk
+from pythonlib.formats.rebirth.theirsce import Theirsce
+from pythonlib.formats.rebirth.theirsce_instructions import (AluOperation,
                                                      InstructionType,
                                                      TheirsceBaseInstruction)