OpMod and OpSetLocal for prime test

5 files changed, 164 insertions, 1 deletions

diff --git a/pkg/lang/vm/code/op.go b/pkg/lang/vm/code/op.go
index 1e501db..8b8ff3a 100644
--- a/pkg/lang/vm/code/op.go
+++ b/pkg/lang/vm/code/op.go
@@ -21,6 +21,7 @@ const (
 
 	OpGetGlobal
 	OpGetLocal
+	OpSetLocal
 	OpGetMember
 	OpSetMember
 
@@ -32,6 +33,7 @@ const (
 
 	OpAdd
 	OpSub
+	OpMod
 	OpIndex
 	OpLte
 	OpCall
diff --git a/pkg/lang/vm/exec.go b/pkg/lang/vm/exec.go
index be72158..4fbc516 100644
--- a/pkg/lang/vm/exec.go
+++ b/pkg/lang/vm/exec.go
@@ -69,6 +69,23 @@ func (vm *VM) execGetLocal(offset int) error {
 	return nil
 }
 
+func (vm *VM) execSetLocal(offset int) error {
+	new, err := vm.stack.Pop()
+	if err != nil {
+		return err
+	}
+
+	local, err := vm.stack.Local(int(offset))
+	if err != nil {
+		return err
+	}
+
+	local.Drop(vm.memory)
+
+	stackIndex := vm.stack.LocalToStackIndex(offset)
+	return vm.stack.Set(stackIndex, new)
+}
+
 func (vm *VM) execPushType(name string) error {
 	ref, err := value.NewType(vm.memory, name)
 	if err != nil {
@@ -556,6 +573,44 @@ func (vm *VM) execSub() error {
 	return nil
 }
 
+func (vm *VM) execMod() error {
+	x, err := vm.popAndDrop()
+	if err != nil {
+		return err
+	}
+	y, err := vm.popAndDrop()
+	if err != nil {
+		return err
+	}
+
+	var res value.Value
+
+	switch x.Type() {
+	case value.IntType:
+		xv := x.Data().(value.IntData).Get()
+		switch y.Type() {
+		case value.IntType:
+			yv := y.Data().(value.IntData).Get()
+			res = value.NewInt(xv % yv)
+		default:
+			return ErrInvalidOperandTypes{
+				Op: code.OpMod,
+				X:  x.Type(),
+				Y:  y.Type(),
+			}
+		}
+	default:
+		return ErrInvalidOperandTypes{
+			Op: code.OpMod,
+			X:  x.Type(),
+			Y:  y.Type(),
+		}
+	}
+
+	vm.stack.Push(res)
+	return nil
+}
+
 func (vm *VM) execIndex() error {
 	v, err := vm.popAndDrop()
 	if err != nil {
diff --git a/pkg/lang/vm/text/op.go b/pkg/lang/vm/text/op.go
index e6af940..0d01bdb 100644
--- a/pkg/lang/vm/text/op.go
+++ b/pkg/lang/vm/text/op.go
@@ -19,6 +19,7 @@ var (
 		code.OpDrop:         "drop",
 		code.OpGetGlobal:    "get_global",
 		code.OpGetLocal:     "get_local",
+		code.OpSetLocal:     "set_local",
 		code.OpGetMember:    "get_member",
 		code.OpSetMember:    "set_member",
 		code.OpGetEnv:       "get_env",
@@ -27,6 +28,7 @@ var (
 		code.OpAnchorType:   "anchor_type",
 		code.OpAdd:          "add",
 		code.OpSub:          "sub",
+		code.OpMod:          "mod",
 		code.OpIndex:        "index",
 		code.OpLte:          "lte",
 		code.OpCall:         "call",
diff --git a/pkg/lang/vm/vm.go b/pkg/lang/vm/vm.go
index 41334a6..e483176 100644
--- a/pkg/lang/vm/vm.go
+++ b/pkg/lang/vm/vm.go
@@ -125,6 +125,11 @@ func (vm *VM) step(op code.Op) (stepDecision, error) {
 		vm.pc += advance
 
 		err = vm.execGetLocal(int(offset))
+	case code.OpSetLocal:
+		offset, advance := vm.code.GetInt(vm.pc)
+		vm.pc += advance
+
+		err = vm.execSetLocal(int(offset))
 	case code.OpGetMember:
 		name, advance := vm.code.GetString(vm.pc)
 		vm.pc += advance
@@ -159,6 +164,8 @@ func (vm *VM) step(op code.Op) (stepDecision, error) {
 		err = vm.execAdd()
 	case code.OpSub:
 		err = vm.execSub()
+	case code.OpMod:
+		err = vm.execMod()
 	case code.OpIndex:
 		err = vm.execIndex()
 	case code.OpLte:
diff --git a/pkg/lang/vm/vm_test.go b/pkg/lang/vm/vm_test.go
index 325f200..d63523f 100644
--- a/pkg/lang/vm/vm_test.go
+++ b/pkg/lang/vm/vm_test.go
@@ -236,7 +236,7 @@ func TestTypeConstruct(t *testing.T) {
 				this.name = name
 				this.age = age
 			}
-		
+
 			fn meow(this) {
 				return this.name + " says Meow!"
 			}
@@ -306,6 +306,103 @@ func TestTypeConstruct(t *testing.T) {
 	test(t, src, "\"Kitty says Meow!\"")
 }
 
+func TestPrimes(t *testing.T) {
+	/*
+		var i = 0
+		var primes = []
+
+		for i <= 10 {
+			var factors = []
+
+			var j = 1
+			for j <= i {
+				if i % j <= 0 {
+					factors.push(j)
+				}
+				j = j + 1
+			}
+
+			if factors.length() <= 2 {
+				primes.push(i)
+			}
+
+			i = i + 1
+		}
+
+		print(primes)
+	*/
+
+	// TODO: This test uses OpLte instead of OpEq, which works,
+	// but once OpEq is implemented, this test should be updated.
+
+	src := `
+	push_int 2
+	push_array
+	@main_loop:
+		push_int 10
+		get_local 0
+		lte
+		jf @end
+
+		push_array
+		push_int 1
+
+	@factor_loop:
+		get_local 0
+		get_local 3
+		lte
+		jf @prime_check
+
+		push_int 0
+
+		get_local 3
+		get_local 0
+		mod
+
+		lte
+		jf @factor_loop_inc
+
+		get_local 3
+		get_local 2
+		temp_arr_push
+
+	@factor_loop_inc:
+		push_int 1
+		get_local 3
+		add
+		set_local 3
+		jmp @factor_loop
+
+	@prime_check:
+		push_int 2
+		get_local 2
+		get_member "length"
+		call 0
+		lte
+
+		jf @main_loop_inc
+
+		get_local 0
+		get_local 1
+		temp_arr_push
+
+	@main_loop_inc:
+		push_int 1
+		get_local 0
+		add
+		set_local 0
+
+		drop
+		drop
+
+		jmp @main_loop
+	@end:
+		halt
+	`
+
+	test(t, src, "[2, 3, 5, 7]")
+}
+
 func test(t *testing.T, src string, expected string) {
 	bc := compile(t, src)
 	vm := vm.New(&bc)