1use std::{
2 collections::{BTreeMap, HashMap, HashSet},
3 fmt::Display,
4 ops::Range,
5};
6
7use anyhow::bail;
8use dialoguer::{theme::ColorfulTheme, BasicHistory, Confirm, FuzzySelect, History, Input, Select};
9use essential_constraint_asm::Op;
10use essential_constraint_vm::{
11 error::OpError, mut_keys_set, transient_data, Access, BytecodeMapped, OpAccess,
12 ProgramControlFlow, Repeat, SolutionAccess, Stack, StateSlotSlice, StateSlots, TransientData,
13};
14use essential_types::{
15 predicate::Predicate,
16 solution::{Solution, SolutionDataIndex},
17 ContentAddress, Key, Value, Word,
18};
19
20pub use source::Source;
21
22mod parse_types;
23mod source;
24mod state;
25
26const PROMPT: &str = "<essential-dbg>";
27const PRIMITIVES: &[&str] = &["int", "bool", "b256"];
28const COMPOUND: &[&str] = &["array", "tuple"];
29const SHOW: &[&str] = &["transient", "pre state", "post state", "decision vars"];
30
31pub struct ConstraintDebugger {
32 stack: Stack,
33 memory: essential_constraint_vm::Memory,
34 repeat: Repeat,
35 pc: usize,
36 code: BytecodeMapped<Op>,
37 solution: Solution,
38 pre_state: Vec<Vec<Word>>,
39 post_state: Vec<Vec<Word>>,
40 index: SolutionDataIndex,
41}
42
43pub struct Session<'a> {
44 solution: &'a Solution,
45 index: SolutionDataIndex,
46 mutable_keys: HashSet<&'a [Word]>,
47 transient_data: TransientData,
48 pre: &'a StateSlotSlice,
49 post: &'a StateSlotSlice,
50 code: &'a mut BytecodeMapped<Op>,
51 stack: &'a mut Stack,
52 memory: &'a mut essential_constraint_vm::Memory,
53 repeat: &'a mut Repeat,
54 pc: &'a mut usize,
55 last_op: Option<essential_constraint_asm::Constraint>,
56 pos: usize,
57}
58
59pub enum Outcome {
60 ProgramEnd,
61 Step,
62 Panic(OpError),
63}
64
65pub async fn run_with_source(
66 solution: Solution,
67 index: SolutionDataIndex,
68 predicate: Predicate,
69 constraint: usize,
70 state: HashMap<ContentAddress, BTreeMap<Key, Value>>,
71 source: Source,
72) -> anyhow::Result<()> {
73 run_inner(solution, index, predicate, constraint, state, Some(source)).await
74}
75
76pub async fn run(
77 solution: Solution,
78 index: SolutionDataIndex,
79 predicate: Predicate,
80 constraint: usize,
81 state: HashMap<ContentAddress, BTreeMap<Key, Value>>,
82) -> anyhow::Result<()> {
83 run_inner(solution, index, predicate, constraint, state, None).await
84}
85
86async fn run_inner(
87 solution: Solution,
88 index: SolutionDataIndex,
89 predicate: Predicate,
90 constraint: usize,
91 state: HashMap<ContentAddress, BTreeMap<Key, Value>>,
92 source: Option<Source>,
93) -> anyhow::Result<()> {
94 let mut debugger =
95 ConstraintDebugger::new(solution, index, predicate, constraint, state).await?;
96 let mut session = debugger.start_session();
97
98 let mut out = String::new();
99
100 let mut history = BasicHistory::new().max_entries(20).no_duplicates(true);
101
102 loop {
103 let command: String = Input::with_theme(&ColorfulTheme::default())
104 .with_prompt(&format!("{}\n{}", out, PROMPT))
105 .history_with(&mut history)
106 .interact_text()?;
107
108 match command.as_str() {
109 "n" | "next" => session.next(&mut out)?,
110 "b" | "back" => session.back(&mut out)?,
111 "e" | "end" => session.play_till_error(&mut out)?,
112 "q" | "quit" | "exit" => break,
113 "h" | "help" => {
114 out = help_msg();
115 }
116 "h t" | "help type" | "h type" | "help t" => {
117 out = types_msg();
118 }
119 "h c" | "help code" | "h code" | "help c" => {
120 out = help_code();
121 }
122 "s" | "show" => {
123 let prompt = format!("{}::show", PROMPT);
124 let selection = FuzzySelect::with_theme(&ColorfulTheme::default())
125 .with_prompt(&format!("What would you like to show?\n{}", prompt))
126 .default(0)
127 .items(SHOW)
128 .interact()?;
129 match SHOW[selection] {
130 "transient" => {
131 let prompt = format!("{}::transient", prompt);
132 let indices = (0..session.solution.data.len()).collect::<Vec<_>>();
133 let selection = Select::with_theme(&ColorfulTheme::default())
134 .with_prompt(&format!("Which solution data?\n{}", prompt))
135 .default(0)
136 .items(&indices)
137 .interact()?;
138
139 let prompt = format!("{}::{}", prompt, selection);
140 let t = session
141 .transient_data
142 .get(&(selection as u16))
143 .expect("Can't be out of bounds");
144 let keys: Vec<String> = t
145 .keys()
146 .map(|k| {
147 k.iter()
148 .map(|i| i.to_string())
149 .collect::<Vec<String>>()
150 .join(" ")
151 })
152 .collect();
153 let selection = FuzzySelect::with_theme(&ColorfulTheme::default())
154 .with_prompt(&format!("Which key would you like to show?\n{}", prompt))
155 .default(0)
156 .items(&keys)
157 .interact()?;
158 let key = keys[selection]
159 .split(' ')
160 .map(|i| i.parse().unwrap())
161 .collect::<Vec<_>>();
162 let v = t.get(&key).unwrap();
163 out = format!("Transient data: {:?} => {:?}", key, v);
164 }
165 "pre state" => {
166 let prompt = format!("{}::pre", prompt);
167 let indices = (0..session.pre.len()).collect::<Vec<_>>();
168 let selection = Select::with_theme(&ColorfulTheme::default())
169 .with_prompt(&format!("Which slot would you like to show?\n{}", prompt))
170 .default(0)
171 .items(&indices)
172 .interact()?;
173 let v = &session.pre[selection];
174 out = format!("Pre state slot {}: {:?}", selection, v);
175 }
176 "post state" => {
177 let prompt = format!("{}::post", prompt);
178 let indices = (0..session.post.len()).collect::<Vec<_>>();
179 let selection = Select::with_theme(&ColorfulTheme::default())
180 .with_prompt(&format!("Which slot would you like to show?\n{}", prompt))
181 .default(0)
182 .items(&indices)
183 .interact()?;
184 let v = &session.post[selection];
185 out = format!("Post state slot {}: {:?}", selection, v);
186 }
187 "decision vars" => {
188 let prompt = format!("{}::decision_vars", prompt);
189 let indices = (0..session.solution.data[index as usize]
190 .decision_variables
191 .len())
192 .collect::<Vec<_>>();
193 let selection = Select::with_theme(&ColorfulTheme::default())
194 .with_prompt(&format!(
195 "Which solution data slot would you like to show?\n{}",
196 prompt
197 ))
198 .default(0)
199 .items(&indices)
200 .interact()?;
201 let v =
202 &session.solution.data[index as usize].decision_variables[selection];
203 out = format!("Decision variable {}: {:?}", selection, v);
204 }
205 _ => unreachable!(),
206 }
207 }
208 _ => {
209 let mut c = command.split(' ');
210
211 let Some(next_command) = c.next() else {
212 out = format!("Unknown command: {}", command);
213 continue;
214 };
215 match next_command {
216 "p" | "play" => {
217 let i = c
218 .next()
219 .and_then(|i| i.parse::<usize>().ok())
220 .unwrap_or_default();
221 session.play(i, &mut out)?;
222 }
223 "l" | "list" => match c.next() {
224 Some(i) => {
225 let start = i.parse::<isize>().unwrap_or(0);
226 let end = c.next().and_then(|i| i.parse::<isize>().ok()).unwrap_or(10);
227 session.list_range(start..end, &mut out);
228 }
229 None => session.list(&mut out),
230 },
231 "t" | "type" => {
232 let rest = c.filter(|s| !s.is_empty()).collect::<Vec<_>>().join(" ");
233 if rest.is_empty() {
234 let prompt = format!("{}::type", PROMPT);
235 let pos: String = Input::with_theme(&ColorfulTheme::default())
236 .with_prompt(&format!("Enter position\n{}", prompt))
237 .default("0".to_string())
238 .history_with(&mut history)
239 .interact_text()?;
240 let pos: usize = pos.trim().parse().unwrap_or_default();
241
242 let prompt = format!("{}::{}", prompt, pos);
243 let mut options = PRIMITIVES.to_vec();
244 options.extend_from_slice(COMPOUND);
245
246 let selection = FuzzySelect::with_theme(&ColorfulTheme::default())
247 .with_prompt(&format!("Select type\n{}", prompt))
248 .default(0)
249 .items(&options[..])
250 .interact()?;
251 if PRIMITIVES.contains(&options[selection]) {
252 let input = format!("{} {}", pos, &options[selection]);
253 out = session.parse_type(&input);
254 } else {
255 let prompt = format!("{}::{}", prompt, options[selection]);
256 let input =
257 match options[selection] {
258 "array" => {
259 let selection =
260 FuzzySelect::with_theme(&ColorfulTheme::default())
261 .with_prompt(&format!(
262 "Select array type\n{}",
263 prompt
264 ))
265 .default(0)
266 .items(PRIMITIVES)
267 .interact()?;
268
269 let prompt =
270 format!("{}::{}", prompt, PRIMITIVES[selection]);
271 let len: String =
272 Input::with_theme(&ColorfulTheme::default())
273 .with_prompt(&format!(
274 "Enter array length\n{}",
275 prompt
276 ))
277 .default("1".to_string())
278 .history_with(&mut history)
279 .interact_text()?;
280 let len: usize = len.trim().parse().unwrap_or_default();
281 format!("{} {}[{}]", pos, PRIMITIVES[selection], len)
282 }
283 "tuple" => {
284 let mut fields = String::new();
285 let mut add_field = true;
286 while add_field {
287 let p = format!("{}::{{ {} }}", prompt, fields);
288 let selection = FuzzySelect::with_theme(
289 &ColorfulTheme::default(),
290 )
291 .with_prompt(&format!("Select field type\n{}", p))
292 .default(0)
293 .items(PRIMITIVES)
294 .interact()?;
295
296 fields.push_str(PRIMITIVES[selection]);
297
298 let p = format!("{}::{{ {} }}", prompt, fields);
299
300 add_field =
301 Confirm::with_theme(&ColorfulTheme::default())
302 .with_prompt(&format!(
303 "Do you want to add another field?\n{}",
304 p
305 ))
306 .default(true)
307 .interact()?;
308 if add_field {
309 fields.push_str(", ");
310 }
311 }
312 format!("{} {{ {} }}", pos, fields)
313 }
314 _ => unreachable!(),
315 };
316
317 let force_hex = Confirm::with_theme(&ColorfulTheme::default())
318 .with_prompt(&format!(
319 "Do you want to force HEX formatting?\n{}",
320 prompt
321 ))
322 .default(false)
323 .interact()?;
324 let input = if force_hex {
325 format!("{} HEX", input)
326 } else {
327 input
328 };
329 history.write(&format!("t {}", input));
330 out = session.parse_type(&input);
331 }
332 } else {
333 out = session.parse_type(&rest);
334 }
335 }
336 "c" | "code" => {
337 out = source::show_code(&source, c.next().into());
338 }
339 _ => {
340 out = format!("Unknown command: {}", command);
341 }
342 }
343 }
344 }
345 }
346
347 Ok(())
348}
349
350fn end(out: &mut String) {
351 out.push_str("\nProgram ended");
352}
353
354fn help_msg() -> String {
355 r#"Commands:
356 n | next: Step forward
357 b | back: Step back
358 p | play [i]: Play to ith op
359 e | end: Play till end or error is hit
360 l | list [start] [end]: List ops from start to end
361 s | show: Show transient data, pre state, or post state
362 c | code: Show source code. See `help code` for more info.
363 t | type <i> [type]: Parse the ith word in the stack as the given type. See `help type` for more info.
364 q | quit | exit: Quit
365 h | help: Show this message
366 "#
367 .to_string()
368}
369
370fn types_msg() -> String {
371 r#"Primitives: int, bool, b256
372 Arrays: primitive[] (e.g. int[])
373 Tuple: { primitive, primitive, ... } (e.g. {int, bool, b256})
374 Note that nesting types is not currently supported.
375 To parse a section of the stack as a type, use `t <i> [type]`
376 e.g. `t 1 int[2]` to parse the second and third word as ints.
377 `b256` is always printed as hex.
378 You can force hex formatting by adding `HEX` to the end of the command.
379 e.g. `t 1 int HEX`
380 "#
381 .to_string()
382}
383
384fn help_code() -> String {
385 r#"Shows source code if present.
386 c | code: equivalent to `code constraint`.
387 `code` can be followed by:
388 Commands:
389 a | all: Show all source code
390 p | predicate: Show predicate source code
391 c | constraint: Show predicate source code with only the constraint
392 that is being debugged. Constraint line number is required.
393 co | constraint_only: Show only the constraint line.
394 Constraint line number is required.
395 "#
396 .to_string()
397}
398
399impl ConstraintDebugger {
400 pub async fn new(
401 solution: Solution,
402 index: SolutionDataIndex,
403 predicate: Predicate,
404 constraint: usize,
405 state: HashMap<ContentAddress, BTreeMap<Key, Value>>,
406 ) -> anyhow::Result<Self> {
407 let slots = state::read_state(&solution, index, &predicate, state.clone()).await?;
408
409 let Some(code) = predicate.constraints.get(constraint).cloned() else {
410 bail!("No constraint found");
411 };
412
413 let code = BytecodeMapped::try_from_bytes(code)?;
414 let s = Self {
415 stack: Default::default(),
416 memory: Default::default(),
417 repeat: Default::default(),
418 pc: 0,
419 code,
420 solution,
421 pre_state: slots.pre,
422 post_state: slots.post,
423 index,
424 };
425 Ok(s)
426 }
427
428 pub fn start_session(&mut self) -> Session<'_> {
429 let mutable_keys = mut_keys_set(&self.solution, self.index);
430 let transient_data = transient_data(&self.solution);
431 Session {
432 code: &mut self.code,
433 stack: &mut self.stack,
434 memory: &mut self.memory,
435 repeat: &mut self.repeat,
436 pc: &mut self.pc,
437 last_op: None,
438 solution: &self.solution,
439 index: self.index,
440 mutable_keys,
441 transient_data,
442 pre: &self.pre_state,
443 post: &self.post_state,
444 pos: 0,
445 }
446 }
447}
448
449fn handle_outcome(outcome: Outcome, out: &mut String) {
450 match outcome {
451 Outcome::ProgramEnd => end(out),
452 Outcome::Panic(e) => {
453 *out = format!("Program panic: {:?}\n{}", e, out);
454 }
455 Outcome::Step => (),
456 }
457}
458
459impl Session<'_> {
460 pub fn reset_session(&mut self) {
461 *self.stack = Default::default();
462 *self.memory = Default::default();
463 *self.repeat = Default::default();
464 *self.pc = 0;
465 self.pos = 0;
466 }
467
468 pub fn next(&mut self, out: &mut String) -> anyhow::Result<()> {
469 let outcome = self.step_forward()?;
470
471 *out = format!("{}", self);
472
473 handle_outcome(outcome, out);
474 Ok(())
475 }
476
477 pub fn back(&mut self, out: &mut String) -> anyhow::Result<()> {
478 let pos = self.pos.saturating_sub(1);
479 self.reset_session();
480 let outcome = self.play_to(pos)?;
481 *out = format!("{}", self);
482
483 handle_outcome(outcome, out);
484 Ok(())
485 }
486
487 pub fn play(&mut self, i: usize, out: &mut String) -> anyhow::Result<()> {
488 self.reset_session();
489 let outcome = self.play_to(i)?;
490 *out = format!("{}", self);
491
492 handle_outcome(outcome, out);
493 Ok(())
494 }
495
496 pub fn play_till_error(&mut self, out: &mut String) -> anyhow::Result<()> {
497 loop {
498 match self.step_forward()? {
499 Outcome::Step => (),
500 Outcome::ProgramEnd => match &self.stack[..] {
501 [1] => {
502 *out = format!("Program ended successfully.\n{}", self);
503 break;
504 }
505 [0] => {
506 *out = format!("Program ended with false!\n{}", self);
507 break;
508 }
509 _ => {
510 *out = format!(
511 "Program ended with unexpected stack: {:?}\n{}",
512 self.stack, self
513 );
514 break;
515 }
516 },
517 Outcome::Panic(e) => {
518 *out = format!("Program panic: {:?}\n{}", e, self);
519 break;
520 }
521 }
522 }
523 Ok(())
524 }
525
526 pub fn step_forward(&mut self) -> anyhow::Result<Outcome> {
527 let Self {
528 code,
529 stack,
530 memory,
531 repeat,
532 pc,
533 last_op,
534 solution,
535 index,
536 mutable_keys,
537 transient_data,
538 pre,
539 post,
540 pos,
541 } = self;
542
543 let access = Access {
544 solution: SolutionAccess::new(solution, *index, mutable_keys, transient_data),
545 state_slots: StateSlots { pre, post },
546 };
547
548 let op = (&**code).op_access(**pc);
549
550 let op = match op {
551 Some(Ok(op)) => op,
552 Some(Err(err)) => {
553 bail!("Error: {:?}", err);
555 }
556 None => {
557 return Ok(Outcome::ProgramEnd);
559 }
560 };
561
562 last_op.replace(op);
563
564 let result = match essential_constraint_vm::step_op(access, op, stack, memory, **pc, repeat)
565 {
566 Ok(r) => r,
567 Err(e) => {
568 *pos += 1;
569 return Ok(Outcome::Panic(e));
570 }
571 };
572 *pos += 1;
573
574 match result {
575 Some(ProgramControlFlow::Pc(new_pc)) => {
576 **pc = new_pc;
577 Ok(Outcome::Step)
578 }
579 Some(ProgramControlFlow::Halt) => Ok(Outcome::ProgramEnd),
580 None => {
581 **pc += 1;
582 Ok(Outcome::Step)
583 }
584 }
585 }
586
587 pub fn play_to(&mut self, i: usize) -> anyhow::Result<Outcome> {
588 let mut out = None;
589 let i = if i == 0 { 1 } else { i };
590 for _ in 0..i {
591 match self.step_forward()? {
592 Outcome::ProgramEnd => return Ok(Outcome::ProgramEnd),
593 Outcome::Panic(e) => return Ok(Outcome::Panic(e)),
594 Outcome::Step => {
595 out = Some(Outcome::Step);
596 }
597 }
598 }
599 let Some(out) = out else {
600 bail!("Program didn't run");
601 };
602 Ok(out)
603 }
604
605 pub fn list_range(&self, range: Range<isize>, out: &mut String) {
606 use std::fmt::Write;
607 let start = (self.pos as isize).saturating_add(range.start).max(0) as usize;
608 let end = (self.pos as isize).saturating_add(range.end).max(0) as usize;
609 let len = end.saturating_sub(start);
610 let this_op = (start..end)
611 .contains(&self.pos)
612 .then_some(self.pos.saturating_sub(start));
613 if let Some(ops) = &self.code.ops_from(start) {
614 *out = ops
615 .ops()
616 .take(len)
617 .enumerate()
618 .fold(String::new(), |mut out, (i, op)| {
619 match &this_op {
620 Some(this_op) if *this_op == i => {
621 let _ = writeln!(
622 out,
623 "{}:Op: {:?}",
624 start + i,
625 dialoguer::console::style(op).cyan()
626 );
627 }
628 _ => {
629 let _ = writeln!(out, "{}:Op: {:?}", start + i, op);
630 }
631 }
632 out
633 });
634 }
635 }
636
637 pub fn list(&self, out: &mut String) {
638 use std::fmt::Write;
639 if let Some(ops) = &self.code.ops_from(0) {
640 *out = ops
641 .ops()
642 .enumerate()
643 .fold(String::new(), |mut out, (i, op)| {
644 if self.pos == i {
645 let _ =
646 writeln!(out, "{}:Op: {:?}", i, dialoguer::console::style(op).cyan());
647 } else {
648 let _ = writeln!(out, "{}:Op: {:?}", i, op);
649 }
650 out
651 });
652 }
653 }
654
655 pub fn parse_type(&self, ty: &str) -> String {
656 parse_types::parse_type(&self.stack[..], ty)
657 }
658}
659
660impl Display for Session<'_> {
661 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
662 if let Some(op) = &self.last_op {
663 writeln!(f, "Op: {:?}", op)?;
664 }
665 writeln!(f, " ├── {:?}\n └── {:?}", self.stack, self.memory)
666 }
667}