foundry_debugger/tui/

context.rs

//! Debugger context and event handler implementation.

use crate::{DebugNode, ExitReason, debugger::DebuggerContext};
use alloy_primitives::{Address, hex};
use crossterm::event::{Event, KeyCode, KeyEvent, KeyModifiers, MouseEvent, MouseEventKind};
use foundry_evm_core::buffer::BufferKind;
use foundry_tui::TuiApp;
use ratatui::Frame;
use revm::bytecode::opcode::OpCode;
use revm_inspectors::tracing::types::{CallKind, CallTraceStep};
use std::ops::ControlFlow;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum StatusKind {
    Info,
    Error,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub(crate) struct StatusMessage {
    pub(crate) kind: StatusKind,
    pub(crate) text: String,
}

/// This is currently used to remember last scroll position so screen doesn't wiggle as much.
#[derive(Default)]
pub(crate) struct DrawMemory {
    pub(crate) inner_call_index: usize,
    pub(crate) current_buf_startline: usize,
    pub(crate) current_stack_startline: usize,
}

pub(crate) struct TUIContext<'a> {
    pub(crate) debugger_context: &'a mut DebuggerContext,

    /// Buffer for keys prior to execution, i.e. '10' + 'k' => move up 10 operations.
    pub(crate) key_buffer: String,
    /// Current goto program counter prompt contents, if the prompt is active.
    pub(crate) pc_input: Option<String>,
    /// Last status or error message to show in the footer.
    pub(crate) status: Option<StatusMessage>,
    /// Current step in the debug steps.
    pub(crate) current_step: usize,
    pub(crate) draw_memory: DrawMemory,
    pub(crate) opcode_list: Vec<String>,
    pub(crate) last_index: usize,

    pub(crate) stack_labels: bool,
    /// Whether to decode active buffer as utf8 or not.
    pub(crate) buf_utf: bool,
    pub(crate) show_shortcuts: bool,
    /// The currently active buffer (memory, calldata, returndata) to be drawn.
    pub(crate) active_buffer: BufferKind,
}

impl<'a> TUIContext<'a> {
    pub(crate) fn new(debugger_context: &'a mut DebuggerContext) -> Self {
        TUIContext {
            debugger_context,

            key_buffer: String::with_capacity(64),
            pc_input: None,
            status: None,
            current_step: 0,
            draw_memory: DrawMemory::default(),
            opcode_list: Vec::new(),
            last_index: 0,

            stack_labels: false,
            buf_utf: false,
            show_shortcuts: true,
            active_buffer: BufferKind::Memory,
        }
    }

    pub(crate) fn init(&mut self) {
        self.gen_opcode_list();
    }

    pub(crate) fn debug_arena(&self) -> &[DebugNode] {
        &self.debugger_context.debug_arena
    }

    pub(crate) fn debug_call(&self) -> &DebugNode {
        &self.debug_arena()[self.draw_memory.inner_call_index]
    }

    /// Returns the current call address.
    pub(crate) fn address(&self) -> &Address {
        &self.debug_call().address
    }

    /// Returns the current call kind.
    pub(crate) fn call_kind(&self) -> CallKind {
        self.debug_call().kind
    }

    /// Returns the current debug steps.
    pub(crate) fn debug_steps(&self) -> &[CallTraceStep] {
        &self.debug_call().steps
    }

    /// Returns the current debug step.
    pub(crate) fn current_step(&self) -> &CallTraceStep {
        &self.debug_steps()[self.current_step]
    }

    fn gen_opcode_list(&mut self) {
        self.opcode_list.clear();
        let debug_steps =
            &self.debugger_context.debug_arena[self.draw_memory.inner_call_index].steps;
        for step in debug_steps {
            self.opcode_list.push(pretty_opcode(step));
        }
    }

    fn gen_opcode_list_if_necessary(&mut self) {
        if self.last_index != self.draw_memory.inner_call_index {
            self.gen_opcode_list();
            self.last_index = self.draw_memory.inner_call_index;
        }
    }

    fn active_buffer(&self) -> &[u8] {
        match self.active_buffer {
            BufferKind::Memory => self.current_step().memory.as_ref().unwrap().as_bytes(),
            BufferKind::Calldata => &self.debug_call().calldata,
            BufferKind::Returndata => &self.current_step().returndata,
        }
    }
}

