rust_simple_stack_processor/
lib.rs

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