Compiler Explorer

Source code

// Type your code here, or load an example.

// As of Rust 1.75, small functions are automatically
// marked as `#[inline]` so they will not show up in
// the output when compiling with optimisations. Use
// `#[unsafe(no_mangle)]` or `#[inline(never)]` to
// work around this issue.
// See https://github.com/compiler-explorer/compiler-explorer/issues/5939
#[unsafe(no_mangle)]
pub fn square(bytecode: Vec<u8>) -> i32 {
    let mut interpreter = SwitchInterpreter::new(bytecode, 5);
    std::hint::black_box(interpreter.run())
}

#[allow(non_camel_case_types)]
#[repr(u8)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Op {
    noop,
    ret,
    goto,
    pop,

if_icmpgt,

iinc,
    iadd,
    isub,

iconst,
    iconst0,
    iconst1,
    iconst2,
    iconst3,
    iconst4,
    iconst5,

iload,
    iload0,
    iload1,
    iload2,
    iload3,

istore,
    istore0,
    istore1,
    istore2,
    istore3,
}

// If you use `main()`, declare it as `pub` to see it in the output:
// pub fn main() { ... }
// const STACK_SIZE: usize = 4096 / 4; // 4KB / 4 bytes per i32 = 1024 elements
const STACK_SIZE: usize = 128 / 4; // 4KB / 4 bytes per i32 = 1024 elements

pub struct SwitchInterpreter {
    bytecode: Vec<u8>,
    // Stack-allocated storage for locals and operand stack
    stack_storage: [i32; STACK_SIZE],
    num_locals: usize,
    // Stack pointer (index into stack_storage, grows downward from STACK_SIZE)
    sp: usize,
    pc: usize,
}

impl SwitchInterpreter {
    pub fn new(bytecode: Vec<u8>, num_locals: u8) -> Self {
        let num_locals = num_locals as usize;
        Self {
            bytecode,
            stack_storage: [0; STACK_SIZE],
            num_locals,
            // Stack starts at the top, locals are at the bottom [0..num_locals)
            sp: num_locals,
            pc: 0,
        }
    }

#[inline(always)]
    fn get_local(&self, idx: usize) -> i32 {
        unsafe { *self.stack_storage.get_unchecked(idx) }
    }

#[inline(always)]
    fn set_local(&mut self, idx: usize, value: i32) {
        unsafe { *self.stack_storage.get_unchecked_mut(idx) = value }
    }

#[inline(always)]
    fn fetch_op(&self) -> Op {
        unsafe { std::mem::transmute(*self.bytecode.get_unchecked(self.pc)) }
    }

#[inline(always)]
    fn fetch_byte(&self, offset: usize) -> u8 {
        unsafe { *self.bytecode.get_unchecked(self.pc + offset) }
    }

#[inline(always)]
    fn fetch_i16(&self, offset: usize) -> i16 {
        unsafe {
            i16::from_le_bytes([
                *self.bytecode.get_unchecked(self.pc + offset),
                *self.bytecode.get_unchecked(self.pc + offset + 1),
            ])
        }
    }

#[inline(always)]
    fn push(&mut self, value: i32) {
        unsafe {
            *self.stack_storage.get_unchecked_mut(self.sp) = value;
            self.sp += 1;
        }
    }

#[inline(always)]
    fn pop(&mut self) -> i32 {
        unsafe {
            self.sp -= 1;
            *self.stack_storage.get_unchecked(self.sp)
        }
    }

pub fn run(&mut self) -> i32 {
        loop {
            let op = self.fetch_op();

match op {
                Op::noop => {
                    self.pc += 1;
                }
                Op::pop => {
                    self.pop();
                }
                Op::ret => {
                    return if self.sp == self.num_locals {
                        0
                    } else {
                        self.pop()
                    };
                }

Op::goto => {
                    let offset = self.fetch_i16(1);
                    self.pc = (self.pc as i32 + offset as i32) as usize;
                    continue;
                }

Op::if_icmpgt => {
                    let b = self.pop();
                    let a = self.pop();
                    let offset = self.fetch_i16(1);

if a > b {
                        self.pc = (self.pc as i32 + offset as i32) as usize;
                        continue;
                    }
                    self.pc += 3;
                }

Op::iinc => {
                    let local_idx = self.fetch_byte(1) as usize;
                    let incr = self.fetch_byte(2) as i8;
                    let val = self.get_local(local_idx);
                    self.set_local(local_idx, val.wrapping_add(incr as i32));
                    self.pc += 3;
                }

Op::iadd => {
                    let value2 = self.pop();
                    let value1 = self.pop();
                    self.push(value1.wrapping_add(value2));
                    self.pc += 1;
                }

Op::isub => {
                    let value2 = self.pop();
                    let value1 = self.pop();
                    self.push(value1.wrapping_sub(value2));
                    self.pc += 1;
                }

Op::iconst => {
                    let val = self.fetch_i16(1);
                    self.push(val as i32);
                    self.pc += 3;
                }

Op::iconst0 => {
                    self.push(0);
                    self.pc += 1;
                }

Op::iconst1 => {
                    self.push(1);
                    self.pc += 1;
                }

Op::iconst2 => {
                    self.push(2);
                    self.pc += 1;
                }

Op::iconst3 => {
                    self.push(3);
                    self.pc += 1;
                }

Op::iconst4 => {
                    self.push(4);
                    self.pc += 1;
                }

Op::iconst5 => {
                    self.push(5);
                    self.pc += 1;
                }

Op::iload => {
                    let local_idx = self.fetch_byte(1) as usize;
                    self.push(self.get_local(local_idx));
                    self.pc += 2;
                }

Op::iload0 => {
                    self.push(self.get_local(0));
                    self.pc += 1;
                }

Op::iload1 => {
                    self.push(self.get_local(1));
                    self.pc += 1;
                }

Op::iload2 => {
                    self.push(self.get_local(2));
                    self.pc += 1;
                }

Op::iload3 => {
                    self.push(self.get_local(3));
                    self.pc += 1;
                }

Op::istore => {
                    let local_idx = self.fetch_byte(1) as usize;
                    let value = self.pop();
                    self.set_local(local_idx, value);
                    self.pc += 2;
                }

Op::istore0 => {
                    let value = self.pop();
                    self.set_local(0, value);
                    self.pc += 1;
                }

Op::istore1 => {
                    let value = self.pop();
                    self.set_local(1, value);
                    self.pc += 1;
                }

Op::istore2 => {
                    let value = self.pop();
                    self.set_local(2, value);
                    self.pc += 1;
                }

Op::istore3 => {
                    let value = self.pop();
                    self.set_local(3, value);
                    self.pc += 1;
                }
            }
        }
    }
}