"""ANSI formatting helpers for monitor REPL output. Provides functions to format SimEvent, snapshot state, and step results with ANSI colour codes for readable terminal display. Key features: - NO_COLOUR flag to disable all colouring (for testing and pipes) - Distinct colours for event types, component IDs, data values, errors, and time - Compact and detailed formatting modes for different use cases """ from __future__ import annotations from emu.events import ( CellWritten, DeferredRead, DeferredSatisfied, Emitted, Executed, IRAMWritten, Matched, ResultSent, SimEvent, TokenReceived, ) from monitor.commands import StepResult from monitor.snapshot import PESnapshot, SMSnapshot, StateSnapshot # Global flag to disable all colouring (for testing and pipes) NO_COLOUR = False # ANSI colour codes COLOURS = { "reset": "\033[0m", "bright_cyan": "\033[96m", # Event types "bright_yellow": "\033[93m", # PE/SM identifiers "white": "\033[97m", # Data values "bright_red": "\033[91m", # Errors "bright_green": "\033[92m", # Sim time "dim": "\033[2m", # Dim (for secondary info) } def colour(text: str, code: str) -> str: """Wrap text in ANSI escape code, respecting NO_COLOUR flag. Args: text: Text to colour code: Colour code key (from COLOURS dict) Returns: Coloured text if NO_COLOUR is False, otherwise plain text """ if NO_COLOUR: return text return f"{COLOURS[code]}{text}{COLOURS['reset']}" def format_event(event: SimEvent) -> str: """Format a single SimEvent as a coloured one-line string. Shows event type and key fields in distinct colours. Args: event: SimEvent instance Returns: Formatted event string with ANSI colours """ time_str = colour(f"[{event.time:.1f}]", "bright_green") component_str = colour(event.component, "bright_yellow") match event: case TokenReceived(token=token): event_type = colour("TokenReceived", "bright_cyan") return f"{time_str} {event_type} {component_str}: {token}" case Matched(left=left, right=right, act_id=act_id, frame_id=frame_id): event_type = colour("Matched", "bright_cyan") return ( f"{time_str} {event_type} {component_str}: " f"left={colour(str(left), 'white')} " f"right={colour(str(right), 'white')} " f"act_id={colour(str(act_id), 'white')} " f"frame_id={colour(str(frame_id), 'white')}" ) case Executed(op=op, result=result, bool_out=bool_out): event_type = colour("Executed", "bright_cyan") return ( f"{time_str} {event_type} {component_str}: " f"op={colour(op.name, 'white')} " f"result={colour(str(result), 'white')} " f"bool={colour(str(bool_out), 'white')}" ) case Emitted(token=token): event_type = colour("Emitted", "bright_cyan") return f"{time_str} {event_type} {component_str}: {token}" case IRAMWritten(offset=offset, count=count): event_type = colour("IRAMWritten", "bright_cyan") return ( f"{time_str} {event_type} {component_str}: " f"offset={colour(str(offset), 'white')} " f"count={colour(str(count), 'white')}" ) case CellWritten(addr=addr, old_pres=old_pres, new_pres=new_pres): event_type = colour("CellWritten", "bright_cyan") return ( f"{time_str} {event_type} {component_str}: " f"addr={colour(str(addr), 'white')} " f"{colour(old_pres.name, 'dim')}→{colour(new_pres.name, 'white')}" ) case DeferredRead(addr=addr): event_type = colour("DeferredRead", "bright_cyan") return ( f"{time_str} {event_type} {component_str}: " f"addr={colour(str(addr), 'white')}" ) case DeferredSatisfied(addr=addr, data=data): event_type = colour("DeferredSatisfied", "bright_cyan") return ( f"{time_str} {event_type} {component_str}: " f"addr={colour(str(addr), 'white')} " f"data={colour(str(data), 'white')}" ) case ResultSent(token=token): event_type = colour("ResultSent", "bright_cyan") return f"{time_str} {event_type} {component_str}: {token}" case _: return f"{time_str} {colour('Unknown', 'bright_red')} {component_str}" def format_snapshot_summary(snapshot: StateSnapshot) -> str: """Format a compact overview of the simulation state. Shows: PE count, SM count, sim time, next event time, queue depths. Args: snapshot: StateSnapshot instance Returns: Compact formatted summary """ pe_count = len(snapshot.pes) sm_count = len(snapshot.sms) sim_time_str = colour(f"{snapshot.sim_time:.1f}", "bright_green") next_time_str = colour(f"{snapshot.next_time:.1f}", "bright_green") # Count total tokens in queues total_input_tokens = sum(len(pe.input_queue) for pe in snapshot.pes.values()) total_input_tokens += sum(len(sm.input_queue) for sm in snapshot.sms.values()) lines = [ f"Simulation: {colour(f'{pe_count} PEs', 'bright_yellow')}, " f"{colour(f'{sm_count} SMs', 'bright_yellow')}", f"Time: {sim_time_str} (next event: {next_time_str})", f"Input queue depth: {colour(str(total_input_tokens), 'white')}", ] return "\n".join(lines) def format_pe_state(pe_snapshot: PESnapshot) -> str: """Format detailed PE state information (multi-line). Shows: IRAM entries, frame state, tag store, free frames, queue depth, output log length. Args: pe_snapshot: PESnapshot instance Returns: Detailed multi-line PE state """ pe_id_str = colour(f"PE {pe_snapshot.pe_id}", "bright_yellow") lines = [pe_id_str] # IRAM if pe_snapshot.iram: lines.append(" IRAM:") for addr in sorted(pe_snapshot.iram.keys()): inst = pe_snapshot.iram[addr] lines.append( f" [{colour(str(addr), 'white')}] {inst}" ) else: lines.append(" IRAM: (empty)") # Tag store if pe_snapshot.tag_store: tag_str = ", ".join( f"{colour(str(k), 'white')}: frame {colour(str(fid), 'white')} lane {colour(str(lane), 'white')}" for k, (fid, lane) in sorted(pe_snapshot.tag_store.items()) ) lines.append(f" Tag store: {{{tag_str}}}") else: lines.append(" Tag store: (empty)") # Free frames if pe_snapshot.free_frames: free_str = ", ".join(colour(str(f), "white") for f in pe_snapshot.free_frames) lines.append(f" Free frames: [{free_str}]") else: lines.append(" Free frames: (none)") # Frames if pe_snapshot.frames: lines.append(" Frames:") for frame_id, frame in enumerate(pe_snapshot.frames): frame_display = [] for slot_id, slot_val in enumerate(frame): if slot_val is None: frame_display.append("None") elif isinstance(slot_val, int): frame_display.append(colour(str(slot_val), "white")) else: # FrameDest object frame_display.append(f"FrameDest(...)") lines.append( f" Frame {colour(str(frame_id), 'white')}: [{', '.join(frame_display)}]" ) else: lines.append(" Frames: (empty)") # Input queue lines.append( f" Input queue: {colour(str(len(pe_snapshot.input_queue)), 'white')} tokens" ) # Output log lines.append( f" Output log: {colour(str(len(pe_snapshot.output_log)), 'white')} entries" ) return "\n".join(lines) def format_sm_state(sm_snapshot: SMSnapshot) -> str: """Format detailed SM state information (multi-line). Shows: non-empty cells with presence, deferred read info, T0 store size. Args: sm_snapshot: SMSnapshot instance Returns: Detailed multi-line SM state """ sm_id_str = colour(f"SM {sm_snapshot.sm_id}", "bright_yellow") lines = [sm_id_str] # Non-empty cells if sm_snapshot.cells: lines.append(" Cells:") for addr in sorted(sm_snapshot.cells.keys()): cell = sm_snapshot.cells[addr] pres_str = colour(cell.pres.name, "white") data_str = f"L={cell.data_l}" if cell.data_l is not None else "" if cell.data_r is not None: data_str += f" R={cell.data_r}" lines.append(f" [{colour(str(addr), 'white')}] {pres_str} {data_str}") else: lines.append(" Cells: (empty)") # Deferred read if sm_snapshot.deferred_read: dr = sm_snapshot.deferred_read lines.append( f" Deferred read: addr={colour(str(dr['cell_addr']), 'white')} " f"route={colour(str(dr['return_route']), 'white')}" ) # T0 store lines.append( f" T0 store: {colour(str(len(sm_snapshot.t0_store)), 'white')} entries" ) # Input queue lines.append( f" Input queue: {colour(str(len(sm_snapshot.input_queue)), 'white')} tokens" ) return "\n".join(lines) def format_step_result(result: StepResult) -> str: """Format a step result combining events and snapshot summary. Shows: events (if any), snapshot summary, sim time, and finished status. Args: result: StepResult instance Returns: Formatted output with events and summary """ lines = [] # Events if result.events: event_count = colour(str(len(result.events)), "white") lines.append(f"Events ({event_count}):") for event in result.events: lines.append(f" {format_event(event)}") else: lines.append("Events: (none)") # Summary if result.snapshot: lines.append("") lines.append(format_snapshot_summary(result.snapshot)) # Finished status if result.finished: lines.append("") lines.append(colour("Simulation finished.", "bright_green")) return "\n".join(lines)