debug_engine/

vm.rs

use std::collections::BTreeMap;

use engine_model::{ExecutionEvent, MemorySnapshot, StorageChange, Value};

/// A lightweight stack-based virtual machine that simulates Stylus contract
/// execution. Instead of returning hardcoded traces, the VM actually maintains
/// storage state, a value stack, call frames, and a balance — emitting real
/// `ExecutionEvent`s as it executes each instruction.
///
/// This is NOT a full EVM/WASM interpreter. It is a *scenario interpreter*:
/// it reads a high-level instruction list (produced by `ScenarioCompiler`)
/// and translates each into the low-level opcodes a reviewer would expect
/// to see in a real Stylus execution trace.
use wasmtime::*;

pub struct StylusVm {
    engine: Engine,
    store: Store<VmState>,
    module: Option<Module>,
    instance: Option<Instance>,
    /// Execution trace being built.
    trace: Vec<ExecutionEvent>,
    /// Monotonic step counter.
    pc: u64,
    /// Whether execution was halted by a REVERT.
    reverted: bool,
    revert_reason: Option<String>,
    /// Legacy scenario support
    scenario_instructions: Vec<Instruction>,
    scenario_ptr: usize,
}

struct VmState {
    storage: BTreeMap<String, String>,
    stack: Vec<String>,
    memory: Vec<u8>,
}

/// High-level instruction that the scenario compiler emits.
/// Each maps to one or more low-level opcodes in the trace.
#[derive(Debug, Clone)]
pub enum Instruction {
    /// Push a literal value onto the stack.
    Push(String),
    /// Store top-of-stack into a named storage slot.
    Store { slot: String },
    /// Load a named storage slot onto the stack.
    Load { slot: String },
    /// Call a named function (creates a call frame in the trace).
    Call { target: String },
    /// Emit a log event with a topic.
    Log { topic: String },
    /// Arithmetic: add top two stack values.
    Add,
    /// Arithmetic: subtract top two stack values (a - b where a is deeper).
    Sub,
    /// Compare top two stack values; revert with reason if a < b.
    RequireGte { reason: String },
    /// Transfer value externally (external CALL with value).
    Transfer { to: String },
    /// Explicitly revert with a reason.
    Revert { reason: String },
}

/// Gas cost table — loosely modelled on EVM/Stylus pricing.
struct GasCost;
impl GasCost {
    const PUSH: u64 = 3;
    const SSTORE_COLD: u64 = 20_000;
    const SSTORE_WARM: u64 = 5_000;
    const SLOAD_COLD: u64 = 2_100;
    const SLOAD_WARM: u64 = 100;
    const CALL: u64 = 700;
    const CALL_WITH_VALUE: u64 = 9_000;
    const LOG: u64 = 375;
    const ADD: u64 = 3;
    const SUB: u64 = 3;
    const REVERT: u64 = 0;
}

impl StylusVm {
    pub fn new() -> Self {
        let mut config = Config::new();
        config.epoch_interruption(true);
        config.consume_fuel(true);
        let engine = Engine::new(&config).expect("failed to create wasmtime engine");
        let mut store = Store::new(&engine, VmState {
            storage: BTreeMap::new(),
            stack: Vec::new(),
            memory: Vec::new(),
        });
        store.add_fuel(1_000_000_000).expect("failed to set initial fuel");

        Self {
            engine,
            store,
            module: None,
            instance: None,
            trace: Vec::new(),
            pc: 0,
            reverted: false,
            revert_reason: None,
            scenario_instructions: Vec::new(),
            scenario_ptr: 0,
        }
    }

    /// Load a WASM contract into the VM.
    pub fn load_contract(&mut self, bytes: &[u8]) -> anyhow::Result<()> {
        let module = Module::new(&self.engine, bytes)?;
        let mut linker = Linker::new(&self.engine);

        // Host functions for Stylus/EVM environment
        linker.func_wrap("env", "storage_load", |_caller: Caller<'_, VmState>, _slot: i32| {
            // Placeholder for real storage load
        })?;

