Implement setting values at index in an array

7 files changed, 89 insertions, 0 deletions

diff --git a/pkg/lang/vm/code/op.go b/pkg/lang/vm/code/op.go
index b4c172a..8a5c6b5 100644
--- a/pkg/lang/vm/code/op.go
+++ b/pkg/lang/vm/code/op.go
@@ -47,6 +47,7 @@ const (
 	OpGte
 
 	OpIndex
+	OpSetAtIndex
 	OpCall
 
 	OpJmp
diff --git a/pkg/lang/vm/exec.go b/pkg/lang/vm/exec.go
index 34cf444..a64751a 100644
--- a/pkg/lang/vm/exec.go
+++ b/pkg/lang/vm/exec.go
@@ -832,6 +832,48 @@ func (vm *VM) execIndex() error {
 	return nil
 }
 
+func (vm *VM) execSetAtIndex() error {
+	e, err := vm.popAndDrop()
+	if err != nil {
+		return err
+	}
+
+	i, err := vm.popAndDrop()
+	if err != nil {
+		return err
+	}
+
+	v, err := vm.popAndDrop()
+	if err != nil {
+		return err
+	}
+
+	if i.Type() != value.IntType || v.Type() != value.ArrayType {
+		return ErrInvalidOperandTypes{
+			Op: code.OpSetAtIndex,
+			X:  v.Type(),
+			Y:  i.Type(),
+		}
+	}
+
+	arr := v.Data().(value.ArrayData)
+	idx := i.Data().(value.IntData).Get()
+
+	len, err := arr.Len(vm.memory)
+	if err != nil {
+		return err
+	}
+
+	if idx < 0 || idx >= int64(len) {
+		return ErrArrayIndexOutOfBounds{
+			Index: int(idx),
+			Len:   int(len),
+		}
+	}
+
+	return arr.Set(vm.memory, int(idx), e)
+}
+
 func (vm *VM) execCall(argCount uint) error {
 	var err error
 
diff --git a/pkg/lang/vm/text/decompiler.go b/pkg/lang/vm/text/decompiler.go
index 8a85d49..df9e1ed 100644
--- a/pkg/lang/vm/text/decompiler.go
+++ b/pkg/lang/vm/text/decompiler.go
@@ -67,6 +67,7 @@ func (d *Decompiler) decompileInstruction(bc code.Raw) (string, code.Raw) {
 		code.OpLte,
 		code.OpGte,
 		code.OpIndex,
+		code.OpSetAtIndex,
 		code.OpRet,
 		code.OpTempArrLen,
 		code.OpTempArrPush:
diff --git a/pkg/lang/vm/text/op.go b/pkg/lang/vm/text/op.go
index ff3c568..213656b 100644
--- a/pkg/lang/vm/text/op.go
+++ b/pkg/lang/vm/text/op.go
@@ -40,6 +40,7 @@ var (
 		code.OpLte:          "lte",
 		code.OpGte:          "gte",
 		code.OpIndex:        "index",
+		code.OpSetAtIndex:   "set_at_index",
 		code.OpCall:         "call",
 		code.OpJmp:          "jmp",
 		code.OpJt:           "jt",
diff --git a/pkg/lang/vm/value/data.go b/pkg/lang/vm/value/data.go
index 1cd258b..59e4742 100644
--- a/pkg/lang/vm/value/data.go
+++ b/pkg/lang/vm/value/data.go
@@ -108,6 +108,20 @@ func (a ArrayData) At(m mem.Mem, i int) (Value, error) {
 	return arr[i], nil
 }
 
+func (a ArrayData) Set(m mem.Mem, i int, v Value) error {
+	data, err := m.Get(a.data)
+	if err != nil {
+		return err
+	}
+	arr := data.(ArrayCell).Get()
+	arr[i] = v
+
+	if err := m.Set(a.data, ArrayCell(arr)); err != nil {
+		return err
+	}
+	return nil
+}
+
 func (a ArrayData) Push(m mem.Mem, v Value) error {
 	data, err := m.Get(a.data)
 	if err != nil {
diff --git a/pkg/lang/vm/vm.go b/pkg/lang/vm/vm.go
index 4b40633..f24095d 100644
--- a/pkg/lang/vm/vm.go
+++ b/pkg/lang/vm/vm.go
@@ -270,6 +270,9 @@ func (vm *VM) step(module *code.Code, op code.Op) (stepDecision, error) {
 
 	case code.OpIndex:
 		err = vm.execIndex()
+	case code.OpSetAtIndex:
+		err = vm.execSetAtIndex()
+
 	case code.OpCall:
 		argCount, advance := module.GetUint(vm.pc())
 		vm.advancePC(advance)
diff --git a/pkg/lang/vm/vm_test.go b/pkg/lang/vm/vm_test.go
index 178d756..8f961ef 100644
--- a/pkg/lang/vm/vm_test.go
+++ b/pkg/lang/vm/vm_test.go
@@ -510,6 +510,33 @@ func TestCoreSay(t *testing.T) {
 	require.Equal(t, "\"Meow!!\"", string(result))
 }
 
+func TestSetAtArray(t *testing.T) {
+	src := `
+	push_array
+
+	get_local 0
+	get_member "push"
+	push_int 1
+	call 1
+	drop 1
+
+	get_local 0
+	get_member "push"
+	push_int 2
+	call 1
+	drop 1
+
+	get_local 0
+	push_int 0
+	push_int 2
+	set_at_index
+
+	get_local 0
+	`
+
+	test(t, src, "[2, 2]")
+}
+
 func test(t *testing.T, src string, expected string) {
 	bc := compile(t, src)
 	vm := vm.New(modules.NewUnknownModule(&bc))