#!/usr/bin/env python3 """ Audit Comment Tool Adds canonical `; CPU:$XXXX PRG:XXXX AA BB` comments to disassembly source files. Uses nesdis.py tracer as ground truth for instruction bytes and addresses. """ import sys import re import argparse from pathlib import Path sys.path.insert(0, str(Path(__file__).parent)) import nesdis # Canonical comment format CANONICAL_FMT = "; CPU:${cpu:04X} PRG:{prg:04X} {bytes}" # Line classification patterns INSTRUCTION_RE = re.compile(r'^(\s*)([A-Za-z_][A-Za-z0-9_]*)\s+(.+?)(?:\s*;.*)?$') LABEL_RE = re.compile(r'^(\s*)([A-Za-z_][A-Za-z0-9_]*):\s*(?:;.*)?$') BYTE_RE = re.compile(r'^(\s*)\.byte\s+(.+?)(?:\s*;.*)?$', re.I) WORD_RE = re.compile(r'^(\s*)\.word\s+(.+?)(?:\s*;.*)?$', re.I) DWORD_RE = re.compile(r'^(\s*)\.dword\s+(.+?)(?:\s*;.*)?$', re.I) ORG_RE = re.compile(r'^(\s*)\.org\s+\$([0-9A-Fa-f]+)', re.I) SEGMENT_RE = re.compile(r'^\.segment\s+"PRG"', re.I) INCLUDE_RE = re.compile(r'^\.include\s+') SET_CPU_RE = re.compile(r'^\.setcpu\s+') EQU_RE = re.compile(r'^(\s*)([A-Za-z_][A-Za-z0-9_]*)\s*=\s*', re.I) COMMENT_RE = re.compile(r'^\s*;') BLANK_RE = re.compile(r'^\s*$') # Existing address comment in source: ; $C036: A9 E2 SRC_ADDR_RE = re.compile(r';\s*\$([0-9A-Fa-f]{4})\s*:\s*([0-9A-Fa-f\s]+)') # Existing audit comment patterns to replace - only match canonical format with byte lists at end of line # Canonical: ; CPU:$XXXX PRG:XXXX AA BB (byte list goes to end of line or before --) # Source address comments: ; $XXXX: description or ; PRG:$XXXX AA BB description EXISTING_AUDIT_RE = re.compile( r'\s*;\s*(?:' r'CPU[\s:]\$?[\dA-F]+\s+PRG[\s:]\$?[\dA-F]+\s+(?:[\dA-F]{2}\s*)+\s*(?:--.*)?$|' # CPU:$XXXX PRG:XXXX AA BB r'PRG[\s:]\$?[\dA-F]+\s*\|\s*CPU[\s:]\$?[\dA-F]+\s*\|\s*(?:[\dA-F]{2}\s*)+\s*$|' # PRG $XXXX | CPU $XXXX | AA BB r'PRG[\s:]\$?[\dA-F]+\s+(?:[\dA-F]{2}\s*)+\s*$|' # PRG:$XXXX AA BB (at end of line) r'PRG[\s:]\$?[\dA-F]+\s*/\s*CPU[\s:]\$?[\dA-F]+:\s*(?:[\dA-F]{2}\s*)+\s*$|' # PRG $XXXX / CPU $XXXX: AA BB r'\$[\dA-F]{4}:\s*(?:[\dA-F]{2}\s*)+\s*$|' # $XXXX: AA BB (short form, at end) r'PRG[\s:]\$?[\dA-F]+\s+Bytes:\s*(?:[\dA-F]{2}\s*)+\s*$|' # PRG:$XXXX Bytes:AA BB r'ROM\s\$[\dA-F]{6}:\s*(?:[\dA-F]{2}\s*)+\s*$' # ROM $XXXXXX: AA BB r')' ) def parse_byte_operands(operand_str): """Parse .byte/.word/.dword operands into list of byte values.""" # Strip trailing comment first operand_str = operand_str.split(';')[0] operands = [] for part in operand_str.split(','): part = part.strip() if not part: continue # Handle $XX or #$XX if part.startswith('$'): val = int(part[1:], 16) operands.append(val & 0xFF) elif part.startswith('#$'): val = int(part[2:], 16) operands.append(val & 0xFF) elif part.isdigit(): operands.append(int(part) & 0xFF) else: # Symbol or expression - can't resolve, placeholder operands.append(None) return operands def count_directive_bytes(line): """Count bytes generated by .byte/.word/.dword directive.""" m = re.match(r'^\s*\.(byte|word|dword)\s+(.+)', line, re.I) if not m: return 0 directive = m.group(1).lower() operands = parse_byte_operands(m.group(2)) size = {'byte': 1, 'word': 2, 'dword': 4}[directive] # For .word/.dword, each operand generates 2/4 bytes # For .byte, each operand generates 1 byte return len(operands) * size def build_trace_map(prg_path, base=0xC000, undoc=False): """Run tracer and return dict mapping PRG offset -> trace info.""" prg = Path(prg_path).read_bytes() nesdis.set_undocumented(undoc) vo = len(prg) - 6 rst = prg[vo + 2] | (prg[vo + 3] << 8) nmi = prg[vo] | (prg[vo + 1] << 8) irq = prg[vo + 4] | (prg[vo + 5] << 8) insns, invalid, indirect, targets, overlap = nesdis.trace(prg, base, [rst, nmi, irq]) nesdis.set_undocumented(False) trace_map = {} for off, insn in insns.items(): cpu = base + off raw_bytes = prg[off:off + insn.length] bytes_str = " ".join(f"{b:02X}" for b in raw_bytes) text = f"{insn.mnemonic} {insn.text}".strip() trace_map[off] = { 'cpu': cpu, 'prg': off, 'bytes': bytes_str, 'length': insn.length, 'text': text, 'mnemonic': insn.mnemonic, 'operand': insn.text, 'is_code': True } # Data regions from scan_orphans - store as 16-byte chunks regions = nesdis.scan_orphans(prg, base, insns, 1) for start, length in regions: for i in range(start, start + length, 16): chunk = prg[i:min(i + 16, start + length)] bytes_str = " ".join(f"{b:02X}" for b in chunk) trace_map[i] = { 'cpu': base + i, 'prg': i, 'bytes': bytes_str, 'length': len(chunk), 'is_data': True } return trace_map, prg def parse_source_for_offsets(source_path, trace_map, base=0xC000): """Parse source file, track PRG offset, and generate annotated lines.""" lines = Path(source_path).read_text().split('\n') annotated = [] prg_offset = 0 in_prg_segment = False org_set = False # First pass: find .org to establish base offset for line in lines: m = ORG_RE.match(line) if m: org_addr = int(m.group(2), 16) prg_offset = org_addr - base org_set = True break if not org_set: prg_offset = 0 for line in lines: stripped = line.strip() # Skip blank lines and pure comments if BLANK_RE.match(line) or COMMENT_RE.match(line): annotated.append((line, None)) continue # .segment "PRG" - start tracking if SEGMENT_RE.match(line): in_prg_segment = True annotated.append((line, None)) continue # .org - reset offset m = ORG_RE.match(line) if m: org_addr = int(m.group(2), 16) prg_offset = org_addr - base annotated.append((line, None)) continue # .include, .setcpu, symbol definitions - no bytes if INCLUDE_RE.match(line) or SET_CPU_RE.match(line) or EQU_RE.match(line): annotated.append((line, None)) continue # Label - no bytes if LABEL_RE.match(stripped): annotated.append((line, None)) continue # Check for existing address comment to sync offset: ; $C036: A9 E2 # Only sync if address looks like a CPU address (>= $8000), not a PRG offset addr_match = SRC_ADDR_RE.search(line) if addr_match: cpu_addr = int(addr_match.group(1), 16) # Only treat as CPU address if >= $8000 (typical PRG mapping) # PRG offsets in "short" format are typically $0000-$3FFF if cpu_addr >= 0x8000: prg_offset = cpu_addr - base # .byte, .word, .dword - data directives m = BYTE_RE.match(line) or WORD_RE.match(line) or DWORD_RE.match(line) if m: byte_count = count_directive_bytes(line) # Try to get bytes from existing source comment first # Formats: ; $XXXX: AA BB or ; PRG:$XXXX AA BB bytes_str = "" src_bytes_match = re.search(r';\s*(?:\$([0-9A-Fa-f]{4})\s*:\s*|PRG:\$([0-9A-Fa-f]{4})\s+)([0-9A-Fa-f\s]+)', line) if src_bytes_match: # Group 1 or 2 has the address, group 3 has the bytes bytes_part = src_bytes_match.group(3).strip() byte_vals = bytes_part.split() if byte_vals: bytes_str = " ".join(f"{int(b, 16):02X}" for b in byte_vals) # Fallback: parse operands if they're hex values if not bytes_str: operands = parse_byte_operands(m.group(2)) bytes_str = " ".join(f"{b:02X}" for b in operands if b is not None) # Fallback to trace map if not bytes_str: trace_info = trace_map.get(prg_offset) if trace_info: # For data, trace_info has 16-byte chunks; take what we need tb = trace_info['bytes'].split() bytes_str = " ".join(tb[:byte_count]) else: bytes_str = "??" canonical = CANONICAL_FMT.format( cpu=base + prg_offset, prg=prg_offset, bytes=bytes_str ) # Remove existing audit comment clean_line = EXISTING_AUDIT_RE.sub('', line).rstrip() new_line = f"{clean_line} {canonical}" annotated.append((new_line, prg_offset)) prg_offset += byte_count continue # Instruction m = INSTRUCTION_RE.match(line) if m and in_prg_segment: indent, mnemonic, operand = m.groups() operand = operand.strip() # Check for existing address comment with bytes: ; $C036: A9 E2 addr_match = SRC_ADDR_RE.search(line) bytes_str = "??" length = 2 if addr_match: cpu_addr = int(addr_match.group(1), 16) # Only sync PRG offset if address looks like a CPU address (>= $8000) # PRG offsets in "short" format are typically $0000-$3FFF if cpu_addr >= 0x8000: # Sync PRG offset to this address prg_offset = cpu_addr - base # Parse bytes from comment regardless bytes_part = addr_match.group(2).strip() byte_vals = bytes_part.split() if byte_vals: bytes_str = " ".join(f"{int(b, 16):02X}" for b in byte_vals) length = len(byte_vals) else: bytes_str = "??" else: # Look up in trace_map trace_info = trace_map.get(prg_offset) if trace_info and trace_info.get('is_code'): bytes_str = trace_info['bytes'] length = trace_info['length'] else: # Fallback: estimate length length = 2 bytes_str = "??" canonical = CANONICAL_FMT.format( cpu=base + prg_offset, prg=prg_offset, bytes=bytes_str ) # Remove existing audit comment clean_line = EXISTING_AUDIT_RE.sub('', line).rstrip() new_line = f"{clean_line} {canonical}" annotated.append((new_line, prg_offset)) # Advance offset prg_offset += length continue # Unknown line type - keep as is annotated.append((line, None)) return annotated def main(): ap = argparse.ArgumentParser(description="Add canonical audit comments to disassembly source") ap.add_argument("project", help="Project directory name (e.g., 'Baseball (UE)')") ap.add_argument("--base", default="0xC000", help="PRG base address") ap.add_argument("--undoc", action="store_true", help="Include undocumented opcodes in trace") ap.add_argument("--dry-run", action="store_true", help="Show changes without writing") args = ap.parse_args() root = Path(__file__).parent.parent project_dir = root / args.project copilot = project_dir / "copilot" # Find source file source_file = None for name in ("prg.s", "main.s"): if (copilot / name).exists(): source_file = copilot / name break if not source_file: print(f"Error: No prg.s or main.s found in {copilot}") return 1 # Find PRG binary prg_file = copilot / "prg.bin" if not prg_file.exists(): # Try to carve from rom.nes rom = copilot / "rom.nes" if rom.exists(): print("Carving PRG from rom.nes...") rom_data = rom.read_bytes() # iNES header is 16 bytes # PRG size = rom[4] * 16384 prg_size = rom_data[4] * 16384 chr_size = rom_data[5] * 8192 has_trainer = bool(rom_data[6] & 4) prg_start = 16 + (512 if has_trainer else 0) prg_data = rom_data[prg_start:prg_start + prg_size] prg_file = copilot / "prg.bin" prg_file.write_bytes(prg_data) print(f" Carved {len(prg_data)} bytes to prg.bin") else: print(f"Error: No PRG binary found in {copilot}") return 1 base = int(args.base, 16) print(f"Building trace map for {args.project}...") trace_map, prg = build_trace_map(prg_file, base, args.undoc) print(f" {len(trace_map)} trace entries") print(f"Parsing source and adding comments...") annotated = parse_source_for_offsets(source_file, trace_map, base) # Count changes old_lines = source_file.read_text().split('\n') changes = 0 for old, (new_line, _) in zip(old_lines, annotated): if new_line != old: changes += 1 print(f" {changes} lines would be modified") if args.dry_run: # Show first 50 changes shown = 0 for i, (old, (new_line, _)) in enumerate(zip(old_lines, annotated)): if old != new_line and shown < 50: print(f"Line {i+1}:") print(f" OLD: {old}") print(f" NEW: {new_line}") shown += 1 return 0 # Write back output = "\n".join(a[0] for a in annotated) + "\n" source_file.write_text(output) print(f"Written to {source_file}") # Verify build import subprocess result = subprocess.run(['make'], cwd=copilot, capture_output=True, text=True) if result.returncode == 0: print("Build: PASS (byte-identical)") else: print("Build: FAIL") print(result.stdout[-500:]) print(result.stderr[-500:]) return 1 return 0 if __name__ == "__main__": sys.exit(main())