rust_simple_stack_processor/
lib.rs

1use std::convert::TryFrom;
2use thiserror::Error;
3
4#[cfg(test)]
5mod tests;
6
7/// Gas limit for execution: unlimited or limited to a number of steps.
8#[derive(Debug, Clone, PartialEq)]
9pub enum GasLimit {
10    Unlimited,
11    Limited(u64),
12}
13
14/// Errors that can occur during stack machine execution.
15#[derive(Error, Debug, Clone, PartialEq)]
16pub enum StackMachineError {
17    #[error("Division by zero")]
18    DivisionByZero { failing_opcode: Opcode },
19
20    #[error("Numeric overflow")]
21    NumericOverflow { failing_opcode: Opcode },
22
23    #[error("TryFromInt error")]
24    TryFromIntError(#[from] std::num::TryFromIntError),
25
26    #[error("The internal number stack has underflowed (do you have too many POPs?)")]
27    NumberStackUnderflow,
28
29    #[error("The internal loop stack has underflowed (are you missing a loop start opcode?)")]
30    LoopStackUnderflow,
31
32    #[error(
33        "The internal scratch stack has underflowed (do you have too many scratch stack POPs?)"
34    )]
35    ScratchStackUnderflow,
36
37    #[error("Invalid Cell Operation (perhaps your parameters are incorrect?)")]
38    InvalidCellOperation,
39
40    #[error("Unhandled trap id: {unhandled_trap_id}")]
41    UnhandledTrap { unhandled_trap_id: i64 },
42
43    #[error("You used too much gas during execution (used {gas_used:?}, gas_limit {gas_limit:?}")]
44    RanOutOfGas { gas_used: u64, gas_limit: GasLimit },
45
46    #[error("Unknown StackMachineError")]
47    UnknownError,
48}
49
50/// Result of trap handling.
51#[derive(Debug, Clone, PartialEq)]
52pub enum TrapHandled {
53    Handled,
54    NotHandled,
55}
56
57/// Trait for trap handlers implementing the Chain of Command pattern.
58pub trait HandleTrap {
59    fn handle_trap(
60        &mut self,
61        trap_id: i64,
62        st: &mut StackMachineState,
63    ) -> Result<TrapHandled, StackMachineError>;
64}
65
66/// A trap handler that handles a specific trap id with a closure.
67pub struct TrapHandler<'a> {
68    handled_trap: i64,
69    to_run: Box<dyn Fn(i64, &mut StackMachineState) -> Result<TrapHandled, StackMachineError> + 'a>,
70}
71
72impl<'a> TrapHandler<'a> {
73    /// Create a new TrapHandler for a specific trap id.
74    pub fn new<C>(handled_trap: i64, f: C) -> TrapHandler<'a>
75    where
76        C: Fn(i64, &mut StackMachineState) -> Result<TrapHandled, StackMachineError> + 'a,
77    {
78        TrapHandler {
79            handled_trap,
80            to_run: Box::new(f),
81        }
82    }
83}
84
85impl<'a> HandleTrap for TrapHandler<'a> {
86    fn handle_trap(
87        &mut self,
88        trap_number: i64,
89        st: &mut StackMachineState,
90    ) -> Result<TrapHandled, StackMachineError> {
91        if trap_number == self.handled_trap {
92            (self.to_run)(self.handled_trap, st)
93        } else {
94            Ok(TrapHandled::NotHandled)
95        }
96    }
97}
98
99/// Opcodes supported by the stack machine.
100#[derive(Debug, Clone, PartialEq)]
101pub enum Opcode {
102    JMP,
103    JR,
104    JRZ,
105    JRNZ,
106    CALL,
107    CMPZ,
108    CMPNZ,
109    LDI(i64),
110    DROP,
111    SWAP,
112    SWAP2,
113    RET,
114    ADD,
115    SUB,
116    MUL,
117    DIV,
118    NOT,
119    DUP,
120    DUP2,
121    TRAP,
122    NOP,
123    PUSHLP,
124    INCLP,
125    ADDLP,
126    GETLP,
127    GETLP2,
128    DROPLP,
129    CMPLOOP,
130    OVER2,
131    GtR,
132    RGt,
133    RAt,
134    GtR2,
135    RGt2,
136    RAt2,
137    AND,
138    NEWCELLS,
139    MOVETOCELLS,
140    MOVEFROMCELLS,
141}
142
143/// Internal state of the stack machine.
144pub struct StackMachineState {
145    pub number_stack: Vec<i64>,
146    pub scratch_stack: Vec<i64>,
147    return_stack: Vec<usize>,
148    // current index, max_index
149    loop_stack: Vec<(i64, i64)>,
150    cells: Vec<i64>,
151    pub opcodes: Vec<Opcode>,
152    pc: usize,
153    gas_used: u64,
154}
155
156impl Default for StackMachineState {
157    fn default() -> Self {
158        StackMachineState {
159            number_stack: Vec::new(),
160            scratch_stack: Vec::new(),
161            return_stack: Vec::new(),
162            loop_stack: Vec::new(),
163            cells: Vec::new(),
164            opcodes: Vec::new(),
165            pc: 0,
166            gas_used: 0,
167        }
168    }
169}
170
171impl StackMachineState {
172    /// Returns the amount of gas used so far.
173    pub fn gas_used(&self) -> u64 {
174        self.gas_used
175    }
176}
177
178/// The stack machine itself, holding state and trap handlers.
179pub struct StackMachine {
180    pub st: StackMachineState,
181    pub trap_handlers: Vec<Box<dyn HandleTrap>>,
182}
183
184impl Default for StackMachine {
185    fn default() -> StackMachine {
186        StackMachine {
187            st: StackMachineState::default(),
188            trap_handlers: Vec::new(),
189        }
190    }
191}
192
193
194impl StackMachine {
195    fn pop_number_stack(&mut self) -> Result<i64, StackMachineError> {
196        self.st.number_stack.pop().ok_or(StackMachineError::NumberStackUnderflow)
197    }
198
199    fn push_number_stack(&mut self, value: i64) {
200        self.st.number_stack.push(value);
201    }
202
203    fn pop_scratch_stack(&mut self) -> Result<i64, StackMachineError> {
204        self.st.scratch_stack.pop().ok_or(StackMachineError::ScratchStackUnderflow)
205    }
206
207    fn push_scratch_stack(&mut self, value: i64) {
208        self.st.scratch_stack.push(value);
209    }
210
211    fn peek_scratch_stack(&self) -> Result<i64, StackMachineError> {
212        self.st.scratch_stack.last().copied().ok_or(StackMachineError::ScratchStackUnderflow)
213    }
214
215    fn execute_binary_op<F>(&mut self, op: F, opcode: &Opcode) -> Result<(), StackMachineError>
216    where
217        F: FnOnce(i64, i64) -> Option<i64>,
218    {
219        let second = self.pop_number_stack()?;
220        let first = self.pop_number_stack()?;
221        let result = op(first, second).ok_or(StackMachineError::NumericOverflow {
222            failing_opcode: opcode.clone(),
223        })?;
224        self.push_number_stack(result);
225        Ok(())
226    }
227
228    fn execute_division(&mut self, opcode: &Opcode) -> Result<(), StackMachineError> {
229        let divisor = self.pop_number_stack()?;
230        let dividend = self.pop_number_stack()?;
231        let result = dividend.checked_div(divisor).ok_or(StackMachineError::DivisionByZero {
232            failing_opcode: opcode.clone(),
233        })?;
234        self.push_number_stack(result);
235        Ok(())
236    }
237
238    fn execute_jump_relative(&mut self, condition: Option<bool>) -> Result<bool, StackMachineError> {
239        let offset = self.pop_number_stack()?;
240        let should_jump = if let Some(cond) = condition {
241            let value = self.pop_number_stack()?;
242            (value == 0) == cond
243        } else {
244            true
245        };
246
247        if should_jump {
248            let new_offset = i64::try_from(self.st.pc)? + offset;
249            self.st.pc = usize::try_from(new_offset)?;
250            Ok(true)
251        } else {
252            Ok(false)
253        }
254    }
255    /// JR(*) is relative from the JR(*) instruction,
256    /// 0 would jump back onto the JR instruction
257    /// -1 Would jump back to the instruction before the JR(*}) instruction
258    /// 1 Would jump to the instruction after the JR(*) instruction
259    ///
260    /// TRAPs always have a numeric code on the number stack to define which TRAP is being called
261    ///
262    /// CMPLOOP
263    /// pushes 1 on the stack if the loop counter is greater than or equal to the max
264    /// pushes 0 on the stack if the loop counter is less than the max
265    pub fn execute(
266        &mut self,
267        starting_point: usize,
268        gas_limit: GasLimit,
269    ) -> Result<(), StackMachineError> {
270        self.st.gas_used = 0;
271        self.st.pc = starting_point;
272        loop {
273            let mut pc_reset = false;
274            let opcode = self
275                .st
276                .opcodes
277                .get(self.st.pc)
278                .ok_or(StackMachineError::UnknownError)?;
279            match opcode {
280                Opcode::JMP => {
281                    let target = usize::try_from(self.pop_number_stack()?)?;
282                    self.st.pc = target;
283                    pc_reset = true;
284                }
285                Opcode::JR => {
286                    pc_reset = self.execute_jump_relative(None)?;
287                }
288                Opcode::CALL => {
289                    self.st.return_stack.push(self.st.pc + 1);
290                    let target = usize::try_from(self.pop_number_stack()?)?;
291                    self.st.pc = target;
292                    pc_reset = true;
293                }
294                Opcode::CMPZ => {
295                    let value = self.pop_number_stack()?;
296                    self.push_number_stack(if value == 0 { -1 } else { 0 });
297                }
298                Opcode::CMPNZ => {
299                    let value = self.pop_number_stack()?;
300                    self.push_number_stack(if value == 0 { 0 } else { -1 });
301                }
302                Opcode::JRZ => {
303                    pc_reset = self.execute_jump_relative(Some(true))?;
304                }
305                Opcode::JRNZ => {
306                    pc_reset = self.execute_jump_relative(Some(false))?;
307                }
308                Opcode::LDI(immediate_value) => self.push_number_stack(*immediate_value),
309                Opcode::DROP => {
310                    let _ = self.pop_number_stack()?;
311                }
312                Opcode::RET => {
313                    if let Some(return_address) = self.st.return_stack.pop() {
314                        self.st.pc = return_address;
315                        pc_reset = true;
316                    } else {
317                        return Ok(());
318                    }
319                }
320                Opcode::GtR => {
321                    let value = self.pop_number_stack()?;
322                    self.push_scratch_stack(value);
323                }
324                Opcode::RGt => {
325                    let value = self.pop_scratch_stack()?;
326                    self.push_number_stack(value);
327                }
328                Opcode::RAt => {
329                    let value = self.peek_scratch_stack()?;
330                    self.push_number_stack(value);
331                }
332                Opcode::GtR2 => {
333                    let x = self.pop_number_stack()?;
334                    let y = self.pop_number_stack()?;
335                    self.push_scratch_stack( y);
336                    self.push_scratch_stack( x);
337                }
338                Opcode::RGt2 => {
339                    let x = self.pop_scratch_stack()?;
340                    let y = self.pop_scratch_stack()?;
341                    self.push_number_stack(y);
342                    self.push_number_stack(x);
343                }
344                Opcode::RAt2 => {
345                    let x = self.pop_scratch_stack()?;
346                    let y = self.pop_scratch_stack()?;
347                    self.push_scratch_stack(y);
348                    self.push_scratch_stack(x);
349                    self.push_number_stack(y);
350                    self.push_number_stack(x);
351                }
352                Opcode::ADD => {
353                    self.execute_binary_op(|a, b| a.checked_add(b), &opcode.clone())?;
354                }
355                Opcode::SUB => {
356                    self.execute_binary_op(|a, b| b.checked_sub(a), &opcode.clone())?;
357                }
358                Opcode::MUL => {
359                    self.execute_binary_op(|a, b| a.checked_mul(b), &opcode.clone())?;
360                }
361                Opcode::DIV => {
362                    self.execute_division(&opcode.clone())?;
363                }
364                Opcode::NOT => {
365                    let x = self.pop_number_stack()?;
366                    self.push_number_stack(if x == 0 { 1 } else { 0 });
367                }
368                Opcode::DUP => {
369                    let x = self.pop_number_stack()?;
370                    self.push_number_stack(x);
371                    self.push_number_stack(x);
372                }
373                Opcode::DUP2 => {
374                    let x = self.pop_number_stack()?;
375                    let y = self.pop_number_stack()?;
376                    self.push_number_stack(y);
377                    self.push_number_stack(x);
378                    self.push_number_stack(y);
379                    self.push_number_stack(x);
380                }
381                Opcode::OVER2 => {
382                    let x4 = self.pop_number_stack()?;
383                    let x3 = self.pop_number_stack()?;
384                    let x2 = self.pop_number_stack()?;
385                    let x1 = self.pop_number_stack()?;
386                    self.push_number_stack(x1);
387                    self.push_number_stack(x2);
388                    self.push_number_stack(x3);
389                    self.push_number_stack(x4);
390                    self.push_number_stack(x1);
391                    self.push_number_stack(x2);
392                }
393                Opcode::SWAP => {
394                    let x = self.pop_number_stack()?;
395                    let y = self.pop_number_stack()?;
396                    self.push_number_stack(x);
397                    self.push_number_stack(y);
398                }
399                Opcode::SWAP2 => {
400                    let x4 = self.pop_number_stack()?;
401                    let x3 = self.pop_number_stack()?;
402                    let x2 = self.pop_number_stack()?;
403                    let x1 = self.pop_number_stack()?;
404                    self.push_number_stack(x3);
405                    self.push_number_stack(x4);
406                    self.push_number_stack(x1);
407                    self.push_number_stack(x2);
408                }
409                Opcode::TRAP => {
410                    let trap_id = self.pop_number_stack()?;
411                    for h in self.trap_handlers.iter_mut() {
412                        if let TrapHandled::Handled = h.handle_trap(trap_id, &mut self.st)? {
413                            return Ok(());
414                        }
415                    }
416                    return Err(StackMachineError::UnhandledTrap {
417                        unhandled_trap_id: trap_id,
418                    });
419                }
420                Opcode::NOP => {}
421                Opcode::PUSHLP => {
422                    let current_index = self.pop_number_stack()?;
423                    let max_index = self.pop_number_stack()?;
424                    self.st.loop_stack.push((current_index, max_index));
425                }
426                Opcode::INCLP => {
427                    if let Some((current_index, _)) = self.st.loop_stack.last_mut() {
428                        *current_index += 1;
429                    } else {
430                        return Err(StackMachineError::LoopStackUnderflow);
431                    }
432                }
433                Opcode::ADDLP => {
434                    let increment = self.pop_number_stack()?;
435                    if let Some((current_index, _)) = self.st.loop_stack.last_mut() {
436                        *current_index += increment;
437                    } else {
438                        return Err(StackMachineError::LoopStackUnderflow);
439                    }
440                }
441                Opcode::GETLP => {
442                    let (current_index, _) = self
443                        .st
444                        .loop_stack
445                        .last()
446                        .ok_or(StackMachineError::LoopStackUnderflow)?;
447                    self.st.number_stack.push(*current_index);
448                }
449                Opcode::GETLP2 => {
450                    if self.st.loop_stack.len() < 2 {
451                        return Err(StackMachineError::LoopStackUnderflow);
452                    }
453                    let (current_index, _) = self
454                        .st
455                        .loop_stack
456                        .get(self.st.loop_stack.len() - 2)
457                        .ok_or(StackMachineError::LoopStackUnderflow)?;
458                    self.st.number_stack.push(*current_index);
459                }
460                Opcode::DROPLP => {
461                    self.st
462                        .loop_stack
463                        .pop()
464                        .ok_or(StackMachineError::LoopStackUnderflow)?;
465                }
466                Opcode::CMPLOOP => {
467                    let (current_index, max_index) = self
468                        .st
469                        .loop_stack
470                        .last()
471                        .ok_or(StackMachineError::LoopStackUnderflow)?;
472                    self.st
473                        .number_stack
474                        .push(if *current_index >= *max_index { 1 } else { 0 });
475                }
476                Opcode::AND => {
477                    let x = self.pop_number_stack()?;
478                    let y = self.pop_number_stack()?;
479                    self.push_number_stack(x & y);
480                }
481                Opcode::NEWCELLS => {
482                    let num_cells = usize::try_from(self.pop_number_stack()?)
483                        .map_err(|_| StackMachineError::InvalidCellOperation)?;
484                    let newaddress = self.st.cells.len();
485                    self.st
486                        .cells
487                        .resize_with(newaddress + num_cells, Default::default);
488                }
489                Opcode::MOVETOCELLS => {
490                    let num_cells = usize::try_from(self.pop_number_stack()?)
491                        .map_err(|_| StackMachineError::InvalidCellOperation)?;
492                    let address = usize::try_from(self.pop_number_stack()?)
493                        .map_err(|_| StackMachineError::InvalidCellOperation)?;
494                    if num_cells < 1 || self.st.cells.len() < address + num_cells {
495                        return Err(StackMachineError::InvalidCellOperation);
496                    }
497                    for i in address..address + num_cells {
498                        self.st.cells[i] = self.pop_number_stack()?;
499                    }
500                }
501                Opcode::MOVEFROMCELLS => {
502                    let num_cells = usize::try_from(self.pop_number_stack()?)
503                        .map_err(|_| StackMachineError::InvalidCellOperation)?;
504                    let address = usize::try_from(self.pop_number_stack()?)
505                        .map_err(|_| StackMachineError::InvalidCellOperation)?;
506                    if num_cells < 1 || self.st.cells.len() < address + num_cells {
507                        return Err(StackMachineError::InvalidCellOperation);
508                    }
509                    for i in (address..address + num_cells).rev() {
510                        self.push_number_stack(self.st.cells[i]);
511                    }
512                }
513            }
514            if !pc_reset {
515                self.st.pc += 1;
516            }
517
518            self.st.gas_used += 1;
519
520            if let GasLimit::Limited(limit) = gas_limit {
521                if self.st.gas_used > limit {
522                    return Err(StackMachineError::RanOutOfGas {
523                        gas_used: self.st.gas_used,
524                        gas_limit,
525                    });
526                }
527            }
528        }
529    }
530}
rust_simple_stack_processor/lib.rs

rust_simple_stack_processor/
lib.rs
