Harden VM Mem

3 files changed, 115 insertions, 37 deletions

diff --git a/pkg/lang/vm/value/cells.go b/pkg/lang/vm/value/cells.go
new file mode 100644
index 0000000..17a0916
--- /dev/null
+++ b/pkg/lang/vm/value/cells.go
@@ -0,0 +1,46 @@
+package value
+
+import "jinx/pkg/lang/vm/mem"
+
+type ArrayCell []Value
+
+func (a ArrayCell) DropCell(m *mem.Mem) {
+	for _, v := range a {
+		v.Drop(m)
+	}
+}
+
+func (a ArrayCell) Get() []Value {
+	return a
+}
+
+type StringCell string
+
+func (s StringCell) DropCell(m *mem.Mem) {
+}
+
+func (s StringCell) Get() string {
+	return string(s)
+}
+
+type OutletCell Value
+
+func (o OutletCell) DropCell(m *mem.Mem) {
+	Value(o).Drop(m)
+}
+
+func (o OutletCell) Get() Value {
+	return Value(o)
+}
+
+type EnvCell Env
+
+func (e EnvCell) DropCell(m *mem.Mem) {
+	for _, v := range e.references {
+		m.Release(v.outlet)
+	}
+}
+
+func (e EnvCell) Get() Env {
+	return Env(e)
+}
diff --git a/pkg/lang/vm/value/data.go b/pkg/lang/vm/value/data.go
index a49753e..2e3b3e6 100644
--- a/pkg/lang/vm/value/data.go
+++ b/pkg/lang/vm/value/data.go
@@ -8,7 +8,7 @@ import (
 )
 
 type Data interface {
-	String(*mem.Mem) string
+	String(*mem.Mem) (string, error)
 }
 
 type IntData int64
@@ -17,8 +17,8 @@ func (i IntData) Get() int64 {
 	return int64(i)
 }
 
