rust_simple_stack_processor/
lib.rs

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