impl TUIContext<'_> {
    pub(crate) fn handle_event(&mut self, event: Event) -> ControlFlow<ExitReason> {
        let ret = match event {
            Event::Key(event) => self.handle_key_event(event),
            Event::Mouse(event) => self.handle_mouse_event(event),
            _ => ControlFlow::Continue(()),
        };
        // Generate the list after the event has been handled.
        self.gen_opcode_list_if_necessary();
        ret
    }

    fn handle_key_event(&mut self, event: KeyEvent) -> ControlFlow<ExitReason> {
        if self.pc_input.is_some() {
            self.handle_pc_input_key_event(event);
            return ControlFlow::Continue(());
        }

        // Breakpoints
        if let KeyCode::Char(c) = event.code
            && c.is_alphabetic()
            && self.key_buffer.starts_with('\'')
        {
            self.handle_breakpoint(c);
            return ControlFlow::Continue(());
        }

        let control = event.modifiers.contains(KeyModifiers::CONTROL);

        match event.code {
            // Exit
            KeyCode::Char('q') => return ControlFlow::Break(ExitReason::CharExit),

            // Scroll up the memory buffer
            KeyCode::Char('k') | KeyCode::Up if control => self.repeat(|this| {
                this.draw_memory.current_buf_startline =
                    this.draw_memory.current_buf_startline.saturating_sub(1);
            }),
            // Scroll down the memory buffer
            KeyCode::Char('j') | KeyCode::Down if control => self.repeat(|this| {
                let max_buf = (this.active_buffer().len() / 32).saturating_sub(1);
                if this.draw_memory.current_buf_startline < max_buf {
                    this.draw_memory.current_buf_startline += 1;
                }
            }),

            // Move up
            KeyCode::Char('k') | KeyCode::Up => self.repeat(Self::step_back),
            // Move down
            KeyCode::Char('j') | KeyCode::Down => self.repeat(Self::step),

            // Scroll up the stack
            KeyCode::Char('K') => self.repeat(|this| {
                this.draw_memory.current_stack_startline =
                    this.draw_memory.current_stack_startline.saturating_sub(1);
            }),
            // Scroll down the stack
            KeyCode::Char('J') => self.repeat(|this| {
                let max_stack =
                    this.current_step().stack.as_ref().map_or(0, |s| s.len()).saturating_sub(1);
                if this.draw_memory.current_stack_startline < max_stack {
                    this.draw_memory.current_stack_startline += 1;
                }
            }),

            // Cycle buffers
            KeyCode::Char('b') => {
                self.active_buffer = self.active_buffer.next();
                self.draw_memory.current_buf_startline = 0;
            }

            // Go to top of file
            KeyCode::Char('g') => {
                self.draw_memory.inner_call_index = 0;
                self.current_step = 0;
            }

            // Go to bottom of file
            KeyCode::Char('G') => {
                self.draw_memory.inner_call_index = self.debug_arena().len() - 1;
                self.current_step = self.n_steps() - 1;
            }

            // Go to previous call
            KeyCode::Char('c') => {
                self.draw_memory.inner_call_index =
                    self.draw_memory.inner_call_index.saturating_sub(1);
                self.current_step = self.n_steps() - 1;
            }

            // Go to next call
            KeyCode::Char('C')
                if self.debug_arena().len() > self.draw_memory.inner_call_index + 1 =>
            {
                self.draw_memory.inner_call_index += 1;
                self.current_step = 0;
            }

            // Step forward
            KeyCode::Char('s') => self.repeat(|this| {
                let remaining_steps = &this.debug_steps()[this.current_step..];
                if let Some((i, _)) =
                    remaining_steps.iter().enumerate().skip(1).find(|(i, step)| {
                        let prev = &remaining_steps[*i - 1];
                        is_jump(step, prev)
                    })
                {
                    this.current_step += i
                }
            }),

            // Step backwards
            KeyCode::Char('a') => self.repeat(|this| {
                let ops = &this.debug_steps()[..this.current_step];
                this.current_step = ops
                    .iter()
                    .enumerate()
                    .skip(1)
                    .rev()
                    .find(|&(i, op)| {
                        let prev = &ops[i - 1];
                        is_jump(op, prev)
                    })
                    .map(|(i, _)| i)
                    .unwrap_or_default();
            }),

            // Toggle stack labels
            KeyCode::Char('t') => self.stack_labels = !self.stack_labels,

            // Toggle memory UTF-8 decoding
            KeyCode::Char('m') => self.buf_utf = !self.buf_utf,

            // Go to program counter
            KeyCode::Char('p') => {
                self.key_buffer.clear();
                self.status = None;
                self.pc_input = Some(String::new());
            }

            // Toggle help notice
            KeyCode::Char('h') => self.show_shortcuts = !self.show_shortcuts,

            // Numbers for repeating commands or breakpoints
            KeyCode::Char(
                other @ ('0' | '1' | '2' | '3' | '4' | '5' | '6' | '7' | '8' | '9' | '\''),
            ) => {
                // Early return to not clear the buffer.
                self.key_buffer.push(other);
                return ControlFlow::Continue(());
            }

            // Unknown/unhandled key code
            _ => {}
        };

        self.key_buffer.clear();
        ControlFlow::Continue(())
    }

    fn handle_pc_input_key_event(&mut self, event: KeyEvent) {
        match event.code {
            KeyCode::Esc => {
                self.pc_input = None;
            }
            KeyCode::Enter => {
                let input = self.pc_input.take().unwrap_or_default();
                self.goto_pc_from_input(&input);
            }
            KeyCode::Backspace => {
                if let Some(input) = &mut self.pc_input {
                    input.pop();
                }
            }
            KeyCode::Char(c)
                if !event.modifiers.contains(KeyModifiers::CONTROL) && is_pc_input_char(c) =>
            {
                if let Some(input) = &mut self.pc_input {
                    input.push(c);
                }
            }
            _ => {}
        }
    }

    fn goto_pc_from_input(&mut self, input: &str) {
        let candidates = match parse_pc_candidates(input) {
            Ok(candidates) => candidates,
            Err(err) => {
                self.set_error(err);
                return;
            }
        };

        let mut found = Vec::new();
        for &candidate in &candidates {
            if let Some(target) = find_pc_target(
                self.debug_arena(),
                self.draw_memory.inner_call_index,
                self.current_step,
                candidate.pc,
            ) {
                found.push((candidate, target));
            }
        }

        match found.as_slice() {
            [] => {
                let current = self.debug_call();
                let outside = candidates.iter().any(|candidate| {
                    pc_exists_outside_code_context(self.debug_arena(), current, candidate.pc)
                });
                let pc = if let [candidate] = candidates.as_slice() {
                    let pc = candidate.pc;
                    format!("PC 0x{pc:x} ({pc})")
                } else {
                    format!("PC `{}`", input.trim())
                };
                let mut msg = format!("{pc} not found in current contract");
                if outside {
                    msg.push_str("; it exists in another contract, switch calls first");
                }
                self.set_error(msg);
            }
            [(candidate, target)] => self.apply_pc_target(*candidate, *target),
            _ => {
                let input = input.trim();
                let options = found
                    .iter()
                    .map(|(candidate, _)| candidate.describe())
                    .collect::<Vec<_>>()
                    .join(" and ");
                self.set_error(format!(
                    "Ambiguous PC `{input}`: {options} both exist; use d:<pc> or 0x<pc>"
                ));
            }
        }
    }

    fn apply_pc_target(&mut self, candidate: PcCandidate, target: PcTarget) {
        let already_at_target = self.draw_memory.inner_call_index == target.node_index
            && self.current_step == target.step_index;

        self.draw_memory.inner_call_index = target.node_index;
        self.current_step = target.step_index;
        self.draw_memory.current_buf_startline = 0;
        self.draw_memory.current_stack_startline = 0;
        self.key_buffer.clear();

        let pc = candidate.pc;
        let scope = match target.scope {
            PcTargetScope::CurrentNode => "current trace",
            PcTargetScope::SameCodeContext => "same contract",
        };
        let action = if already_at_target { "Already at" } else { "Jumped to" };
        self.set_info(format!("{action} PC 0x{pc:x} ({pc}) in {scope}"));
    }

    fn set_info(&mut self, text: String) {
        self.status = Some(StatusMessage { kind: StatusKind::Info, text });
    }

    fn set_error(&mut self, text: String) {
        self.status = Some(StatusMessage { kind: StatusKind::Error, text });
    }

    fn handle_breakpoint(&mut self, c: char) {
        // Find the location of the called breakpoint in the whole debug arena (at this address with
        // this pc)
        if let Some((caller, pc)) = self.debugger_context.breakpoints.get(&c) {
            for (i, node) in self.debug_arena().iter().enumerate() {
                if node.address == *caller
                    && let Some(step) = node.steps.iter().position(|step| step.pc == *pc)
                {
                    self.draw_memory.inner_call_index = i;
                    self.current_step = step;
                    break;
                }
            }
        }
        self.key_buffer.clear();
    }

    fn handle_mouse_event(&mut self, event: MouseEvent) -> ControlFlow<ExitReason> {
        if self.pc_input.is_some() {
            return ControlFlow::Continue(());
        }

        match event.kind {
            MouseEventKind::ScrollUp => self.step_back(),
            MouseEventKind::ScrollDown => self.step(),
            _ => {}
        }

        ControlFlow::Continue(())
    }

    fn step_back(&mut self) {
        if self.current_step > 0 {
            self.current_step -= 1;
        } else if self.draw_memory.inner_call_index > 0 {
            self.draw_memory.inner_call_index -= 1;
            self.current_step = self.n_steps() - 1;
        }
    }

    fn step(&mut self) {
        if self.current_step < self.n_steps() - 1 {
            self.current_step += 1;
        } else if self.draw_memory.inner_call_index < self.debug_arena().len() - 1 {
            self.draw_memory.inner_call_index += 1;
            self.current_step = 0;
        }
    }

    /// Calls a closure `f` the number of times specified in the key buffer, and at least once.
    fn repeat(&mut self, mut f: impl FnMut(&mut Self)) {
        for _ in 0..buffer_as_number(&self.key_buffer) {
            f(self);
        }
    }

    fn n_steps(&self) -> usize {
        self.debug_steps().len()
    }
}

