jak-project/decompiler/gui/decompiler_gui.py

from PyQt5.QtWidgets import QApplication, QLabel, QMainWindow, QTreeView, QVBoxLayout, QWidget, QPlainTextEdit, QLineEdit, QListView, QDialog, QSplitter, QSizePolicy
from PyQt5.Qt import QStandardItemModel, QStandardItem, QFont, QModelIndex
import json
import re
import os

def get_monospaced_font():
	"""
	Get a monospaced font. Should work on both windows and linux.
	"""
	font = QFont("monospace")
	font.setStyleHint(QFont.TypeWriter)
	return font

def get_jak_path():
	"""
	Get a path to jak-project/
	"""
	return os.path.join(os.path.dirname(os.path.realpath(__file__)), "..", "..")

def segment_id_to_name(sid):
	"""
	GOAL segment ID to name string
	"""
	if sid == 0:
		return "main"
	elif sid == 1:
		return "debug"
	elif sid == 2:
		return "top-level"
	else:
		return "INVALID-SEGMENT"

# Hold all the metadata for an object file
class ObjFile():
	def __init__(self, obj):
		"""
		Convert the json data in obj.txt.
		If the format changes or we add new fields, they should be added here.
		"""
		self.unique_name = obj[0] # The unique name that's used by the decompiler.
		self.name_in_dgo = obj[1] # Name in the game.
		self.version = obj[2]     # GOAL object file format version

	def get_description(self):
		return "Name: {}\n  Version: {}\n  Name in game: {}".format(self.unique_name, self.version, self.name_in_dgo)

# Hold all of the object files in a dgo.
class DgoFile():
	def __init__(self):
		self.obj_files = dict() 
	def add_obj(self, obj):
		self.obj_files[obj[0]] = ObjFile(obj)

# Hold all DGOs/Object files.
class FileMap():
	def __init__(self):
		self.dgo_files = dict()
		self.all_objs = dict()

	def add_obj_to_dgo(self, dgo, obj):
		if not(dgo in self.dgo_files):
			self.dgo_files[dgo] = DgoFile()
		self.dgo_files[dgo].add_obj(obj)
		self.all_objs[obj[0]] = ObjFile(obj)

	def get_objs_matching_regex(self, regex):
		"""
		Get a list of object files with a name that matches the given regex.
		"""
		try:
			r = re.compile(regex)
		except:
			return []
		return list(filter(r.match, self.all_objs.keys()))


def load_obj_map_file(file_path):
	"""
	Load the obj.txt file generated by the decompiler.
	Return a FileMap.
	"""
	file_map = FileMap()
	with open(file_path) as f:
		json_data = json.loads(f.read())

	for obj_file in json_data:
		for dgo in obj_file[3]:
			file_map.add_obj_to_dgo(dgo, obj_file)
	return file_map

