#!/usr/bin/env python3 # /// script # dependencies = [ # "duckdb>=1.5.3,<2", # "tqdm>=4.67,<5", # "unicorn==2.1.4", # ] # /// # SPDX-License-Identifier: Unlicense # # This example is released into the public domain under the Unlicense # (https://unlicense.org/). You may copy, modify, and use it for any purpose, # with or without attribution. from __future__ import annotations import argparse import json import tempfile import urllib.request from dataclasses import dataclass from pathlib import Path import duckdb from tqdm.auto import tqdm import unicorn import unicorn.x86_const as x86_const DEFAULT_BASE_URL = "https://binit.wazab.in" FILES = ("instructions.parquet", "test_cases.parquet", "test_results.parquet") MASK64 = (1 << 64) - 1 ENTRY_ADDR = 0x666666660000 PAGE_SIZE = 0x1000 MEM0_ADDR = 0x666666010100 MEM0_SIZE = 8 MEM0_PAGE_ADDR = MEM0_ADDR & ~(PAGE_SIZE - 1) CONTROL_FLOW_TARGET_ADDR = 0x666666661042 CONTROL_FLOW_TARGET_PAGE_ADDR = CONTROL_FLOW_TARGET_ADDR & ~(PAGE_SIZE - 1) SCALAR_REGISTER_KEYS = ("rax", "rbx", "rcx", "rdx", "r8", "r9", "rbp", "rsp") STATE_KEYS = (*SCALAR_REGISTER_KEYS, "flag", "mem0_value") STATUS_FLAG_BITS = {"CF": 0, "PF": 2, "AF": 4, "ZF": 6, "SF": 7, "OF": 11} SCALAR_REG_IDS = { name: getattr(x86_const, f"UC_X86_REG_{name.upper()}") for name in SCALAR_REGISTER_KEYS } @dataclass class ExpectationRow: test_case_id: int instruction_id: int instruction: str opcode: str state_index: int initial_state: dict[str, int] expected_final_state: dict[str, int] | None expected_exception_kind: str | None undefined_flags: list[str] instruction_tags: list[str] total_states: int = 0 @dataclass class ExecutionResult: outcome_class: str final_state: dict[str, int] | None exception_detail: str | None = None backend_error: str | None = None class UnicornRunner: """Replays states on a single, reused emulator. Constructing a fresh ``unicorn.Uc`` and re-mapping memory for every state is by far the dominant cost (engine setup is in the millisecond range, versus microseconds to actually execute one instruction). We map the fixed pages once and, per state, only rewrite the registers, flags, and the (rare) memory operand. The instruction bytes are rewritten only when the opcode changes, which both avoids needless work and keeps Unicorn's translation cache warm across the many states that share a single test case. Reuse is safe because the dataset only touches the scalar registers, flags, and ``mem0`` (see ``STATE_KEYS``): every run rewrites all of those, so no state leaks between runs. Registers outside that set are never written by these instructions, so they stay zero as on a fresh engine. """ def __init__(self) -> None: emu = unicorn.Uc(unicorn.UC_ARCH_X86, unicorn.UC_MODE_64) for page in {ENTRY_ADDR, CONTROL_FLOW_TARGET_PAGE_ADDR, MEM0_PAGE_ADDR}: emu.mem_map(page, PAGE_SIZE) emu.mem_write(CONTROL_FLOW_TARGET_PAGE_ADDR, b"\x90" * PAGE_SIZE) self.emu = emu self._loaded: tuple[int, bytes] | None = None def run(self, row: ExpectationRow) -> ExecutionResult: try: insn_bytes = bytes.fromhex(row.opcode.strip()) except ValueError as exc: return ExecutionResult("backend_error", None, backend_error=str(exc)) emu = self.emu start = ENTRY_ADDR + PAGE_SIZE - len(insn_bytes) # Only touch code memory when the instruction changes; rewriting it # would invalidate Unicorn's translation cache for the block. if self._loaded != (start, insn_bytes): emu.mem_write(start, insn_bytes) self._loaded = (start, insn_bytes) try: for name, reg_id in SCALAR_REG_IDS.items(): emu.reg_write(reg_id, row.initial_state.get(name, 0) & MASK64) emu.reg_write( x86_const.UC_X86_REG_RFLAGS, row.initial_state.get("flag", 0) & MASK64, ) if uses_mem0(row): emu.mem_write( MEM0_ADDR, int_to_mem0(row.initial_state.get("mem0_value", 0)), ) emu.reg_write(x86_const.UC_X86_REG_RIP, start) emu.emu_start(start, start + len(insn_bytes), count=1) final_state = { name: emu.reg_read(reg_id) & MASK64 for name, reg_id in SCALAR_REG_IDS.items() } final_state["flag"] = emu.reg_read(x86_const.UC_X86_REG_RFLAGS) & MASK64 if uses_mem0(row): final_state["mem0_value"] = mem0_to_int( emu.mem_read(MEM0_ADDR, MEM0_SIZE) ) return ExecutionResult("normal", final_state) except unicorn.UcError as exc: return ExecutionResult( "exception", None, exception_detail=type(exc).__name__ ) def compare(row: ExpectationRow, actual: ExecutionResult) -> dict[str, object]: expected_outcome = "exception" if row.expected_exception_kind else "normal" if expected_outcome != actual.outcome_class: diff = { "expected_outcome": expected_outcome, "actual_outcome": actual.outcome_class, } if actual.backend_error: diff["backend_error"] = actual.backend_error return diff if expected_outcome == "exception": return {} if row.expected_final_state is None or actual.final_state is None: return { "expected_final_state": row.expected_final_state, "actual_final_state": actual.final_state, } ignored_flags = set(row.undefined_flags) diff = {} for name, expected in row.expected_final_state.items(): actual_value = actual.final_state.get(name) if actual_value is None: diff[name] = {"expected": hex(expected & MASK64), "actual": None} elif name == "flag": expected_flags = status_flags(expected, ignored_flags) actual_flags = status_flags(actual_value, ignored_flags) if expected_flags != actual_flags: diff[name] = { "expected": hex(expected_flags), "actual": hex(actual_flags), } elif (expected & MASK64) != (actual_value & MASK64): diff[name] = { "expected": hex(expected & MASK64), "actual": hex(actual_value & MASK64), } return diff def validate_row(row: ExpectationRow) -> str | None: bad_initial = sorted(set(row.initial_state) - set(STATE_KEYS)) if bad_initial: return f"unsupported initial_state keys: {', '.join(bad_initial)}" if row.expected_final_state is not None: bad_final = sorted(set(row.expected_final_state) - set(STATE_KEYS)) if bad_final: return f"unsupported expected_final_state keys: {', '.join(bad_final)}" return None def uses_mem0(row: ExpectationRow) -> bool: return ( "mem0_value" in row.initial_state or ( row.expected_final_state is not None and "mem0_value" in row.expected_final_state ) ) def int_to_mem0(value: int) -> bytes: return (value & MASK64).to_bytes(MEM0_SIZE, byteorder="little", signed=False) def mem0_to_int(value) -> int: return int.from_bytes(bytes(value[:MEM0_SIZE]), byteorder="little", signed=False) def status_flags(value: int, ignored_flags: set[str]) -> int: mask = 0 for name, bit in STATUS_FLAG_BITS.items(): if name not in ignored_flags: mask |= 1 << bit return value & mask def download_data(base_url: str, out_dir: Path) -> None: for name in FILES: url = f"{base_url.rstrip('/')}/data/{name}" urllib.request.urlretrieve(url, out_dir / name) def open_parquet_db(data_dir: Path) -> duckdb.DuckDBPyConnection: con = duckdb.connect(database=":memory:") for table in ("instructions", "test_cases", "test_results"): con.read_parquet(str(data_dir / f"{table}.parquet")).create_view(table) return con def iter_rows( con: duckdb.DuckDBPyConnection, limit: int | None, *, chunk_size: int, ): seen = 0 last_test_case_id = -1 while limit is None or seen < limit: case_rows = con.execute( """ SELECT tc.test_case_id, tc.instruction_id, tc.instruction, tc.opcode, tc.initial_states, i.undefined_flags, i.tags FROM test_cases tc JOIN instructions i ON i.instruction_id = tc.instruction_id WHERE NOT list_contains(i.tags, 'avx') AND tc.test_case_id > ? ORDER BY tc.test_case_id LIMIT ? """, [last_test_case_id, chunk_size], ).fetchall() if not case_rows: return cases = {} for r in case_rows: last_test_case_id = r[0] cases[r[0]] = { "instruction_id": r[1], "instruction": r[2], "opcode": r[3], "initial_states": json.loads(r[4] or "[]"), "undefined_flags": list(r[5] or []), "instruction_tags": list(r[6] or []), } result_rows = con.execute( """ SELECT test_case_id, state_index, exception_kind, final_state FROM test_results WHERE test_case_id = ANY(?) ORDER BY test_case_id, state_index """, [list(cases)], ).fetchall() for r in result_rows: if limit is not None and seen >= limit: return case = cases[r[0]] states = case["initial_states"] final_state = json.loads(r[3]) if r[3] else None seen += 1 yield ExpectationRow( test_case_id=r[0], instruction_id=case["instruction_id"], instruction=case["instruction"], opcode=case["opcode"], initial_state={k: int(v) for k, v in states[r[1]].items()}, state_index=r[1], expected_exception_kind=r[2], expected_final_state=( {k: int(v) for k, v in final_state.items()} if final_state is not None else None ), undefined_flags=case["undefined_flags"], instruction_tags=case["instruction_tags"], total_states=len(states), ) def main() -> int: parser = argparse.ArgumentParser() parser.add_argument("--base-url", default=DEFAULT_BASE_URL) parser.add_argument("--limit", type=int, default=None) parser.add_argument("--chunk-size", type=int, default=256) args = parser.parse_args() if args.chunk_size <= 0: parser.error("--chunk-size must be greater than 0") processed = mismatches = backend_errors = 0 with tempfile.TemporaryDirectory(prefix="binit-parquet-") as tmp: data_dir = Path(tmp) download_data(args.base_url, data_dir) con = open_parquet_db(data_dir) runner = UnicornRunner() # Once a mnemonic mismatches, the rest of its states almost always do # too, so we record the first failure and skip the mnemonic afterwards. failed_instructions: set[int] = set() rows = iter_rows(con, args.limit, chunk_size=args.chunk_size) for row in tqdm(rows, total=args.limit, unit="state", dynamic_ncols=True): if row.instruction_id in failed_instructions: continue processed += 1 error = validate_row(row) actual = runner.run(row) if error is None else None if actual is None: tqdm.write( f"contract error: {row.test_case_id}:{row.state_index}: {error}" ) mismatches += 1 failed_instructions.add(row.instruction_id) continue if actual.outcome_class == "backend_error": backend_errors += 1 failed_instructions.add(row.instruction_id) diff = compare(row, actual) if diff: mismatches += 1 failed_instructions.add(row.instruction_id) tqdm.write( f"mismatch {row.test_case_id}:{row.state_index} " f"{row.instruction}: {diff}" ) print( f"processed={processed} mismatches={mismatches} backend_errors={backend_errors}" ) return 1 if mismatches or backend_errors else 0 if __name__ == "__main__": raise SystemExit(main())