impl TuiApp for TUIContext<'_> {
    type Exit = ExitReason;

    fn draw(&mut self, frame: &mut Frame<'_>) {
        self.draw_layout(frame);
    }

    fn handle_event(&mut self, event: Event) -> ControlFlow<Self::Exit> {
        TUIContext::handle_event(self, event)
    }
}

/// Grab number from buffer. Used for something like '10k' to move up 10 operations
fn buffer_as_number(s: &str) -> usize {
    const MIN: usize = 1;
    const MAX: usize = 100_000;
    s.parse().unwrap_or(MIN).clamp(MIN, MAX)
}

const fn is_pc_input_char(c: char) -> bool {
    c.is_ascii_hexdigit() || matches!(c, 'x' | 'X' | ':')
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum PcBase {
    Hex,
    Decimal,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct PcCandidate {
    pc: usize,
    base: PcBase,
}

impl PcCandidate {
    fn describe(self) -> String {
        match self.base {
            PcBase::Hex => format!("hex 0x{:x}", self.pc),
            PcBase::Decimal => format!("decimal {}", self.pc),
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum PcTargetScope {
    CurrentNode,
    SameCodeContext,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct PcTarget {
    node_index: usize,
    step_index: usize,
    scope: PcTargetScope,
}

fn parse_pc_candidates(input: &str) -> Result<Vec<PcCandidate>, String> {
    let input = input.trim();
    if input.is_empty() {
        return Err("Enter a program counter".to_string());
    }

    if let Some(rest) = input.strip_prefix("0x").or_else(|| input.strip_prefix("0X")) {
        return parse_pc_candidate(rest, 16, PcBase::Hex, input);
    }

    if let Some(rest) = input.strip_prefix("d:").or_else(|| input.strip_prefix("dec:")) {
        return parse_pc_candidate(rest, 10, PcBase::Decimal, input);
    }

    if input.chars().any(|c| c.is_ascii_hexdigit() && c.is_ascii_alphabetic()) {
        return parse_pc_candidate(input, 16, PcBase::Hex, input);
    }

    if input.chars().all(|c| c.is_ascii_digit()) {
        let decimal = parse_pc(input, 10, input)?;
        let hex = parse_pc(input, 16, input)?;
        if decimal == hex {
            return Ok(vec![PcCandidate { pc: decimal, base: PcBase::Decimal }]);
        }
        return Ok(vec![
            PcCandidate { pc: decimal, base: PcBase::Decimal },
            PcCandidate { pc: hex, base: PcBase::Hex },
        ]);
    }

    Err(format!("Invalid PC `{input}`; use 0x2a, 2a, or d:42"))
}

fn parse_pc_candidate(
    input: &str,
    radix: u32,
    base: PcBase,
    original: &str,
) -> Result<Vec<PcCandidate>, String> {
    Ok(vec![PcCandidate { pc: parse_pc(input, radix, original)?, base }])
}

fn parse_pc(input: &str, radix: u32, original: &str) -> Result<usize, String> {
    if input.is_empty() {
        return Err(format!("Invalid PC `{original}`; use 0x2a, 2a, or d:42"));
    }
    usize::from_str_radix(input, radix)
        .map_err(|_| format!("Invalid PC `{original}`; use 0x2a, 2a, or d:42"))
}

fn find_pc_target(
    arena: &[DebugNode],
    current_node_index: usize,
    current_step: usize,
    pc: usize,
) -> Option<PcTarget> {
    let current_node = arena.get(current_node_index)?;

    if let Some(step_index) = find_pc_in_current_node(&current_node.steps, current_step, pc) {
        return Some(PcTarget {
            node_index: current_node_index,
            step_index,
            scope: PcTargetScope::CurrentNode,
        });
    }

    for (node_index, node) in arena.iter().enumerate().skip(current_node_index + 1) {
        if same_code_context(current_node, node)
            && let Some(step_index) = node.steps.iter().position(|step| step.pc == pc)
        {
            return Some(PcTarget {
                node_index,
                step_index,
                scope: PcTargetScope::SameCodeContext,
            });
        }
    }

    for (node_index, node) in arena.iter().enumerate().take(current_node_index).rev() {
        if same_code_context(current_node, node)
            && let Some(step_index) = node.steps.iter().rposition(|step| step.pc == pc)
        {
            return Some(PcTarget {
                node_index,
                step_index,
                scope: PcTargetScope::SameCodeContext,
            });
        }
    }

    None
}

fn find_pc_in_current_node(
    steps: &[CallTraceStep],
    current_step: usize,
    pc: usize,
) -> Option<usize> {
    if steps.get(current_step).is_some_and(|step| step.pc == pc) {
        return Some(current_step);
    }

    steps
        .iter()
        .enumerate()
        .skip(current_step.saturating_add(1))
        .find_map(|(i, step)| (step.pc == pc).then_some(i))
        .or_else(|| {
            steps[..current_step.min(steps.len())]
                .iter()
                .enumerate()
                .rev()
                .find_map(|(i, step)| (step.pc == pc).then_some(i))
        })
}

fn same_code_context(a: &DebugNode, b: &DebugNode) -> bool {
    a.address == b.address && a.kind.is_any_create() == b.kind.is_any_create()
}

fn pc_exists_outside_code_context(arena: &[DebugNode], current: &DebugNode, pc: usize) -> bool {
    arena.iter().any(|node| {
        !same_code_context(current, node) && node.steps.iter().any(|step| step.pc == pc)
    })
}

fn pretty_opcode(step: &CallTraceStep) -> String {
    if let Some(immediate) = step.immediate_bytes.as_ref().filter(|b| !b.is_empty()) {
        format!("{}(0x{})", step.op, hex::encode(immediate))
    } else {
        step.op.to_string()
    }
}

fn is_jump(step: &CallTraceStep, prev: &CallTraceStep) -> bool {
    if !matches!(prev.op, OpCode::JUMP | OpCode::JUMPI) {
        return false;
    }

    let immediate_len = prev.immediate_bytes.as_ref().map_or(0, |b| b.len());

    step.pc != prev.pc + 1 + immediate_len
}

#[cfg(test)]
mod tests {
    use super::*;
    use alloy_primitives::Bytes;
    use foundry_evm_core::Breakpoints;
    use foundry_evm_traces::debug::ContractSources;
    use revm::interpreter::InstructionResult;

    fn step(pc: usize) -> CallTraceStep {
        CallTraceStep {
            pc,
            op: OpCode::STOP,
            stack: None,
            push_stack: None,
            memory: None,
            returndata: Bytes::new(),
            gas_remaining: 0,
            gas_refund_counter: 0,
            gas_used: 0,
            gas_cost: 0,
            storage_change: None,
            status: Some(InstructionResult::Stop),
            immediate_bytes: None,
            decoded: None,
        }
    }

    fn node(address: Address, kind: CallKind, pcs: &[usize]) -> DebugNode {
        DebugNode::new(address, kind, pcs.iter().copied().map(step).collect(), Bytes::new(), 0)
    }

    fn context_with_arena(arena: Vec<DebugNode>) -> DebuggerContext {
        DebuggerContext {
            debug_arena: arena,
            stats: None,
            identified_contracts: Default::default(),
            contracts_sources: ContractSources::default(),
            breakpoints: Breakpoints::default(),
        }
    }

    fn key(code: KeyCode) -> KeyEvent {
        KeyEvent::new(code, KeyModifiers::empty())
    }

    #[test]
    fn parses_prefixed_hex_pc() {
        assert_eq!(
            parse_pc_candidates("0x2a").unwrap(),
            vec![PcCandidate { pc: 42, base: PcBase::Hex }]
        );
        assert_eq!(
            parse_pc_candidates("0X2A").unwrap(),
            vec![PcCandidate { pc: 42, base: PcBase::Hex }]
        );
    }

    #[test]
    fn parses_bare_hex_pc_with_letters() {
        assert_eq!(
            parse_pc_candidates("2a").unwrap(),
            vec![PcCandidate { pc: 42, base: PcBase::Hex }]
        );
    }

    #[test]
    fn parses_explicit_decimal_pc() {
        assert_eq!(
            parse_pc_candidates("d:42").unwrap(),
            vec![PcCandidate { pc: 42, base: PcBase::Decimal }]
        );
        assert_eq!(
            parse_pc_candidates("dec:42").unwrap(),
            vec![PcCandidate { pc: 42, base: PcBase::Decimal }]
        );
    }

    #[test]
    fn parses_bare_digits_as_decimal_and_hex_candidates() {
        assert_eq!(
            parse_pc_candidates("10").unwrap(),
            vec![
                PcCandidate { pc: 10, base: PcBase::Decimal },
                PcCandidate { pc: 16, base: PcBase::Hex },
            ]
        );
        assert_eq!(
            parse_pc_candidates("9").unwrap(),
            vec![PcCandidate { pc: 9, base: PcBase::Decimal }]
        );
    }

    #[test]
    fn rejects_invalid_pc_input() {
        assert!(parse_pc_candidates("").is_err());
        assert!(parse_pc_candidates("0x").is_err());
        assert!(parse_pc_candidates("xyz").is_err());
        assert!(parse_pc_candidates("184467440737095516160").is_err());
    }

    #[test]
    fn finds_pc_in_current_node() {
        let address = Address::repeat_byte(1);
        let arena = vec![node(address, CallKind::Call, &[1, 2, 3])];

        assert_eq!(
            find_pc_target(&arena, 0, 0, 3),
            Some(PcTarget { node_index: 0, step_index: 2, scope: PcTargetScope::CurrentNode })
        );
    }

    #[test]
    fn repeated_pc_stays_current_then_prefers_next_then_previous() {
        let address = Address::repeat_byte(1);
        let arena = vec![node(address, CallKind::Call, &[1, 2, 3, 2])];

        assert_eq!(find_pc_target(&arena, 0, 1, 2).unwrap().step_index, 1);
        assert_eq!(find_pc_target(&arena, 0, 0, 2).unwrap().step_index, 1);
        assert_eq!(find_pc_target(&arena, 0, 3, 2).unwrap().step_index, 3);
        assert_eq!(find_pc_target(&arena, 0, 2, 2).unwrap().step_index, 3);
    }

    #[test]
    fn searches_later_then_earlier_same_code_context() {
        let address = Address::repeat_byte(1);
        let arena = vec![
            node(address, CallKind::Call, &[1]),
            node(address, CallKind::Call, &[2]),
            node(address, CallKind::Call, &[3]),
        ];

        assert_eq!(find_pc_target(&arena, 1, 0, 3).unwrap().node_index, 2);
        assert_eq!(find_pc_target(&arena, 1, 0, 1).unwrap().node_index, 0);
    }

    #[test]
    fn does_not_search_different_address_or_creation_context() {
        let address = Address::repeat_byte(1);
        let other = Address::repeat_byte(2);
        let arena = vec![
            node(address, CallKind::Call, &[1]),
            node(other, CallKind::Call, &[2]),
            node(address, CallKind::Create, &[3]),
        ];

        assert!(find_pc_target(&arena, 0, 0, 2).is_none());
        assert!(find_pc_target(&arena, 0, 0, 3).is_none());
        assert!(pc_exists_outside_code_context(&arena, &arena[0], 2));
        assert!(pc_exists_outside_code_context(&arena, &arena[0], 3));
    }

    #[test]
    fn goto_resolves_unambiguous_bare_digits_and_reports_ambiguity() {
        let address = Address::repeat_byte(1);
        let mut context = context_with_arena(vec![node(address, CallKind::Call, &[10, 16, 42])]);
        let mut tui = TUIContext::new(&mut context);
        tui.init();

        tui.goto_pc_from_input("2a");
        assert_eq!(tui.current_step, 2);
        assert_eq!(tui.status.as_ref().unwrap().kind, StatusKind::Info);

        tui.current_step = 0;
        tui.goto_pc_from_input("10");
        assert_eq!(tui.current_step, 0);
        assert!(tui.status.as_ref().unwrap().text.contains("Ambiguous PC"));

        tui.goto_pc_from_input("d:10");
        assert_eq!(tui.current_step, 0);
        assert_eq!(tui.status.as_ref().unwrap().kind, StatusKind::Info);
    }

    #[test]
    fn goto_reports_pc_in_other_contract_without_moving() {
        let address = Address::repeat_byte(1);
        let other = Address::repeat_byte(2);
        let mut context = context_with_arena(vec![
            node(address, CallKind::Call, &[1]),
            node(other, CallKind::Call, &[42]),
        ]);
        let mut tui = TUIContext::new(&mut context);
        tui.init();

        tui.goto_pc_from_input("2a");
        assert_eq!(tui.draw_memory.inner_call_index, 0);
        assert_eq!(tui.current_step, 0);
        let status = tui.status.as_ref().unwrap();
        assert_eq!(status.kind, StatusKind::Error);
        assert!(status.text.contains("exists in another contract"));
    }

    #[test]
    fn goto_reports_ambiguous_input_in_other_contract_without_choosing_first_candidate() {
        let address = Address::repeat_byte(1);
        let other = Address::repeat_byte(2);
        let mut context = context_with_arena(vec![
            node(address, CallKind::Call, &[1]),
            node(other, CallKind::Call, &[16]),
        ]);
        let mut tui = TUIContext::new(&mut context);
        tui.init();

        tui.goto_pc_from_input("10");
        let status = tui.status.as_ref().unwrap();
        assert_eq!(status.kind, StatusKind::Error);
        assert!(status.text.starts_with("PC `10` not found"));
        assert!(status.text.contains("exists in another contract"));
    }

    #[test]
    fn pc_input_mode_handles_keys_and_blocks_normal_commands() {
        let address = Address::repeat_byte(1);
        let mut context = context_with_arena(vec![node(address, CallKind::Call, &[1, 42])]);
        let mut tui = TUIContext::new(&mut context);
        tui.init();

        assert!(matches!(tui.handle_key_event(key(KeyCode::Char('p'))), ControlFlow::Continue(())));
        assert_eq!(tui.pc_input.as_deref(), Some(""));

        let _ = tui.handle_key_event(key(KeyCode::Char('q')));
        assert_eq!(tui.pc_input.as_deref(), Some(""));
        assert_eq!(tui.current_step, 0);

        let _ = tui.handle_key_event(key(KeyCode::Char('2')));
        let _ = tui.handle_key_event(key(KeyCode::Char('a')));
        assert_eq!(tui.pc_input.as_deref(), Some("2a"));

        let _ = tui.handle_key_event(key(KeyCode::Backspace));
        assert_eq!(tui.pc_input.as_deref(), Some("2"));
        let _ = tui.handle_key_event(key(KeyCode::Char('a')));
        let _ = tui.handle_key_event(key(KeyCode::Enter));

        assert_eq!(tui.pc_input, None);
        assert_eq!(tui.current_step, 1);
        assert_eq!(tui.status.as_ref().unwrap().kind, StatusKind::Info);
    }

    #[test]
    fn pc_input_escape_cancels_without_moving() {
        let address = Address::repeat_byte(1);
        let mut context = context_with_arena(vec![node(address, CallKind::Call, &[1, 42])]);
        let mut tui = TUIContext::new(&mut context);
        tui.init();

        let _ = tui.handle_key_event(key(KeyCode::Char('p')));
        let _ = tui.handle_key_event(key(KeyCode::Char('2')));
        let _ = tui.handle_key_event(key(KeyCode::Esc));

        assert_eq!(tui.pc_input, None);
        assert_eq!(tui.current_step, 0);
        assert_eq!(tui.status, None);
    }
}