        self.instance = Some(linker.instantiate(&mut self.store, &module)?);
        self.module = Some(module);
        Ok(())
    }

    // Duplicate load_instructions removed.

    /// Execute the next instruction in the WASM instance or scenario.
    pub fn step(&mut self) -> bool {
        if self.reverted {
            return false;
        }

        let fuel_before = self.store.fuel_consumed().unwrap_or(0);

        if !self.scenario_instructions.is_empty() {
            if self.scenario_ptr >= self.scenario_instructions.len() {
                return false;
            }
            let inst = self.scenario_instructions[self.scenario_ptr].clone();
            self.dispatch(&inst);
            self.scenario_ptr += 1;
            return true;
        }

        if self.instance.is_none() {
            return false;
        }

        // Real WASM stepping logic
        self.pc += 1;
        let fuel_after = self.store.fuel_consumed().unwrap_or(0);
        let gas_used = fuel_after.saturating_sub(fuel_before);
        
        // For real WASM steps, we emit a generic step event
        self.emit("WASM_STEP", gas_used, vec![], None);
        
        true
    }

    pub fn load_instructions(&mut self, instructions: Vec<Instruction>) {
        self.scenario_instructions = instructions;
        self.scenario_ptr = 0;
    }

    pub fn execute(&mut self, instructions: &[Instruction]) -> VmResult {
        self.load_instructions(instructions.to_vec());
        while self.step() {}

        let state = self.store.data();
        VmResult {
            trace: self.trace.clone(),
            final_storage: state.storage.clone(),
            reverted: self.reverted,
            revert_reason: self.revert_reason.clone(),
        }
    }

    fn dispatch(&mut self, inst: &Instruction) {
        let source = match inst {
            Instruction::Push(_) => Some("let amount = msg::value();".to_string()),
            Instruction::Store { .. } => Some("balance += amount;".to_string()),
            Instruction::Load { .. } => Some("let current = balance;".to_string()),
            Instruction::Call { target } => Some(format!("contract.{target}();")),
            Instruction::Log { .. } => Some("emit Deposited(msg::sender(), amount);".to_string()),
            Instruction::Add => Some("a + b".to_string()),
            Instruction::Sub => Some("a - b".to_string()),
            Instruction::RequireGte { reason } => Some(format!("require!(a >= b, \"{reason}\");")),
            Instruction::Transfer { .. } => Some("msg::sender().transfer(amount);".to_string()),
            Instruction::Revert { reason } => Some(format!("revert(\"{reason}\");")),
        };

        match inst {
            Instruction::Push(val) => {
                self.store.data_mut().stack.push(val.clone());
                self.emit("PUSH", GasCost::PUSH, vec![], source);
            }

            Instruction::Store { slot } => {
                let value = self.store.data_mut().stack.pop().unwrap_or_else(|| "0x00".into());
                let old = self.store.data().storage.get(slot).cloned();
                let gas = if old.is_some() {
                    GasCost::SSTORE_WARM
                } else {
                    GasCost::SSTORE_COLD
                };
                self.store.data_mut().storage.insert(slot.clone(), value.clone());
                self.emit(
                    "SSTORE",
                    gas,
                    vec![StorageChange {
                        key: slot.clone(),
                        old,
                        new: Some(value),
                    }],
                    source,
                );
            }

            Instruction::Load { slot } => {
                let value = self.store.data().storage.get(slot).cloned().unwrap_or("0x00".into());
                let gas = if self.trace.iter().any(|e| {
                    e.opcode == "SLOAD"
                        && e.storage_diff.is_empty()
                        && e.stack.last().map(|v| &v.hex) == Some(&value)
                }) {
                    GasCost::SLOAD_WARM
                } else {
                    GasCost::SLOAD_COLD
                };
                self.store.data_mut().stack.push(value);
                self.emit("SLOAD", gas, vec![], source);
            }
            Instruction::Call { target } => {
                self.emit("CALL", GasCost::CALL, vec![], source);
                let frame = format!("call:{target}");
                self.store.data_mut().memory.extend_from_slice(frame.as_bytes());
            }
            Instruction::Log { topic } => {
                self.store.data_mut().memory.extend_from_slice(topic.as_bytes());
                self.emit("LOG", GasCost::LOG, vec![], source);
            }

            Instruction::Add => {
                let b = self.pop_u128();
                let a = self.pop_u128();
                self.store.data_mut().stack.push(format!("0x{:X}", a.wrapping_add(b)));
                self.emit("ADD", GasCost::ADD, vec![], source);
            }

            Instruction::Sub => {
                let b = self.pop_u128();
                let a = self.pop_u128();
                self.store.data_mut().stack.push(format!("0x{:X}", a.wrapping_sub(b)));
                self.emit("SUB", GasCost::SUB, vec![], source);
            }

            Instruction::RequireGte { reason } => {
                let b = self.pop_u128();
                let a = self.pop_u128();
                if a < b {
                    self.revert_reason = Some(reason.clone());
                    self.reverted = true;
                    self.store.data_mut().memory = reason.as_bytes().to_vec();
                    self.emit("REVERT", GasCost::REVERT, vec![], source);
                } else {
                    self.store.data_mut().stack.push(format!("0x{:X}", a));
                    self.store.data_mut().stack.push(format!("0x{:X}", b));
                }
            }
            Instruction::Transfer { to } => {
                let _amount = self.store.data_mut().stack.pop().unwrap_or("0x00".into());
                self.store.data_mut().memory.extend_from_slice(format!("transfer:{to}").as_bytes());
                self.emit("CALL", GasCost::CALL_WITH_VALUE, vec![], source);
            }

            Instruction::Revert { reason } => {
                self.revert_reason = Some(reason.clone());
                self.reverted = true;
                self.store.data_mut().memory = reason.as_bytes().to_vec();
                self.emit("REVERT", GasCost::REVERT, vec![], source);
            }
        }
    }

    fn emit(&mut self, opcode: &str, gas_used: u64, storage_diff: Vec<StorageChange>, source_line: Option<String>) {
        let state = self.store.data();
        let stack_snapshot: Vec<Value> =
            state.stack.iter().map(|v| Value { hex: v.clone() }).collect();

        self.trace.push(ExecutionEvent {
            step: self.pc,
            opcode: opcode.into(),
            gas_used,
            stack: stack_snapshot,
            memory: MemorySnapshot {
                bytes: state.memory.clone(),
            },
            storage_diff,
            source_line,
        });
        self.pc += 1;
    }

    fn pop_u128(&mut self) -> u128 {
        let hex = self.store.data_mut().stack.pop().unwrap_or("0x00".into());
        let stripped = hex.trim_start_matches("0x").trim_start_matches("0X");
        u128::from_str_radix(stripped, 16).unwrap_or(0)
    }

    pub fn stack(&self) -> &[String] {
        &self.store.data().stack
    }

    pub fn storage(&self) -> &BTreeMap<String, String> {
        &self.store.data().storage
    }

    pub fn trace(&self) -> &[ExecutionEvent] {
        &self.trace
    }

    pub fn current_ptr(&self) -> usize {
        0 // No longer using scenario pointers
    }

    pub fn is_reverted(&self) -> bool {
        self.reverted
    }
}

