OR-1 dataflow CPU sketch
1"""ANSI formatting helpers for monitor REPL output.
2
3Provides functions to format SimEvent, snapshot state, and step results
4with ANSI colour codes for readable terminal display.
5
6Key features:
7- NO_COLOUR flag to disable all colouring (for testing and pipes)
8- Distinct colours for event types, component IDs, data values, errors, and time
9- Compact and detailed formatting modes for different use cases
10"""
11
12from __future__ import annotations
13
14from emu.events import (
15 CellWritten,
16 DeferredRead,
17 DeferredSatisfied,
18 Emitted,
19 Executed,
20 IRAMWritten,
21 Matched,
22 ResultSent,
23 SimEvent,
24 TokenReceived,
25)
26from monitor.commands import StepResult
27from monitor.snapshot import PESnapshot, SMSnapshot, StateSnapshot
28
29# Global flag to disable all colouring (for testing and pipes)
30NO_COLOUR = False
31
32# ANSI colour codes
33COLOURS = {
34 "reset": "\033[0m",
35 "bright_cyan": "\033[96m", # Event types
36 "bright_yellow": "\033[93m", # PE/SM identifiers
37 "white": "\033[97m", # Data values
38 "bright_red": "\033[91m", # Errors
39 "bright_green": "\033[92m", # Sim time
40 "dim": "\033[2m", # Dim (for secondary info)
41}
42
43
44def colour(text: str, code: str) -> str:
45 """Wrap text in ANSI escape code, respecting NO_COLOUR flag.
46
47 Args:
48 text: Text to colour
49 code: Colour code key (from COLOURS dict)
50
51 Returns:
52 Coloured text if NO_COLOUR is False, otherwise plain text
53 """
54 if NO_COLOUR:
55 return text
56 return f"{COLOURS[code]}{text}{COLOURS['reset']}"
57
58
59def format_event(event: SimEvent) -> str:
60 """Format a single SimEvent as a coloured one-line string.
61
62 Shows event type and key fields in distinct colours.
63
64 Args:
65 event: SimEvent instance
66
67 Returns:
68 Formatted event string with ANSI colours
69 """
70 time_str = colour(f"[{event.time:.1f}]", "bright_green")
71 component_str = colour(event.component, "bright_yellow")
72
73 match event:
74 case TokenReceived(token=token):
75 event_type = colour("TokenReceived", "bright_cyan")
76 return f"{time_str} {event_type} {component_str}: {token}"
77
78 case Matched(left=left, right=right, ctx=ctx, offset=offset):
79 event_type = colour("Matched", "bright_cyan")
80 return (
81 f"{time_str} {event_type} {component_str}: "
82 f"left={colour(str(left), 'white')} "
83 f"right={colour(str(right), 'white')} "
84 f"ctx={colour(str(ctx), 'white')} "
85 f"offset={colour(str(offset), 'white')}"
86 )
87
88 case Executed(op=op, result=result, bool_out=bool_out):
89 event_type = colour("Executed", "bright_cyan")
90 return (
91 f"{time_str} {event_type} {component_str}: "
92 f"op={colour(op.name, 'white')} "
93 f"result={colour(str(result), 'white')} "
94 f"bool={colour(str(bool_out), 'white')}"
95 )
96
97 case Emitted(token=token):
98 event_type = colour("Emitted", "bright_cyan")
99 return f"{time_str} {event_type} {component_str}: {token}"
100
101 case IRAMWritten(offset=offset, count=count):
102 event_type = colour("IRAMWritten", "bright_cyan")
103 return (
104 f"{time_str} {event_type} {component_str}: "
105 f"offset={colour(str(offset), 'white')} "
106 f"count={colour(str(count), 'white')}"
107 )
108
109 case CellWritten(addr=addr, old_pres=old_pres, new_pres=new_pres):
110 event_type = colour("CellWritten", "bright_cyan")
111 return (
112 f"{time_str} {event_type} {component_str}: "
113 f"addr={colour(str(addr), 'white')} "
114 f"{colour(old_pres.name, 'dim')}→{colour(new_pres.name, 'white')}"
115 )
116
117 case DeferredRead(addr=addr):
118 event_type = colour("DeferredRead", "bright_cyan")
119 return (
120 f"{time_str} {event_type} {component_str}: "
121 f"addr={colour(str(addr), 'white')}"
122 )
123
124 case DeferredSatisfied(addr=addr, data=data):
125 event_type = colour("DeferredSatisfied", "bright_cyan")
126 return (
127 f"{time_str} {event_type} {component_str}: "
128 f"addr={colour(str(addr), 'white')} "
129 f"data={colour(str(data), 'white')}"
130 )
131
132 case ResultSent(token=token):
133 event_type = colour("ResultSent", "bright_cyan")
134 return f"{time_str} {event_type} {component_str}: {token}"
135
136 case _:
137 return f"{time_str} {colour('Unknown', 'bright_red')} {component_str}"
138
139
140def format_snapshot_summary(snapshot: StateSnapshot) -> str:
141 """Format a compact overview of the simulation state.
142
143 Shows: PE count, SM count, sim time, next event time, queue depths.
144
145 Args:
146 snapshot: StateSnapshot instance
147
148 Returns:
149 Compact formatted summary
150 """
151 pe_count = len(snapshot.pes)
152 sm_count = len(snapshot.sms)
153 sim_time_str = colour(f"{snapshot.sim_time:.1f}", "bright_green")
154 next_time_str = colour(f"{snapshot.next_time:.1f}", "bright_green")
155
156 # Count total tokens in queues
157 total_input_tokens = sum(len(pe.input_queue) for pe in snapshot.pes.values())
158 total_input_tokens += sum(len(sm.input_queue) for sm in snapshot.sms.values())
159
160 lines = [
161 f"Simulation: {colour(f'{pe_count} PEs', 'bright_yellow')}, "
162 f"{colour(f'{sm_count} SMs', 'bright_yellow')}",
163 f"Time: {sim_time_str} (next event: {next_time_str})",
164 f"Input queue depth: {colour(str(total_input_tokens), 'white')}",
165 ]
166 return "\n".join(lines)
167
168
169def format_pe_state(pe_snapshot: PESnapshot) -> str:
170 """Format detailed PE state information (multi-line).
171
172 Shows: IRAM entries, matching store occupancy, gen counters, queue depth, output log length.
173
174 Args:
175 pe_snapshot: PESnapshot instance
176
177 Returns:
178 Detailed multi-line PE state
179 """
180 pe_id_str = colour(f"PE {pe_snapshot.pe_id}", "bright_yellow")
181 lines = [pe_id_str]
182
183 # IRAM
184 if pe_snapshot.iram:
185 lines.append(" IRAM:")
186 for addr in sorted(pe_snapshot.iram.keys()):
187 inst = pe_snapshot.iram[addr]
188 lines.append(
189 f" [{colour(str(addr), 'white')}] {inst}"
190 )
191 else:
192 lines.append(" IRAM: (empty)")
193
194 # Matching store occupancy
195 total_occupied = 0
196 for ctx_row in pe_snapshot.matching_store:
197 for entry in ctx_row:
198 if entry.get("occupied", False):
199 total_occupied += 1
200 lines.append(
201 f" Matching store: {colour(str(total_occupied), 'white')} entries occupied"
202 )
203
204 # Gen counters
205 if pe_snapshot.gen_counters:
206 gens_str = ", ".join(colour(str(g), "white") for g in pe_snapshot.gen_counters)
207 lines.append(f" Gen counters: [{gens_str}]")
208
209 # Input queue
210 lines.append(
211 f" Input queue: {colour(str(len(pe_snapshot.input_queue)), 'white')} tokens"
212 )
213
214 # Output log
215 lines.append(
216 f" Output log: {colour(str(len(pe_snapshot.output_log)), 'white')} entries"
217 )
218
219 return "\n".join(lines)
220
221
222def format_sm_state(sm_snapshot: SMSnapshot) -> str:
223 """Format detailed SM state information (multi-line).
224
225 Shows: non-empty cells with presence, deferred read info, T0 store size.
226
227 Args:
228 sm_snapshot: SMSnapshot instance
229
230 Returns:
231 Detailed multi-line SM state
232 """
233 sm_id_str = colour(f"SM {sm_snapshot.sm_id}", "bright_yellow")
234 lines = [sm_id_str]
235
236 # Non-empty cells
237 if sm_snapshot.cells:
238 lines.append(" Cells:")
239 for addr in sorted(sm_snapshot.cells.keys()):
240 cell = sm_snapshot.cells[addr]
241 pres_str = colour(cell.pres.name, "white")
242 data_str = f"L={cell.data_l}" if cell.data_l is not None else ""
243 if cell.data_r is not None:
244 data_str += f" R={cell.data_r}"
245 lines.append(f" [{colour(str(addr), 'white')}] {pres_str} {data_str}")
246 else:
247 lines.append(" Cells: (empty)")
248
249 # Deferred read
250 if sm_snapshot.deferred_read:
251 dr = sm_snapshot.deferred_read
252 lines.append(
253 f" Deferred read: addr={colour(str(dr['cell_addr']), 'white')} "
254 f"route={colour(str(dr['return_route']), 'white')}"
255 )
256
257 # T0 store
258 lines.append(
259 f" T0 store: {colour(str(len(sm_snapshot.t0_store)), 'white')} tokens"
260 )
261
262 # Input queue
263 lines.append(
264 f" Input queue: {colour(str(len(sm_snapshot.input_queue)), 'white')} tokens"
265 )
266
267 return "\n".join(lines)
268
269
270def format_step_result(result: StepResult) -> str:
271 """Format a step result combining events and snapshot summary.
272
273 Shows: events (if any), snapshot summary, sim time, and finished status.
274
275 Args:
276 result: StepResult instance
277
278 Returns:
279 Formatted output with events and summary
280 """
281 lines = []
282
283 # Events
284 if result.events:
285 event_count = colour(str(len(result.events)), "white")
286 lines.append(f"Events ({event_count}):")
287 for event in result.events:
288 lines.append(f" {format_event(event)}")
289 else:
290 lines.append("Events: (none)")
291
292 # Summary
293 if result.snapshot:
294 lines.append("")
295 lines.append(format_snapshot_summary(result.snapshot))
296
297 # Finished status
298 if result.finished:
299 lines.append("")
300 lines.append(colour("Simulation finished.", "bright_green"))
301
302 return "\n".join(lines)