From e99387beb4609a923e92281e170d1b4722aa3238 Mon Sep 17 00:00:00 2001
From: Ian Gulliver <git@flamingcow.io>
Date: Thu, 18 Nov 2021 17:53:38 -1000
Subject: [PATCH] Full test (which fails)

---
 assembler/assemble_test.go |  4 ++++
 vm/instruction.go          |  7 -------
 vm/memory.go               | 14 +++++++-------
 vm/stackframe.go           |  4 ++--
 vm/state.go                | 20 ++++++++++++++------
 5 files changed, 27 insertions(+), 22 deletions(-)

diff --git a/assembler/assemble_test.go b/assembler/assemble_test.go
index 30570f7..7fda311 100644
--- a/assembler/assemble_test.go
+++ b/assembler/assemble_test.go
@@ -27,4 +27,8 @@ functions:
 	if err != nil {
 		t.Fatal(err)
 	}
+
+	if state.GlobalMemory().MustReadUnsigned(0) != 3 {
+		t.Fatalf("Expected 3, found %d\n", state.GlobalMemory().MustReadUnsigned(0))
+	}
 }
diff --git a/vm/instruction.go b/vm/instruction.go
index fffddba..3b21965 100644
--- a/vm/instruction.go
+++ b/vm/instruction.go
@@ -41,12 +41,5 @@ func NewInstructionsFromByteCode(byteCode []byte) ([]*Instruction, error) {
 		instrs = append(instrs, instr)
 	}
 
-	if len(instrs) == 0 || instrs[len(instrs)-1].OpCode != OpReturn {
-		// Add implicit return at the end of each function
-		instrs = append(instrs, &Instruction{
-			OpCode: OpReturn,
-		})
-	}
-
 	return instrs, nil
 }
diff --git a/vm/memory.go b/vm/memory.go
index bd38168..70d8a74 100644
--- a/vm/memory.go
+++ b/vm/memory.go
@@ -2,17 +2,17 @@ package vm
 
 import "fmt"
 
-type memory struct {
+type Memory struct {
 	entries []uint64
 }
 
-func newMemory(size uint64) *memory {
-	return &memory{
+func NewMemory(size uint64) *Memory {
+	return &Memory{
 		entries: make([]uint64, size),
 	}
 }
 
-func (mem *memory) readUnsigned(index uint64) (uint64, error) {
+func (mem *Memory) ReadUnsigned(index uint64) (uint64, error) {
 	if index >= uint64(len(mem.entries)) {
 		return 0, fmt.Errorf("Invalid memory index: %016x", index)
 	}
@@ -20,15 +20,15 @@ func (mem *memory) readUnsigned(index uint64) (uint64, error) {
 	return mem.entries[index], nil
 }
 
-func (mem *memory) mustReadUnsigned(index uint64) uint64 {
-	value, err := mem.readUnsigned(index)
+func (mem *Memory) MustReadUnsigned(index uint64) uint64 {
+	value, err := mem.ReadUnsigned(index)
 	if err != nil {
 		panic(err)
 	}
 	return value
 }
 
-func (mem *memory) writeUnsigned(index uint64, value uint64) error {
+func (mem *Memory) WriteUnsigned(index uint64, value uint64) error {
 	if index >= uint64(len(mem.entries)) {
 		return fmt.Errorf("Invalid memory index: %016x", index)
 	}
diff --git a/vm/stackframe.go b/vm/stackframe.go
index b91fe73..bd93c2e 100644
--- a/vm/stackframe.go
+++ b/vm/stackframe.go
@@ -3,13 +3,13 @@ package vm
 type stackFrame struct {
 	previousFunctionIndex    int64
 	previousInstructionIndex int64
-	functionMemory           *memory
+	functionMemory           *Memory
 }
 
 func newStackFrame(state *State) *stackFrame {
 	return &stackFrame{
 		previousFunctionIndex:    state.functionIndex,
 		previousInstructionIndex: state.instructionIndex,
-		functionMemory:           newMemory(16),
+		functionMemory:           NewMemory(16),
 	}
 }
diff --git a/vm/state.go b/vm/state.go
index f615751..eea02cf 100644
--- a/vm/state.go
+++ b/vm/state.go
@@ -13,14 +13,14 @@ type State struct {
 	instructionIndex int64
 
 	comparisonResult bool
-	globalMemory     *memory
+	globalMemory     *Memory
 	stack            []*stackFrame
 }
 
 func NewState(functions [][]*Instruction) (*State, error) {
 	return &State{
 		functions:    functions,
-		globalMemory: newMemory(16),
+		globalMemory: NewMemory(16),
 	}, nil
 }
 
@@ -51,6 +51,10 @@ func (state *State) Execute() error {
 	return state.err
 }
 
+func (state *State) GlobalMemory() *Memory {
+	return state.globalMemory
+}
+
 func (state *State) stackFrame() *stackFrame {
 	return state.stack[len(state.stack)-1]
 }
@@ -83,6 +87,10 @@ func (state *State) processInstruction() {
 	instr := fnc[state.instructionIndex]
 	state.instructionIndex += 1
 	instr.opHandler(state, instr)
+
+	if state.instructionIndex >= int64(len(fnc)) {
+		state.ret()
+	}
 }
 
 func (state *State) readUnsigned(op *Operand) uint64 {
@@ -91,14 +99,14 @@ func (state *State) readUnsigned(op *Operand) uint64 {
 		return op.Value
 
 	case FunctionMemoryIndex:
-		value, err := state.stackFrame().functionMemory.readUnsigned(op.Value)
+		value, err := state.stackFrame().functionMemory.ReadUnsigned(op.Value)
 		if err != nil {
 			state.setError(err)
 		}
 		return value
 
 	case GlobalMemoryIndex:
-		value, err := state.globalMemory.readUnsigned(op.Value)
+		value, err := state.globalMemory.ReadUnsigned(op.Value)
 		if err != nil {
 			state.setError(err)
 		}
@@ -120,13 +128,13 @@ func (state *State) writeUnsigned(op *Operand, value uint64) {
 		state.setError(fmt.Errorf("Write to literal operand"))
 
 	case FunctionMemoryIndex:
-		err := state.stackFrame().functionMemory.writeUnsigned(op.Value, value)
+		err := state.stackFrame().functionMemory.WriteUnsigned(op.Value, value)
 		if err != nil {
 			state.setError(err)
 		}
 
 	case GlobalMemoryIndex:
-		err := state.globalMemory.writeUnsigned(op.Value, value)
+		err := state.globalMemory.WriteUnsigned(op.Value, value)
 		if err != nil {
 			state.setError(err)
 		}