mirror of
https://github.com/pret/pmd-sky.git
synced 2026-03-22 01:35:35 -05:00
280 lines
11 KiB
Python
280 lines
11 KiB
Python
import os
|
|
import argparse
|
|
|
|
from write_inc_file import write_inc_file
|
|
|
|
# Use this script to extract a function out of assembly and create a blank C function as a placeholder.
|
|
# If the function is in the middle, the assembly file is split in two. If the function is at either end, it is added to the next/previous source file.
|
|
# This is intended for use after decompiling a function, so that you can place the decompiled C code in the placeholder after running the script.
|
|
# Since this script updates files, it's recommended to have no uncommitted changes when running this script, in case something goes wrong and needs a revert.
|
|
# Example usage (auto-generated file name): python extract_function.py overlay_29_0234EC38 "u8 ov29_0234FCA8(u8 arg0)"
|
|
# Example usage (custom file name): python extract_function.py overlay_29_0234EC38 "u8 ov29_0234FCA8(u8 arg0)" -f my_new_file
|
|
# Example usage (nonmatching): python extract_function.py overlay_29_0234EC38 "u8 ov29_0234FCA8(u8 arg0)" -n
|
|
|
|
parser = argparse.ArgumentParser()
|
|
|
|
parser.add_argument('asm_file')
|
|
parser.add_argument('function_header')
|
|
parser.add_argument('-f', '--extract_file_name')
|
|
parser.add_argument('-n', '--nonmatching', action='store_true')
|
|
|
|
args = parser.parse_args()
|
|
function_location = args.asm_file
|
|
function_header = args.function_header
|
|
extract_file_name = args.extract_file_name
|
|
nonmatching = args.nonmatching
|
|
|
|
if function_location.endswith('.s'):
|
|
function_location = function_location[:-2]
|
|
if function_location.startswith('./asm/'):
|
|
function_location = function_location[6:]
|
|
if function_header.endswith(';'):
|
|
function_header = function_header[:-1]
|
|
|
|
is_itcm = function_location.startswith('itcm')
|
|
|
|
# Extract the function name from the function header argument.
|
|
left_parentheses_index = function_header.find('(')
|
|
if left_parentheses_index >= 0:
|
|
function_name = function_header[function_header.rfind(' ', None, left_parentheses_index) + 1 : left_parentheses_index]
|
|
else:
|
|
print('Unable to find function name in function header.')
|
|
exit(1)
|
|
|
|
if os.getcwd().endswith('extract_function'):
|
|
os.chdir(os.path.join('..', '..'))
|
|
|
|
ASM_FOLDER = 'asm'
|
|
INCLUDE_FOLDER = os.path.join(ASM_FOLDER, 'include')
|
|
HEADER_FOLDER = 'include'
|
|
SRC_FOLDER = 'src'
|
|
LSF_FILE_PATH = 'main.lsf'
|
|
|
|
original_file_path = os.path.join(ASM_FOLDER, f'{function_location}.s')
|
|
original_inc_path = os.path.join(INCLUDE_FOLDER, f'{function_location}.inc')
|
|
with open(original_file_path, 'r') as original_file:
|
|
original_lines = original_file.readlines()
|
|
|
|
if function_location.startswith('main'):
|
|
file_prefix = 'main_'
|
|
elif function_location.startswith('itcm'):
|
|
file_prefix = 'itcm_'
|
|
else:
|
|
file_prefix = function_location[:len('overlay_00_')]
|
|
if file_prefix[-1] != '_':
|
|
file_prefix += '_'
|
|
|
|
function_start_line = None
|
|
function_end_line = None
|
|
extract_function_address = None
|
|
new_file_address = None
|
|
ADDRESS_FIND = '; 0x'
|
|
ARM_FUNC_START = '\tarm_func_start '
|
|
first_function_start_line = None
|
|
|
|
def get_line_address(line: str):
|
|
return line[line.index(ADDRESS_FIND) + len(ADDRESS_FIND) : -1]
|
|
|
|
# Find the start and end of the function within the ASM file.
|
|
for i, line in enumerate(original_lines):
|
|
if first_function_start_line is None and line.startswith(ARM_FUNC_START):
|
|
first_function_start_line = i
|
|
if line.strip() == f'{ARM_FUNC_START}{function_name}'.strip():
|
|
function_start_line = i
|
|
elif line.strip() == f'arm_func_end {function_name}'.strip():
|
|
function_end_line = i
|
|
|
|
if function_start_line is not None and extract_function_address is None and ADDRESS_FIND in line:
|
|
extract_function_address = get_line_address(line)
|
|
if function_end_line is not None and ADDRESS_FIND in line:
|
|
new_file_address = get_line_address(line)
|
|
break
|
|
|
|
if function_start_line is None or function_end_line is None or extract_function_address is None:
|
|
print(f'Failed to find function. Start line {function_start_line}, end line {function_end_line}, extract address {extract_function_address}, new address {new_file_address}.')
|
|
exit(1)
|
|
|
|
new_asm_base_name = f"{file_prefix}{new_file_address}"
|
|
|
|
remove_orig_file = first_function_start_line == function_start_line
|
|
include_new_asm_file = new_file_address is not None
|
|
|
|
new_inc_file_name = f"{new_asm_base_name}.inc"
|
|
section_name = 'text'
|
|
if is_itcm:
|
|
section_name = 'itcm,4,1,4'
|
|
new_asm_header = f"""\t.include "asm/macros.inc"
|
|
\t.include "{new_inc_file_name}"
|
|
|
|
\t.{section_name}
|
|
"""
|
|
new_asm_name = f'{new_asm_base_name}.s'
|
|
|
|
new_asm_lines = original_lines[function_end_line + 1:]
|
|
original_asm_lines = original_lines[:function_start_line - 1]
|
|
|
|
with open(LSF_FILE_PATH, 'r') as lsf_file:
|
|
lsf_lines = lsf_file.readlines()
|
|
|
|
if extract_file_name is None:
|
|
extract_file_name = f'{file_prefix}{extract_function_address}'
|
|
|
|
# If needed, add the extracted function's new .o file to main.lsf.
|
|
merge_prev_file = None
|
|
merge_next_file = None
|
|
lsf_suffix = ''
|
|
if is_itcm:
|
|
lsf_suffix = ' (.itcm)'
|
|
SRC_LSF_PREFIX = '\tObject src/'
|
|
for i, line in enumerate(lsf_lines):
|
|
if line.endswith(f'{function_location}.o{lsf_suffix}\n'):
|
|
if remove_orig_file:
|
|
lsf_lines[i] = ''
|
|
prev_line = lsf_lines[i - 1]
|
|
if prev_line.startswith(SRC_LSF_PREFIX):
|
|
merge_prev_file = prev_line[len(SRC_LSF_PREFIX):]
|
|
if not include_new_asm_file and merge_prev_file is None:
|
|
next_line = lsf_lines[i + 1]
|
|
if next_line.startswith(SRC_LSF_PREFIX):
|
|
merge_next_file = next_line[len(SRC_LSF_PREFIX):]
|
|
if merge_prev_file is None and merge_next_file is None:
|
|
lsf_lines[i] += f'\tObject src/{extract_file_name}.o{lsf_suffix}\n'
|
|
if include_new_asm_file:
|
|
lsf_lines[i] += f'\tObject asm/{new_asm_base_name}.o{lsf_suffix}\n'
|
|
break
|
|
|
|
def trim_merge_file_name(file_name: str) -> str:
|
|
if file_name is None:
|
|
return None
|
|
|
|
line_end_index = len('.o\n')
|
|
if file_name.endswith('.o (.itcm)\n'):
|
|
line_end_index = len('.o (.itcm)\n')
|
|
return file_name[:-line_end_index]
|
|
|
|
merge_prev_file = trim_merge_file_name(merge_prev_file)
|
|
merge_next_file = trim_merge_file_name(merge_next_file)
|
|
|
|
print('Updating', LSF_FILE_PATH)
|
|
with open(LSF_FILE_PATH, 'w') as lsf_file:
|
|
lsf_file.writelines(lsf_lines)
|
|
|
|
if remove_orig_file:
|
|
print('Removing', original_file_path)
|
|
os.remove(original_file_path)
|
|
print('Removing', original_inc_path)
|
|
os.remove(original_inc_path)
|
|
else:
|
|
print('Updating', original_file_path)
|
|
with open(original_file_path, 'w') as original_file:
|
|
original_file.writelines(original_asm_lines)
|
|
print('Updating', original_inc_path)
|
|
write_inc_file(original_asm_lines, original_inc_path)
|
|
|
|
if include_new_asm_file:
|
|
new_asm_file_path = os.path.join(ASM_FOLDER, new_asm_name)
|
|
print('Creating', new_asm_file_path)
|
|
with open(os.path.join(ASM_FOLDER, new_asm_name), 'w') as new_asm_file:
|
|
new_asm_file.write(new_asm_header)
|
|
new_asm_file.writelines(new_asm_lines)
|
|
|
|
new_asm_inc_path = os.path.join(INCLUDE_FOLDER, f'{new_asm_base_name}.inc')
|
|
print('Creating', new_asm_inc_path)
|
|
write_inc_file(new_asm_lines, new_asm_inc_path)
|
|
|
|
function_body = f"""{function_header}
|
|
{{
|
|
|
|
}}"""
|
|
|
|
if nonmatching:
|
|
asm_lines = original_lines[function_start_line + 2 : function_end_line - 1]
|
|
modified_asm_lines = []
|
|
for line in asm_lines:
|
|
# .align is not needed in embedded ASM
|
|
if '.align' in line:
|
|
continue
|
|
|
|
# Replace some register mnemonics with numbered registers. These don't work when ASM is embedded in C.
|
|
if not line.startswith('bl') and line.startswith('\t'):
|
|
space_index = line.find(' ')
|
|
line = line[:space_index] + line[space_index:].replace('sb', 'r9').replace('sl', 'r10').replace('fp', 'r11')
|
|
# Replacing "sl" can corrupt lsl instructions, so correct these.
|
|
line = line.replace('lr10', 'lsl')
|
|
|
|
semicolon_index = line.find(';')
|
|
if semicolon_index >= 0:
|
|
if 'jump table' in line or 'case' in line:
|
|
# Jump tables have comments, but semicolons are not recognized as comment markers in embedded ASM, so remove them.
|
|
line = line[:semicolon_index - 1] + '\n'
|
|
else:
|
|
# Replace word values at the end of the function.
|
|
# They are conveniently already included within the ASM as comments.
|
|
line = line[:line.find('_')] + line[semicolon_index + 2:]
|
|
|
|
modified_asm_lines.append(line)
|
|
|
|
asm_string = str.join('', modified_asm_lines)
|
|
|
|
function_body = f"""#ifdef NONMATCHING
|
|
{function_body}
|
|
#else
|
|
asm {function_header}
|
|
{{
|
|
{asm_string}}}
|
|
#endif"""
|
|
|
|
# Add the extracted function to a .h and .c file.
|
|
# If there is an existing C file adjacent to the extracted function, add the function to that file.
|
|
# Otherwise, make a new set of files.
|
|
if merge_prev_file:
|
|
header_file_path = os.path.join(HEADER_FOLDER, f'{merge_prev_file}.h')
|
|
if os.path.exists(header_file_path):
|
|
with open(header_file_path, 'r') as header_file:
|
|
header_lines = header_file.readlines()
|
|
for i, line in reversed(list(enumerate(header_lines))):
|
|
if line != '\n' and not line.startswith('#'):
|
|
header_lines[i] += f'{function_header};\n'
|
|
break
|
|
print('Updating', header_file_path)
|
|
with open(header_file_path, 'w') as header_file:
|
|
header_file.writelines(header_lines)
|
|
|
|
src_file_path = os.path.join(SRC_FOLDER, f'{merge_prev_file}.c')
|
|
print('Updating', src_file_path)
|
|
with open(src_file_path, 'a') as src_file:
|
|
src_file.write(f'\n{function_body}\n')
|
|
|
|
elif merge_next_file:
|
|
def add_in_body_start(file_path: str, add_content: str):
|
|
with open(file_path, 'r') as add_file:
|
|
file_lines = add_file.readlines()
|
|
for i, line in enumerate(file_lines):
|
|
if line != '\n' and not line.startswith('#'):
|
|
file_lines[i] = add_content + file_lines[i]
|
|
break
|
|
with open(file_path, 'w') as add_file:
|
|
add_file.writelines(file_lines)
|
|
|
|
header_file_path = os.path.join(HEADER_FOLDER, f'{merge_next_file}.h')
|
|
print('Updating', header_file_path)
|
|
add_in_body_start(header_file_path, f'{function_header};\n')
|
|
|
|
src_file_path = os.path.join(SRC_FOLDER, f'{merge_next_file}.c')
|
|
print('Updating', src_file_path)
|
|
add_in_body_start(src_file_path, f'{function_body}\n\n')
|
|
|
|
else:
|
|
header_file_path = os.path.join(HEADER_FOLDER, f'{extract_file_name}.h')
|
|
print('Creating', header_file_path)
|
|
with open(header_file_path, 'w') as new_header_file:
|
|
file_guard = f'PMDSKY_{extract_file_name.upper()}_H'
|
|
new_header_file.write(f'#ifndef {file_guard}\n#define {file_guard}\n\n')
|
|
new_header_file.write(f'{function_header};\n\n')
|
|
new_header_file.write(f'#endif //{file_guard}\n')
|
|
|
|
src_file_path = os.path.join(SRC_FOLDER, f'{extract_file_name}.c')
|
|
print('Creating', src_file_path)
|
|
with open(src_file_path, 'w') as new_src_file:
|
|
new_src_file.write(f'#include "{extract_file_name}.h\"\n\n{function_body}\n')
|