-func (i IntData) String(_ *mem.Mem) string {
-	return strconv.FormatInt(int64(i), 10)
+func (i IntData) String(_ *mem.Mem) (string, error) {
+	return strconv.FormatInt(int64(i), 10), nil
 }
 
 type FloatData float64
@@ -27,17 +27,20 @@ func (f FloatData) Get() float64 {
 	return float64(f)
 }
 
-func (f FloatData) String(_ *mem.Mem) string {
-	return strconv.FormatFloat(float64(f), 'f', -1, 64)
+func (f FloatData) String(_ *mem.Mem) (string, error) {
+	return strconv.FormatFloat(float64(f), 'f', -1, 64), nil
 }
 
 type StringData struct {
 	data mem.Ptr
 }
 
-func (s StringData) String(m *mem.Mem) string {
-	data := m.Get(s.data)
-	return "\"" + data.(string) + "\""
+func (s StringData) String(m *mem.Mem) (string, error) {
+	if data, err := m.Get(s.data); err == nil {
+		return "\"" + data.(StringCell).Get() + "\"", nil
+	} else {
+		return "", err
+	}
 }
 
 type BoolData bool
@@ -46,16 +49,21 @@ func (b BoolData) Get() bool {
 	return bool(b)
 }
 
-func (b BoolData) String(_ *mem.Mem) string {
-	return strconv.FormatBool(bool(b))
+func (b BoolData) String(_ *mem.Mem) (string, error) {
+	return strconv.FormatBool(bool(b)), nil
 }
 
 type ArrayData struct {
 	data mem.Ptr
 }
 
-func (a ArrayData) String(m *mem.Mem) string {
-	arr := m.Get(a.data).([]Value)
+func (a ArrayData) String(m *mem.Mem) (string, error) {
+	val, err := m.Get(a.data)
+	if err != nil {
+		return "", err
+	}
+
+	arr := val.(ArrayCell).Get()
 
 	builder := strings.Builder{}
 	builder.WriteString("[")
@@ -63,29 +71,45 @@ func (a ArrayData) String(m *mem.Mem) string {
 		if i > 0 {
 			builder.WriteString(", ")
 		}
-		builder.WriteString(v.Data().String(m))
+		if s, err := v.Data().String(m); err == nil {
+			builder.WriteString(s)
+		} else {
+			return "", err
+		}
 	}
 	builder.WriteString("]")
-	return builder.String()
+	return builder.String(), nil
 }
 
-func (a ArrayData) Len(m *mem.Mem) int {
-	data := m.Get(a.data)
-	arr := data.([]Value)
-	return len(arr)
+func (a ArrayData) Len(m *mem.Mem) (int, error) {
+	data, err := m.Get(a.data)
+	if err != nil {
+		return 0, err
+	}
+	arr := data.(ArrayCell).Get()
+	return len(arr), nil
 }
 
-func (a ArrayData) At(m *mem.Mem, i int) Value {
-	data := m.Get(a.data)
-	arr := data.([]Value)
-	return arr[i]
+func (a ArrayData) At(m *mem.Mem, i int) (Value, error) {
+	data, err := m.Get(a.data)
+	if err != nil {
+		return Value{}, err
+	}
+	arr := data.(ArrayCell).Get()
+	return arr[i], nil
 }
 
-func (a ArrayData) Push(m *mem.Mem, v Value) {
-	data := m.Get(a.data)
-	arr := data.([]Value)
+func (a ArrayData) Push(m *mem.Mem, v Value) error {
+	data, err := m.Get(a.data)
+	if err != nil {
+		return err
+	}
+
+	arr := data.(ArrayCell).Get()
 	arr = append(arr, v)
-	m.Set(a.data, arr)
+	m.Set(a.data, ArrayCell(arr))
+
+	return nil
 }
 
 type NullData struct{}
@@ -107,8 +131,8 @@ func (f FunctionData) WithEnv(env mem.Ptr) FunctionData {
 	return FunctionData{pc: f.pc, env: env}
 }
 
-func (f FunctionData) String(_ *mem.Mem) string {
-	return fmt.Sprintf("<fn %d>", f.pc)
+func (f FunctionData) String(_ *mem.Mem) (string, error) {
+	return fmt.Sprintf("<fn %d>", f.pc), nil
 }
 
 type ObjectData struct{} // TODO
diff --git a/pkg/lang/vm/value/value.go b/pkg/lang/vm/value/value.go
index d19bbd6..f98740e 100644
--- a/pkg/lang/vm/value/value.go
+++ b/pkg/lang/vm/value/value.go
@@ -20,13 +20,17 @@ func NewFloat(x float64) Value {
 	return Value{t: t, d: FloatData(x)}
 }
 
-func NewString(m *mem.Mem, str string) Value {
+func NewString(m *mem.Mem, str string) (Value, error) {
 	t := Type{Kind: StringType}
 
-	ptr := m.Allocate(mem.CellKindString)
-	m.Set(ptr, str)
+	ptr, err := m.Allocate(mem.CellKindString)
+	if err != nil {
+		return Value{}, err
+	}
+
+	m.Set(ptr, StringCell(str))
 
-	return Value{t: t, d: StringData{data: ptr}}
+	return Value{t: t, d: StringData{data: ptr}}, nil
 }
 
 func NewBool(b bool) Value {
@@ -34,13 +38,17 @@ func NewBool(b bool) Value {
 	return Value{t: t, d: BoolData(b)}
 }
 
-func NewArray(m *mem.Mem, arr []Value) Value {
+func NewArray(m *mem.Mem, arr []Value) (Value, error) {
 	t := Type{Kind: ArrayType}
 
-	ptr := m.Allocate(mem.CellKindArray)
-	m.Set(ptr, arr)
+	ptr, err := m.Allocate(mem.CellKindArray)
+	if err != nil {
+		return Value{}, err
+	}
+
+	m.Set(ptr, ArrayCell(arr))
 
-	return Value{t: t, d: ArrayData{data: ptr}}
+	return Value{t: t, d: ArrayData{data: ptr}}, nil
 }
 
 func NewNull() Value {
@@ -99,7 +107,7 @@ func (v Value) Clone(m *mem.Mem) Value {
 func (v Value) Drop(m *mem.Mem) {
 	// If value has an outlet, don't drop it and instead move it to the outlet.
 	if !v.outlet.IsNull() {
-		m.Set(v.outlet, v)
+		m.Set(v.outlet, OutletCell(v))
 		return
 	}