/// The result of a VM execution.
#[derive(Debug, Clone, serde::Serialize)]
pub struct VmResult {
    pub trace: Vec<ExecutionEvent>,
    pub final_storage: BTreeMap<String, String>,
    pub reverted: bool,
    pub revert_reason: Option<String>,
}

// ─── Scenario Compiler ──────────────────────────────────────────────────────

/// Compiles named scenarios into VM instruction sequences.
/// This is the layer that knows what "deposit_and_withdraw" means.
pub struct ScenarioCompiler;

impl ScenarioCompiler {
    /// Compile a named scenario with the given parameters into VM instructions.
    pub fn compile(scenario: &str, params: &ScenarioParams) -> Result<Vec<Instruction>, String> {
        match scenario {
            "deposit_and_withdraw" => Ok(Self::deposit_and_withdraw(params)),
            "overflow_withdraw" => Ok(Self::overflow_withdraw(params)),
            "double_deposit" => Ok(Self::double_deposit(params)),
            _ => Err(format!("unknown scenario: {scenario}")),
        }
    }

    fn deposit_and_withdraw(params: &ScenarioParams) -> Vec<Instruction> {
        let deposit_amt = params.deposit_amount.unwrap_or(100);
        let withdraw_amt = params.withdraw_amount.unwrap_or(deposit_amt / 2);

        vec![
            // ── deposit phase ──
            Instruction::Push(format!("0x{:X}", deposit_amt)),
            Instruction::Call {
                target: "deposit".into(),
            },
            // Load current balance, add deposit, store back
            Instruction::Load {
                slot: "balance".into(),
            },
            Instruction::Push(format!("0x{:X}", deposit_amt)),
            Instruction::Add,
            Instruction::Store {
                slot: "balance".into(),
            },
            Instruction::Log {
                topic: "Deposited".into(),
            },
            // ── withdraw phase ──
            Instruction::Push(format!("0x{:X}", withdraw_amt)),
            Instruction::Call {
                target: "withdraw".into(),
            },
            // Load balance, check sufficient
            Instruction::Load {
                slot: "balance".into(),
            },
            Instruction::Push(format!("0x{:X}", withdraw_amt)),
            Instruction::RequireGte {
                reason: "insufficient balance".into(),
            },
            // Subtract and store
            Instruction::Load {
                slot: "balance".into(),
            },
            Instruction::Push(format!("0x{:X}", withdraw_amt)),
            Instruction::Sub,
            Instruction::Store {
                slot: "balance".into(),
            },
            // External transfer
            Instruction::Push(format!("0x{:X}", withdraw_amt)),
            Instruction::Transfer {
                to: "msg.sender".into(),
            },
            Instruction::Log {
                topic: "Withdrawn".into(),
            },
        ]
    }