class ObjectFileView(QDialog):
	def __init__(self, name):
		super().__init__()
		self.setWindowTitle(name)
		with open(os.path.join(get_jak_path(), "decompiler_out", "{}_asm.json".format(name))) as f:
			self.asm_data = json.loads(f.read())

		main_layout = QVBoxLayout()
		monospaced_font = get_monospaced_font()
		self.header_label = QLabel()

		main_layout.addWidget(self.header_label)

		function_splitter = QSplitter()
		function_splitter.setSizePolicy(QSizePolicy(QSizePolicy.Expanding,QSizePolicy.Expanding))


		self.function_list = QTreeView()
		self.function_list_model = QStandardItemModel()
		self.functions_by_name = dict()
		root = self.function_list_model.invisibleRootItem()
		seg_roots = []

		for i in range(3):
			seg_entry = QStandardItem(segment_id_to_name(i))
			seg_entry.setFont(monospaced_font)
			seg_entry.setEditable(False)
			root.appendRow(seg_entry)
			seg_roots.append(seg_entry)

		for f in self.asm_data["functions"]:
			function_entry = QStandardItem(f["name"])
			function_entry.setFont(monospaced_font)
			function_entry.setEditable(False)
			seg_roots[f["segment"]].appendRow(function_entry)
			self.functions_by_name[f["name"]] = f

		self.header_label.setText("Object File {} Functions ({} total):".format(name, len(self.asm_data["functions"])))

		self.function_list.setModel(self.function_list_model)
		self.function_list.clicked.connect(self.display_function)
		function_splitter.addWidget(self.function_list)



		layout = QVBoxLayout()

		self.function_header_label = QLabel("No function selected")
		self.function_header_label.setFont(monospaced_font)
		self.header_label.setSizePolicy(QSizePolicy(QSizePolicy.Minimum,
                                                 QSizePolicy.Minimum))
		layout.addWidget(self.function_header_label)

		self.op_asm_split_view = QSplitter()
		self.op_asm_split_view.setSizePolicy(QSizePolicy(QSizePolicy.Expanding,QSizePolicy.Expanding))

		self.basic_op_pane = QListView()
		self.basic_op_pane.clicked.connect(self.basic_op_clicked)
		#layout.addWidget(self.basic_op_pane)
		self.op_asm_split_view.addWidget(self.basic_op_pane)

		self.asm_pane = QListView()
		self.op_asm_split_view.addWidget(self.asm_pane)

		layout.addWidget(self.op_asm_split_view)

		self.asm_display = QPlainTextEdit()
		self.asm_display.setMaximumHeight(80)
		layout.addWidget(self.asm_display)

		self.warnings_label = QLabel()
		layout.addWidget(self.warnings_label)

		widget = QWidget()
		widget.setLayout(layout)
		function_splitter.addWidget(widget)
		main_layout.addWidget(function_splitter)

		# add it to the window!
		self.setLayout(main_layout)

	def display_function(self, item):
		name = item.data()
		monospaced_font = get_monospaced_font()
		func = self.functions_by_name[name]
		basic_op_model = QStandardItemModel()
		basic_op_root = basic_op_model.invisibleRootItem()
		asm_model = QStandardItemModel()
		asm_root = asm_model.invisibleRootItem()

		self.basic_id_to_asm = []
		self.current_function = name
		op_idx = 0
		basic_idx = 0
		for op in func["asm"]:
			if "label" in op:
				asm_item = QStandardItem(op["label"] + "\n    " + op["asm_op"])
			else:
				asm_item = QStandardItem("    " + op["asm_op"])
			asm_item.setFont(monospaced_font)
			asm_item.setEditable(False)
			asm_root.appendRow(asm_item)

			if "basic_op" in op:
				if "label" in op:
					basic_item = QStandardItem(op["label"] + "\n    " + op["basic_op"])
				else:
					basic_item = QStandardItem("    " + op["basic_op"])
				basic_item.setFont(monospaced_font)
				basic_item.setEditable(False)
				basic_op_root.appendRow(basic_item)
				self.basic_id_to_asm.append(op_idx)
				basic_idx = basic_idx + 1
			op_idx = op_idx + 1
		self.basic_id_to_asm.append(op_idx)
		self.basic_op_pane.setModel(basic_op_model)
		self.asm_pane.setModel(asm_model)
		self.warnings_label.setText(func["warnings"])
		self.asm_display.setPlainText("")
		self.function_header_label.setText("{}, type: {}\nfunc: {} obj: {}".format(name, func["type"], func["name"], func["parent_object"]))

	def basic_op_clicked(self, item):
		text = ""
		added_reg = 0
		asm_idx = self.basic_id_to_asm[item.row()]

		asm_op = self.functions_by_name[self.current_function]["asm"][asm_idx]
		if "type_map" in asm_op:
			for reg, type_name in asm_op["type_map"].items():
				text += "{}: {} ".format(reg, type_name)
				added_reg += 1
				if added_reg >= 4:
					text += "\n"
					added_reg = 0
			text += "\n"

		for i in range(asm_idx, self.basic_id_to_asm[item.row() + 1]):
			text += self.functions_by_name[self.current_function]["asm"][i]["asm_op"] + "\n"
		op = self.functions_by_name[self.current_function]["asm"][asm_idx]
		if "referenced_string" in op:
			text += op["referenced_string"]
		self.asm_display.setPlainText(text)
		self.asm_display.setFont(get_monospaced_font())
		self.asm_pane.setCurrentIndex(self.asm_pane.model().index(asm_idx, 0))


