8 files changed, 229 insertions, 60 deletions

diff --git a/pkg/lang/vm/exec.go b/pkg/lang/vm/exec.go
index 655f7f1..90b79a9 100644
--- a/pkg/lang/vm/exec.go
+++ b/pkg/lang/vm/exec.go
@@ -14,7 +14,7 @@ func (vm *VM) execPushFloat(x float64) {
 }
 
 func (vm *VM) execPushString(str string) {
-	vm.stack.Push(value.NewString(str))
+	vm.stack.Push(value.NewString(&vm.memory, str))
 }
 
 func (vm *VM) execPushBool(b bool) {
@@ -26,7 +26,7 @@ func (vm *VM) execPushNull() {
 }
 
 func (vm *VM) execPushArray() {
-	vm.stack.Push(value.NewArray([]value.Value{}))
+	vm.stack.Push(value.NewArray(&vm.memory, []value.Value{}))
 }
 
 func (vm *VM) execPushFunction(pc int) {
@@ -39,16 +39,18 @@ func (vm *VM) execGetLocal(offset int) error {
 		return err
 	}
 
-	vm.stack.Push(local)
+	v := local.Clone(&vm.memory)
+
+	vm.stack.Push(v)
 	return nil
 }
 
 func (vm *VM) execAdd() error {
-	x, err := vm.stack.Pop()
+	x, err := vm.popAndDrop()
 	if err != nil {
 		return err
 	}
-	y, err := vm.stack.Pop()
+	y, err := vm.popAndDrop()
 	if err != nil {
 		return err
 	}
@@ -91,7 +93,8 @@ func (vm *VM) execAdd() error {
 	case value.StringType:
 		switch y.Type().Kind {
 		case value.StringType:
-			res = value.NewString(x.Data().(value.StringData).Get() + y.Data().(value.StringData).Get())
+			panic("not implemented")
+			// res = value.NewString(x.Data().(value.StringData).Get() + y.Data().(value.StringData).Get())
 		default:
 			return ErrInvalidOperandTypes{
 				Op: code.OpAdd,
@@ -112,11 +115,11 @@ func (vm *VM) execAdd() error {
 }
 
 func (vm *VM) execSub() error {
-	x, err := vm.stack.Pop()
+	x, err := vm.popAndDrop()
 	if err != nil {
 		return err
 	}
-	y, err := vm.stack.Pop()
+	y, err := vm.popAndDrop()
 	if err != nil {
 		return err
 	}
@@ -169,11 +172,11 @@ func (vm *VM) execSub() error {
 }
 
 func (vm *VM) execIndex() error {
-	v, err := vm.stack.Pop()
+	v, err := vm.popAndDrop()
 	if err != nil {
 		return err
 	}
-	i, err := vm.stack.Pop()
+	i, err := vm.popAndDrop()
 	if err != nil {
 		return err
 	}
@@ -184,13 +187,16 @@ func (vm *VM) execIndex() error {
 		switch i.Type().Kind {
 		case value.IntType:
 			idx := i.Data().(value.IntData).Get()
-			if idx < 0 || idx >= int64(arr.Len()) {
+			len := int64(arr.Len(&vm.memory))
+			if idx < 0 || idx >= len {
 				return ErrArrayIndexOutOfBounds{
 					Index: int(idx),
-					Len:   arr.Len(),
+					Len:   int(len),
 				}
 			}
-			vm.stack.Push(arr.At(int(idx)))
+
+			val := arr.At(&vm.memory, int(idx)).Clone(&vm.memory)
+			vm.stack.Push(val)
 		default:
 			return ErrInvalidOperandTypes{
 				Op: code.OpIndex,
@@ -210,11 +216,11 @@ func (vm *VM) execIndex() error {
 }
 
 func (vm *VM) execLte() error {
-	x, err := vm.stack.Pop()
+	x, err := vm.popAndDrop()
 	if err != nil {
 		return err
 	}
-	y, err := vm.stack.Pop()
+	y, err := vm.popAndDrop()
 	if err != nil {
 		return err
 	}
@@ -267,7 +273,7 @@ func (vm *VM) execLte() error {
 }
 
 func (vm *VM) execCall() error {
-	f, err := vm.stack.Pop()
+	f, err := vm.popAndDrop()
 	if err != nil {
 		return err
 	}
@@ -291,7 +297,7 @@ func (vm *VM) execCall() error {
 }
 
 func (vm *VM) execJumpIf(pc int, cond bool) error {
-	b, err := vm.stack.Pop()
+	b, err := vm.popAndDrop()
 	if err != nil {
 		return err
 	}
@@ -325,7 +331,7 @@ func (vm *VM) execRet() error {
 		return err
 	}
 
-	pc, err := vm.stack.PopCall()
+	pc, err := vm.popCallAndDrop()
 	if err != nil {
 		return err
 	}
@@ -336,7 +342,7 @@ func (vm *VM) execRet() error {
 }
 
 func (vm *VM) execTempArrLen() error {
-	a, err := vm.stack.Pop()
+	a, err := vm.popAndDrop()
 	if err != nil {
 		return err
 	}
@@ -344,7 +350,8 @@ func (vm *VM) execTempArrLen() error {
 	switch a.Type().Kind {
 	case value.ArrayType:
 		arr := a.Data().(value.ArrayData)
-		res := value.NewInt(int64(arr.Len()))
+		len := int64(arr.Len(&vm.memory))
+		res := value.NewInt(len)
 		vm.stack.Push(res)
 	default:
 		return ErrInvalidOperandTypes{
@@ -357,11 +364,11 @@ func (vm *VM) execTempArrLen() error {
 }
 
 func (vm *VM) execTempArrPush() error {
-	a, err := vm.stack.Pop()
+	a, err := vm.popAndDrop()
 	if err != nil {
 		return err
 	}
-	e, err := vm.stack.Pop()
+	e, err := vm.popAndDrop()
 	if err != nil {
 		return err
 	}
@@ -369,7 +376,7 @@ func (vm *VM) execTempArrPush() error {
 	switch a.Type().Kind {
 	case value.ArrayType:
 		arr := a.Data().(value.ArrayData)
-		arr.Push(e)
+		arr.Push(&vm.memory, e)
 	default:
 		return ErrInvalidOperandType{
 			Op: code.OpTempArrPush,
diff --git a/pkg/lang/vm/mem/cell.go b/pkg/lang/vm/mem/cell.go
new file mode 100644
index 0000000..044a45c
--- /dev/null
+++ b/pkg/lang/vm/mem/cell.go
@@ -0,0 +1,16 @@
+package mem
+
+type CellKind int
+
+const (
+	CellKindEmpty CellKind = iota
+	CellKindString
+	CellKindArray
+	CellKindEscaped
+)
+
+type cell struct {
+	kind CellKind
+	refs int
+	data any
+}
diff --git a/pkg/lang/vm/mem/mem.go b/pkg/lang/vm/mem/mem.go
new file mode 100644
index 0000000..4335347
--- /dev/null
+++ b/pkg/lang/vm/mem/mem.go
@@ -0,0 +1,68 @@
+package mem
+
+type Mem struct {
+	cells []cell
+	free  []int
+}
+
+func New() Mem {
+	return Mem{
+		cells: make([]cell, 0),
+		free:  make([]int, 0),
+	}
+}
+
+func (m *Mem) Allocate(kind CellKind) Ptr {
+	if len(m.free) > 0 {
+		idx := m.free[len(m.free)-1]
+		m.free = m.free[:len(m.free)-1]
+
+		if m.cells[idx].kind != CellKindEmpty {
+			panic("invalid free cell")
+		}
+
+		m.cells[idx].kind = kind
+		return Ptr(idx)
+	} else {
+		idx := len(m.cells)
+		m.cells = append(m.cells, cell{kind: kind})
+		return Ptr(idx)
+	}
+
+}
+
+func (m *Mem) Set(ptr Ptr, v any) {
+	if ptr >= Ptr(len(m.cells)) {
+		panic("out of bounds")
+	}
+
+	m.cells[ptr].data = v
+}
+
+func (m *Mem) Get(ptr Ptr) any {
+	if ptr >= Ptr(len(m.cells)) {
+		panic("out of bounds")
+	}
+
+	return m.cells[ptr].data
+}
+
+func (m *Mem) Retain(ptr Ptr) {
+	if ptr >= Ptr(len(m.cells)) {
+		panic("out of bounds")
+	}
+
+	m.cells[ptr].refs++
+}
+
+func (m *Mem) Release(ptr Ptr) {
+	if ptr >= Ptr(len(m.cells)) {
+		panic("out of bounds")
+	}
+
+	m.cells[ptr].refs--
+	if m.cells[ptr].refs == 0 {
+		m.cells[ptr].kind = CellKindEmpty
+		m.free = append(m.free, int(ptr))
+	}
+}
diff --git a/pkg/lang/vm/mem/ptr.go b/pkg/lang/vm/mem/ptr.go
new file mode 100644
index 0000000..5e3bf3d
--- /dev/null
+++ b/pkg/lang/vm/mem/ptr.go
@@ -0,0 +1,10 @@
+package mem
+
+import "strconv"
+
+type Ptr uint64
+
+func (p Ptr) String() string {
+	val := strconv.FormatUint(uint64(p), 10)
+	return "@" + val
+}
diff --git a/pkg/lang/vm/value/data.go b/pkg/lang/vm/value/data.go
index 5d1398d..fecf8cf 100644
--- a/pkg/lang/vm/value/data.go
+++ b/pkg/lang/vm/value/data.go
@@ -2,11 +2,14 @@ package value
 
 import (
 	"fmt"
+	"jinx/pkg/lang/vm/mem"
 	"strconv"
 	"strings"
 )
 
-type Data fmt.Stringer
+type Data interface {
+	String(*mem.Mem) string
+}
 
 type IntData int64
 
@@ -14,7 +17,7 @@ func (i IntData) Get() int64 {
 	return int64(i)
 }
 
-func (i IntData) String() string {
+func (i IntData) String(_ *mem.Mem) string {
 	return strconv.FormatInt(int64(i), 10)
 }
 
@@ -24,18 +27,17 @@ func (f FloatData) Get() float64 {
 	return float64(f)
 }
 
-func (f FloatData) String() string {
+func (f FloatData) String(_ *mem.Mem) string {
 	return strconv.FormatFloat(float64(f), 'f', -1, 64)
 }
 
-type StringData string
-
-func (s StringData) Get() string {
-	return string(s)
+type StringData struct {
+	data mem.Ptr
 }
 
-func (s StringData) String() string {
-	return "\"" + string(s) + "\""
+func (s StringData) String(m *mem.Mem) string {
+	data := m.Get(s.data)
+	return "\"" + data.(string) + "\""
 }
 
 type BoolData bool
@@ -44,41 +46,46 @@ func (b BoolData) Get() bool {
 	return bool(b)
 }
 
-func (b BoolData) String() string {
+func (b BoolData) String(_ *mem.Mem) string {
 	return strconv.FormatBool(bool(b))
 }
 
 type ArrayData struct {
-	arr *[]Value
+	data mem.Ptr
 }
 
-func (a ArrayData) Get() *[]Value {
-	return a.arr
-}
+func (a ArrayData) String(m *mem.Mem) string {
+	arr := m.Get(a.data).([]Value)
 
-func (a ArrayData) String() string {
 	builder := strings.Builder{}
 	builder.WriteString("[")
-	for i, v := range *a.arr {
+	for i, v := range arr {
 		if i > 0 {
 			builder.WriteString(", ")
 		}
-		builder.WriteString(v.Data().String())
+		builder.WriteString(v.Data().String(m))
 	}
 	builder.WriteString("]")
 	return builder.String()
 }
 
-func (a ArrayData) Len() int {
-	return len(*a.arr)
+func (a ArrayData) Len(m *mem.Mem) int {
+	data := m.Get(a.data)
+	arr := data.([]Value)
+	return len(arr)
 }
 
-func (a ArrayData) At(i int) Value {
-	return (*a.arr)[i]
+func (a ArrayData) At(m *mem.Mem, i int) Value {
+	data := m.Get(a.data)
+	arr := data.([]Value)
+	return arr[i]
 }
 
-func (a ArrayData) Push(v Value) {
-	*a.arr = append(*a.arr, v)
+func (a ArrayData) Push(m *mem.Mem, v Value) {
+	data := m.Get(a.data)
+	arr := data.([]Value)
+	arr = append(arr, v)
+	m.Set(a.data, arr)
 }
 
 type NullData struct{}
@@ -91,7 +98,7 @@ func (f FunctionData) Pc() int {
 	return f.pc
 }
 
-func (f FunctionData) String() string {
+func (f FunctionData) String(_ *mem.Mem) string {
 	return fmt.Sprintf("<fn %d>", f.pc)
 }
 
diff --git a/pkg/lang/vm/value/value.go b/pkg/lang/vm/value/value.go
index 93989cf..06306df 100644
--- a/pkg/lang/vm/value/value.go
+++ b/pkg/lang/vm/value/value.go
@@ -1,5 +1,9 @@
 package value
 
+import (
+	"jinx/pkg/lang/vm/mem"
+)
+
 type Value struct {
 	t Type
 	d Data
@@ -15,9 +19,13 @@ func NewFloat(x float64) Value {
 	return Value{t: t, d: FloatData(x)}
 }
 
-func NewString(str string) Value {
+func NewString(m *mem.Mem, str string) Value {
 	t := Type{Kind: StringType}
-	return Value{t: t, d: StringData(str)}
+
+	ptr := m.Allocate(mem.CellKindString)
+	m.Set(ptr, str)
+
+	return Value{t: t, d: StringData{data: ptr}}
 }
 
 func NewBool(b bool) Value {
@@ -25,9 +33,13 @@ func NewBool(b bool) Value {
 	return Value{t: t, d: BoolData(b)}
 }
 
-func NewArray(arr []Value) Value {
+func NewArray(m *mem.Mem, arr []Value) Value {
 	t := Type{Kind: ArrayType}
-	return Value{t: t, d: ArrayData{arr: &arr}}
+
+	ptr := m.Allocate(mem.CellKindArray)
+	m.Set(ptr, arr)
+
+	return Value{t: t, d: ArrayData{data: ptr}}
 }
 
 func NewNull() Value {
@@ -51,3 +63,29 @@ func (v Value) Type() Type {
 func (v Value) Data() Data {
 	return v.d
 }
+
+func (v Value) Clone(m *mem.Mem) Value {
+	if v.t.Kind == StringType {
+		str := v.d.(StringData)
+		m.Retain(str.data)
+	}
+
+	if v.t.Kind == ArrayType {
+		arr := v.d.(ArrayData)
+		m.Retain(arr.data)
+	}
+
+	return v
+}
+
+func (v Value) Drop(m *mem.Mem) {
+	if v.t.Kind == StringType {
+		str := v.d.(StringData)
+		m.Release(str.data)
+	}
+
+	if v.t.Kind == ArrayType {
+		arr := v.d.(ArrayData)
+		m.Release(arr.data)
+	}
+}
diff --git a/pkg/lang/vm/vm.go b/pkg/lang/vm/vm.go
index dfe1ae2..d896d88 100644
--- a/pkg/lang/vm/vm.go
+++ b/pkg/lang/vm/vm.go
@@ -1,21 +1,24 @@
 package vm
 
 import (
-	"fmt"
 	"jinx/pkg/lang/vm/code"
+	"jinx/pkg/lang/vm/mem"
+	"jinx/pkg/lang/vm/value"
 )
 
 type VM struct {
-	code  *code.Code
-	pc    int
-	stack Stack
+	code   *code.Code
+	pc     int
+	stack  Stack
+	memory mem.Mem
 }
 
 func New(code *code.Code) *VM {
 	return &VM{
-		code:  code,
-		pc:    0,
-		stack: NewStack(),
+		code:   code,
+		pc:     0,
+		stack:  NewStack(),
+		memory: mem.New(),
 	}
 }
 
@@ -25,7 +28,8 @@ func (vm *VM) GetResult() (string, error) {
 		return "", err
 	}
 
-	return fmt.Sprintf("%v", res.Data()), nil
+	str := res.Data().String(&vm.memory)
+	return str, nil
 }
 
 func (vm *VM) Run() error {
@@ -150,3 +154,23 @@ func (vm *VM) step(op code.Op) (stepDecision, error) {
 
 	return stepDecisionContinue, err
 }
+
+func (vm *VM) popAndDrop() (value.Value, error) {
+	v, err := vm.stack.Pop()
+	if err != nil {
+		return value.Value{}, err
+	}
+	v.Drop(&vm.memory)
+	return v, nil
+}
+
+func (vm *VM) popCallAndDrop() (int, error) {
+	for !vm.stack.ReachedBaseOfCall() {
+		_, err := vm.popAndDrop()
+		if err != nil {
+			return 0, err
+		}
+	}
+
+	return vm.stack.PopCall()
+}
diff --git a/pkg/lang/vm/vm_test.go b/pkg/lang/vm/vm_test.go
index 298cfb2..74e64d0 100644
--- a/pkg/lang/vm/vm_test.go
+++ b/pkg/lang/vm/vm_test.go
@@ -108,7 +108,6 @@ func TestFunction(t *testing.T) {
 	test(t, src, "42")
 }
 
-
 func test(t *testing.T, src string, expected string) {
 	bc := compile(t, src)
 	vm := vm.New(&bc)