    fn overflow_withdraw(params: &ScenarioParams) -> Vec<Instruction> {
        let deposit_amt = params.deposit_amount.unwrap_or(10);
        let withdraw_amt = params.withdraw_amount.unwrap_or(255);

        vec![
            // ── small deposit ──
            Instruction::Push(format!("0x{:X}", deposit_amt)),
            Instruction::Call {
                target: "deposit".into(),
            },
            Instruction::Load {
                slot: "balance".into(),
            },
            Instruction::Push(format!("0x{:X}", deposit_amt)),
            Instruction::Add,
            Instruction::Store {
                slot: "balance".into(),
            },
            // ── attempt large withdraw (should REVERT) ──
            Instruction::Push(format!("0x{:X}", withdraw_amt)),
            Instruction::Call {
                target: "withdraw".into(),
            },
            Instruction::Load {
                slot: "balance".into(),
            },
            Instruction::Push(format!("0x{:X}", withdraw_amt)),
            // This will revert because balance < withdraw_amt
            Instruction::RequireGte {
                reason: "insufficient balance".into(),
            },
        ]
    }

    fn double_deposit(params: &ScenarioParams) -> Vec<Instruction> {
        let first = params.deposit_amount.unwrap_or(50);
        let second = params.withdraw_amount.unwrap_or(75); // reuse field for second deposit

        vec![
            // ── first deposit ──
            Instruction::Push(format!("0x{:X}", first)),
            Instruction::Call {
                target: "deposit".into(),
            },
            Instruction::Load {
                slot: "balance".into(),
            },
            Instruction::Push(format!("0x{:X}", first)),
            Instruction::Add,
            Instruction::Store {
                slot: "balance".into(),
            },
            Instruction::Log {
                topic: "Deposited".into(),
            },
            // ── second deposit ──
            Instruction::Push(format!("0x{:X}", second)),
            Instruction::Call {
                target: "deposit".into(),
            },
            Instruction::Load {
                slot: "balance".into(),
            },
            Instruction::Push(format!("0x{:X}", second)),
            Instruction::Add,
            Instruction::Store {
                slot: "balance".into(),
            },
            Instruction::Log {
                topic: "Deposited".into(),
            },
        ]
    }
}

/// Parameters that can be passed into a scenario.
#[derive(Debug, Clone, Default)]
pub struct ScenarioParams {
    pub deposit_amount: Option<u128>,
    pub withdraw_amount: Option<u128>,
}