Simple VM in Go

simple-vm.go

package main

import (
	"fmt"
	"log"
	"strings"
)

// Ops
const (
	IPUSH = iota // Push arg onto stack
	IADD         // Add top two items on the stack
	PRINT        // Print top item on the stack
	STORE        // Store top item on the stack in data mem
	FETCH        // Fetch item from data mem, push onto stack
	ISLT         // Is top of stack less than arg
	ISGT         // Is top of stack greater than arg
	JMP          // Jump to arg
	JMPT         // Jump if true
	JMPF         // Jump if false
	CALL         // Call a sub
	ARG          // Push an arg onto the stack
	RET          // Return from sub
	HALT         // SCHTAP!
)

type Op struct {
	name  string
	nargs int
}

var ops = map[int]Op{
	IPUSH: Op{"IPUSH", 1},
	IADD:  Op{"IADD", 0},
	PRINT: Op{"PRINT", 0},
	STORE: Op{"STORE", 1},
	FETCH: Op{"FETCH", 1},
	ISLT:  Op{"ISLT", 1},
	ISGT:  Op{"ISGT", 1},
	JMP:   Op{"JMP", 1},
	JMPT:  Op{"JMPT", 1},
	JMPF:  Op{"JMPF", 1},
	CALL:  Op{"CALL", 2},
	ARG:   Op{"ARG", 1},
	RET:   Op{"RET", 0},
	HALT:  Op{"HALT", 0},
}

func (o Op) String() string {
	return o.name
}

// Virtual MASHEEEEN
type VM struct {
	pc    int   // Program Counter
	sp    int   // Stack Pointer
	fp    int   // Frame Pointer
	stack []int // Stack memory
	code  []int // Code memory
	data  []int // Data memory
}

func newVM(code []int, initialPc int, dataSize int) *VM {
	vm := &VM{
		pc:    initialPc,
		sp:    -1,
		fp:    -1,
		stack: make([]int, 100), // 100 should be enough for anybody
		code:  code,
		data:  make([]int, dataSize),
	}

	return vm
}

// Print current state of the VM
func (vm *VM) trace() {
	op := ops[vm.code[vm.pc]]
	args := make([]string, op.nargs)
	for i := vm.pc + 1; i < vm.pc+1+op.nargs; i++ {
		args = append(args, fmt.Sprintf("%d", vm.code[i]))
	}
	stack := vm.stack[0 : vm.sp+1]

	fmt.Printf("%04d: %s%v\t%v\n", vm.pc, op, strings.Join(args, " "), stack)
}

func (vm *VM) dumpMem() {
	fmt.Printf("Data Memory: %v", vm.data)
}

// Push something onto the stack
func (vm *VM) sPush(i int) {
	vm.sp++
	if vm.sp > len(vm.stack) {
		log.Fatal("Stack overflow")
	}
	vm.stack[vm.sp] = i
}

// Pop something onto the stack
func (vm *VM) sPop() int {
	if vm.sp < 0 {
		log.Fatal("Stack underflow")
	}
	val := vm.stack[vm.sp]
	vm.sp--
	return val
}

// Grab the next op and increment the program counter
func (vm *VM) nextOp() int {
	val := vm.code[vm.pc]
	vm.pc++
	return val
}

// ACTIVE THE QUE^WVM
func (vm *VM) run() {
	for vm.pc < len(vm.code) {
		vm.trace()
		instr := vm.nextOp()

		switch instr {
		case IPUSH:
			val := vm.nextOp()
			vm.sPush(val)
		case IADD:
			a := vm.sPop()
			b := vm.sPop()
			vm.sPush(a + b)
		case PRINT:
			fmt.Println(vm.sPop())
		case STORE:
			addr := vm.nextOp()
			vm.data[addr] = vm.sPop()
		case FETCH:
			addr := vm.nextOp()
			vm.sPush(vm.data[addr])
		case ISLT:
			arg := vm.nextOp()
			val := vm.sPop()
			if val < arg {
				vm.sPush(1)
			} else {
				vm.sPush(0)
			}
		case ISGT:
			arg := vm.nextOp()
			val := vm.sPop()
			if val > arg {
				vm.sPush(1)
			} else {
				vm.sPush(0)
			}
		case JMP:
			vm.pc = vm.nextOp()
		case JMPT:
			addr := vm.nextOp()
			if vm.sPop() == 1 {
				vm.pc = addr
			}
		case JMPF:
			addr := vm.nextOp()
			if vm.sPop() == 0 {
				vm.pc = addr
			}
		case CALL:
			// Get the addr to call
			addr := vm.nextOp()

			// Get the number or args to call it with
			nArgs := vm.nextOp()

			// Store the number of args, frame pointer and program counter
			vm.sPush(nArgs)
			vm.sPush(vm.fp)
			vm.sPush(vm.pc)

			// Set the frame pointer to the top of the stack
			vm.fp = vm.sp

			// Set the program counter to the call addr
			vm.pc = addr
		case ARG:
			pos := vm.fp - vm.nextOp() - 3 // Args are 3 below the frame pointer
			vm.sPush(vm.stack[pos])
		case RET:
			// Get the return val
			res := vm.sPop()

			// Set the stack pointer to the frame pointer
			vm.sp = vm.fp

			// Pop off the previous program counter, frame pointer
			vm.pc = vm.sPop()
			vm.fp = vm.sPop()

			// Pop off the args
			nArgs := vm.sPop()
			for i := 0; i < nArgs; i++ {
				vm.sPop()
			}

			// Push the result onto the stack
			vm.sPush(res)

		case HALT:
			vm.dumpMem()
			return
		default:
			log.Fatal("Unknown Op")
		}
	}
}

// Meat goes here
func main() {

	// Sub calling :)
	code := []int{
		// Add two ints
		/* 0000 */ ARG, 0, // Push arg 0 onto the stack
		/* 0002 */ ARG, 1, // Push arg 1 onto the stack
		/* 0004 */ IADD, // Add them together
		/* 0005 */ RET, // Return the result

		// Add 2 to an int
		/* 0006 */ ARG, 0, // Get arg 0
		/* 0008 */ IPUSH, 2, // Push 2 onto the stack
		/* 0010 */ CALL, 0, 2, // Add them together using the sub
		/* 0013 */ RET, // Return the result

		// Main
		/* 0014 */ IPUSH, 3, // Push 3 onto the stack
		/* 0016 */ CALL, 6, 1, // Add 2 to it by calling addr 6 with 1 arg
		/* 0019 */ PRINT, // Print the result
		/* 0020 */ HALT, // Schtap!
	}
	vm := newVM(code, 14, 4)

	vm.run()
}