# A window for browsing all the object files.
# Doesn't actually know anything about what's in the files, it's just used to select a file.
class ObjectFileBrowser(QMainWindow):
	def __init__(self, obj_map):
		self.obj_map = obj_map
		super().__init__()
		self.setWindowTitle("Object File Browser")
		self.childen_windows = []
		layout = QVBoxLayout()
		monospaced_font = get_monospaced_font()

		layout.addWidget(QLabel("Browse object files by dgo..."))

		# Set up the tree view
		self.tree = QTreeView()
		self.tree_model = QStandardItemModel()
		self.tree_root = self.tree_model.invisibleRootItem()
		for dgo_name, dgo in obj_map.dgo_files.items():
			dgo_entry = QStandardItem(dgo_name)
			dgo_entry.setFont(monospaced_font)
			dgo_entry.setEditable(False)
			for obj_name, obj in dgo.obj_files.items():
				obj_entry = QStandardItem(obj_name)
				obj_entry.setFont(monospaced_font)
				obj_entry.setEditable(False)
				dgo_entry.appendRow(obj_entry)
			self.tree_root.appendRow(dgo_entry)

		self.tree.setModel(self.tree_model)
		self.tree.clicked.connect(self.handle_tree_click)
		self.tree.doubleClicked.connect(self.handle_tree_double_click)
		layout.addWidget(self.tree)

		# Set up the Search Box
		layout.addWidget(QLabel("Or search for object (regex):"))
		self.search_box = QLineEdit()
		self.search_box.textChanged.connect(self.handle_search_change)
		layout.addWidget(self.search_box)

		# Set up Search Results
		self.search_result = QListView()
		layout.addWidget(self.search_result)
		self.search_result.clicked.connect(self.handle_search_result_click)
		self.search_result.doubleClicked.connect(self.handle_search_result_double_click)
		self.search_result.setMaximumHeight(200)

		# Set up the info box at the bottom
		self.text_box = QPlainTextEdit()
		self.text_box.setReadOnly(True)
		self.text_box.setFont(monospaced_font)
		layout.addWidget(self.text_box)
		self.text_box.setMaximumHeight(100)
		self.text_box.setPlainText("Select an object file to see details. Double click to open.")

		# add it to the window!
		widget = QWidget()
		widget.setLayout(layout)
		self.setCentralWidget(widget)

	def handle_tree_click(self, val):
		if not(val.parent().isValid()):
			return
		dgo = val.parent().data()
		obj = val.data()
		obj_info = self.obj_map.dgo_files[dgo].obj_files[obj]
		self.text_box.setPlainText("{}\n  DGO: {}".format(obj_info.get_description(), dgo))

	def handle_search_change(self, text):
		objs = self.obj_map.get_objs_matching_regex(text)
		model = QStandardItemModel()
		root = model.invisibleRootItem()
		monospaced_font = get_monospaced_font()

		for x in objs:
			entry = QStandardItem(x)
			entry.setFont(monospaced_font)
			entry.setEditable(False)
			root.appendRow(entry)
		self.search_result.setModel(model)

	def handle_search_result_click(self, val):		
		obj = val.data()
		obj_info = self.obj_map.all_objs[obj]
		self.text_box.setPlainText(obj_info.get_description())

	def handle_search_result_double_click(self, val):
		obj = val.data()
		window = ObjectFileView(obj)
		window.show()
		# prevents window from being GC'd and closed.
		self.childen_windows.append(window)

	def handle_tree_double_click(self, val):
		if not(val.parent().isValid()):
			return
		obj = val.data()
		window = ObjectFileView(obj)
		window.show()
		# prevents window from being GC'd and closed.
		self.childen_windows.append(window)



map_file = load_obj_map_file(os.path.join(get_jak_path(), "decompiler_out", "obj.txt"))

app = QApplication([])
app.setStyle('Windows')
window = ObjectFileBrowser(map_file)
window.show()
app.exec_()