126 lines
3.4 KiB
Python
126 lines
3.4 KiB
Python
import os
|
|
import re
|
|
import argparse
|
|
import pydantic
|
|
|
|
from pathlib import Path
|
|
from typing import Protocol, cast
|
|
|
|
|
|
CONFIG_ROOT = Path(__file__).parent.parent.parent.resolve() / "config"
|
|
|
|
# enumerate languages from config dir
|
|
LANGUAGES = [Path(f.path).name for f in os.scandir(CONFIG_ROOT) if f.is_dir()]
|
|
|
|
RE_INSTR = re.compile(r"^\t(?:[^.]|\.p2align \d+)")
|
|
|
|
|
|
class InstructionPatch(pydantic.RootModel[tuple[int, str, str]]):
|
|
root: tuple[int, str, str]
|
|
|
|
@property
|
|
def instr_no(self):
|
|
return self.root[0]
|
|
|
|
@property
|
|
def instr_org(self):
|
|
return self.root[1]
|
|
|
|
@property
|
|
def instr_patch(self):
|
|
return self.root[2]
|
|
|
|
|
|
class ASMPatch(pydantic.RootModel[dict[str, list[InstructionPatch]]]):
|
|
root: dict[str, list[InstructionPatch]]
|
|
|
|
|
|
class PatchDB(pydantic.RootModel[dict[str, ASMPatch]]):
|
|
root: dict[str, ASMPatch]
|
|
|
|
|
|
def fix_asm(asm_file: Path, asm_patch: ASMPatch):
|
|
lines = asm_file.read_text().splitlines()
|
|
|
|
def find_func(func: str):
|
|
for i, line in enumerate(lines):
|
|
if re.match(rf"^{func}:", line):
|
|
return i
|
|
|
|
return -1
|
|
|
|
def find_line(start: int, num: int):
|
|
offset = 0
|
|
instr_no = -1
|
|
for line in lines[start:]:
|
|
if RE_INSTR.match(line):
|
|
instr_no += 1
|
|
|
|
if instr_no == num:
|
|
return start + offset
|
|
|
|
offset += 1
|
|
|
|
raise RuntimeError("cannot find instruction!")
|
|
|
|
for func, patch_lst in asm_patch.root.items():
|
|
n = find_func(func)
|
|
|
|
if n == -1:
|
|
print(f"WARNING: cannot find function {func} in {asm_file.name}")
|
|
continue
|
|
|
|
for patch in patch_lst:
|
|
line_no = find_line(n, patch.instr_no)
|
|
line_org = lines[line_no]
|
|
|
|
line_org_clean = re.sub(r"\s+", " ", line_org).strip()
|
|
instr_org_clean = re.sub(r"\s+", " ", patch.instr_org).strip()
|
|
instr_patch_clean = re.sub(r"\s+", " ", patch.instr_patch).strip()
|
|
|
|
if line_org_clean != instr_org_clean:
|
|
print(f"WARNING: wrong line: {asm_file.name}:{func}:{line_no + 1}: {line_org} != {patch.instr_org}")
|
|
continue
|
|
|
|
lines[line_no] = f"\t{instr_patch_clean}"
|
|
|
|
asm_file.write_text("\n".join(lines))
|
|
|
|
|
|
def main():
|
|
class ArgsProtocol(Protocol):
|
|
language: str
|
|
asm_file: Path
|
|
|
|
parser = argparse.ArgumentParser(description="apply asm patches to assembly files")
|
|
parser.add_argument("language", type=str, choices=LANGUAGES, help="language of the asm that is being patched")
|
|
parser.add_argument("asm_file", type=Path, help="generated assembly file to patch (relative to build dir)")
|
|
|
|
args = cast(ArgsProtocol, parser.parse_args())
|
|
|
|
asm_file = CONFIG_ROOT / args.language / args.asm_file
|
|
|
|
if not asm_file.exists():
|
|
print(f"ERROR: cannot find assembly file {asm_file}")
|
|
exit(1)
|
|
|
|
patch_db_path = CONFIG_ROOT / args.language / "asm_patches.json"
|
|
|
|
if not patch_db_path.exists():
|
|
print("ERROR: patch db file doesn't exist. aborting patching asm")
|
|
exit(1)
|
|
|
|
patch_db = PatchDB.model_validate_json(patch_db_path.read_text())
|
|
|
|
if asm_file.name not in patch_db.root:
|
|
print(f"ERROR: no patches found for file {asm_file.name}")
|
|
exit(1)
|
|
|
|
asm_patch = patch_db.root[asm_file.name]
|
|
|
|
fix_asm(asm_file, asm_patch)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|