5 files changed, 300 insertions, 8 deletions

diff --git a/pkg/lang/vm/mem/cell.go b/pkg/lang/vm/mem/cell.go
index 830d382..562d4da 100644
--- a/pkg/lang/vm/mem/cell.go
+++ b/pkg/lang/vm/mem/cell.go
@@ -40,4 +40,5 @@ type cell struct {
 
 type CellData interface {
 	DropCell(Mem)
+	MatchingCellKind() CellKind
 }
diff --git a/pkg/lang/vm/mem/errors.go b/pkg/lang/vm/mem/errors.go
index e7a3f8f..f5dcc86 100644
--- a/pkg/lang/vm/mem/errors.go
+++ b/pkg/lang/vm/mem/errors.go
@@ -1,17 +1,36 @@
 package mem
 
-import "errors"
+import (
+	"errors"
+	"fmt"
+)
 
 var (
 	ErrMemOverflow = errors.New("memory overflow, cannot allocate more than 10000 memory cells")
-
-	ErrFatalNonFreeCell = errors.New("non-free cell marked as free")
 )
 
+type ErrInvalidCellKind struct {
+	Kind CellKind
+}
+
+func (e ErrInvalidCellKind) Error() string {
+	return fmt.Sprintf("cannot allocate cell of kind %s", e.Kind)
+}
+
 type ErrInvalidMemAccess struct {
 	Ptr Ptr
 }
 
 func (e ErrInvalidMemAccess) Error() string {
-	return "invalid memory access at " + e.Ptr.String()
+	return fmt.Sprintf("invalid memory access at %d", e.Ptr)
+}
+
+type ErrDifferingCellKind struct {
+	Ptr      Ptr
+	Expected CellKind
+	Got      CellKind
+}
+
+func (e ErrDifferingCellKind) Error() string {
+	return fmt.Sprintf("tried assigning %v to mem cell of type %v at %d", e.Got, e.Expected, e.Ptr)
 }
diff --git a/pkg/lang/vm/mem/mem.go b/pkg/lang/vm/mem/mem.go
index a645ee2..0f15743 100644
--- a/pkg/lang/vm/mem/mem.go
+++ b/pkg/lang/vm/mem/mem.go
@@ -10,6 +10,7 @@ type Mem interface {
 
 	Retain(ptr Ptr) error
 	Release(ptr Ptr) error
+	RefCount(ptr Ptr) int
 }
 
 type memImpl struct {
@@ -30,15 +31,20 @@ func New() Mem {
 }
 
 func (m *memImpl) Allocate(kind CellKind) (Ptr, error) {
+	if kind == CellKindForbidden || kind == CellKindEmpty {
+		return NullPtr, ErrInvalidCellKind{kind}
+	}
+
 	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 {
-			return NullPtr, ErrFatalNonFreeCell
+			// This should never happen.
+			panic("cell marked as free was not empty")
 		}
 
-		m.cells[idx].kind = kind
+		m.cells[idx] = cell{kind: kind, refs: 1}
 		return Ptr(idx), nil
 	} else {
 		if len(m.cells) > 10000 {
@@ -57,6 +63,10 @@ func (m *memImpl) Set(ptr Ptr, v CellData) error {
 		return err
 	}
 
+	if m.cells[ptr].kind != v.MatchingCellKind() {
+		return ErrDifferingCellKind{Ptr: ptr, Expected: m.cells[ptr].kind, Got: v.MatchingCellKind()}
+	}
+
 	m.cells[ptr].data = v
 	return nil
 }
@@ -71,7 +81,7 @@ func (m *memImpl) Get(ptr Ptr) (CellData, error) {
 
 func (m *memImpl) Is(ptr Ptr, kind CellKind) bool {
 	if ptr >= Ptr(len(m.cells)) {
-		return false
+		return kind == CellKindEmpty
 	}
 
 	return m.cells[ptr].kind == kind
@@ -79,7 +89,7 @@ func (m *memImpl) Is(ptr Ptr, kind CellKind) bool {
 
 func (m *memImpl) Kind(ptr Ptr) CellKind {
 	if ptr >= Ptr(len(m.cells)) {
-		return CellKindForbidden
+		return CellKindEmpty
 	}
 
 	return m.cells[ptr].kind
@@ -111,6 +121,14 @@ func (m *memImpl) Release(ptr Ptr) error {
 	return nil
 }
 
+func (m *memImpl) RefCount(ptr Ptr) int {
+	if err := m.validPtr(ptr); err != nil {
+		return 0
+	}
+
+	return m.cells[ptr].refs
+}
+
 func (m *memImpl) validPtr(ptr Ptr) error {
 	if ptr >= Ptr(len(m.cells)) {
 		return ErrInvalidMemAccess{ptr}
diff --git a/pkg/lang/vm/mem/mem_test.go b/pkg/lang/vm/mem/mem_test.go
new file mode 100644
index 0000000..a75753c
--- /dev/null
+++ b/pkg/lang/vm/mem/mem_test.go
@@ -0,0 +1,238 @@
+package mem_test
+
+import (
+	"jinx/pkg/lang/vm/mem"
+	"jinx/pkg/lang/vm/value"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestAllocation(t *testing.T) {
+	m := mem.New()
+
+	one, err := m.Allocate(mem.CellKindString)
+	assert.NoError(t, err)
+
+	two, err := m.Allocate(mem.CellKindArray)
+	assert.NoError(t, err)
+
+	assert.Equal(t, one, mem.Ptr(1))
+	assert.Equal(t, two, mem.Ptr(2))
+}
+
+func TestGetSet(t *testing.T) {
+	m := mem.New()
+
+	one, err := m.Allocate(mem.CellKindString)
+	assert.NoError(t, err)
+	assert.NoError(t, m.Set(one, value.StringCell("one")))
+
+	two, err := m.Allocate(mem.CellKindString)
+	assert.NoError(t, err)
+	assert.NoError(t, m.Set(two, value.StringCell("two")))
+
+	oneCell, err := m.Get(one)
+	assert.NoError(t, err)
+	assert.Equal(t, oneCell.(value.StringCell).Get(), "one")
+
+	twoCell, err := m.Get(two)
+	assert.NoError(t, err)
+	assert.Equal(t, twoCell.(value.StringCell).Get(), "two")
+
+	assert.NoError(t, m.Set(one, value.StringCell("one, but changed")))
+
+	oneCell, err = m.Get(one)
+	assert.NoError(t, err)
+	assert.Equal(t, oneCell.(value.StringCell).Get(), "one, but changed")
+}
+
+func TestNullIsForbidden(t *testing.T) {
+	m := mem.New()
+
+	_, err := m.Get(mem.NullPtr)
+	assert.Error(t, err)
+
+	assert.Error(t, m.Set(mem.NullPtr, value.StringCell("one")))
+
+	assert.Error(t, m.Retain(mem.NullPtr))
+	assert.Error(t, m.Release(mem.NullPtr))
+	assert.Equal(t, 0, m.RefCount(mem.NullPtr))
+}
+
+func TestCantGetNorSetUnallocatedCell(t *testing.T) {
+	m := mem.New()
+
+	_, err := m.Get(mem.Ptr(123))
+	assert.Error(t, err)
+
+	err = m.Set(mem.Ptr(123), value.StringCell("one"))
+	assert.Error(t, err)
+}
+
+func TestCantAllocateMarkerCells(t *testing.T) {
+	m := mem.New()
+
+	_, err := m.Allocate(mem.CellKindEmpty)
+	assert.Error(t, err)
+
+	_, err = m.Allocate(mem.CellKindForbidden)
+	assert.Error(t, err)
+}
+
+func TestCantAssignWrongCellKind(t *testing.T) {
+	m := mem.New()
+
+	one, err := m.Allocate(mem.CellKindString)
+	assert.NoError(t, err)
+	assert.Error(t, m.Set(one, value.EnvCell(value.NewEnv())))
+}
+
+func TestIsAndKind(t *testing.T) {
+	m := mem.New()
+
+	one, err := m.Allocate(mem.CellKindString)
+	assert.NoError(t, err)
+
+	assert.True(t, m.Is(one, mem.CellKindString))
+	assert.False(t, m.Is(one, mem.CellKindArray))
+
+	assert.Equal(t, mem.CellKindString, m.Kind(one))
+}
+
+func TestIsAndKindAlwaysAgree(t *testing.T) {
+	m := mem.New()
+
+	_, err := m.Allocate(mem.CellKindString)
+	assert.NoError(t, err)
+	_, err = m.Allocate(mem.CellKindArray)
+	assert.NoError(t, err)
+	_, err = m.Allocate(mem.CellKindOutlet)
+	assert.NoError(t, err)
+	_, err = m.Allocate(mem.CellKindEnv)
+	assert.NoError(t, err)
+
+	for i := 0; i < 100; i++ {
+		ptr := mem.Ptr(i)
+		assert.True(t, m.Is(ptr, m.Kind(ptr)))
+	}
+}
+
+func TestRetainRelease(t *testing.T) {
+	m := mem.New()
+
+	ptr, err := m.Allocate(mem.CellKindString)
+	assert.NoError(t, err)
+	assert.NoError(t, m.Set(ptr, value.StringCell("an allocated string")))
+
+	assert.True(t, m.Is(ptr, mem.CellKindString))
+	assert.Equal(t, 1, m.RefCount(ptr))
+
+	assert.NoError(t, m.Retain(ptr))
+	assert.NoError(t, m.Retain(ptr))
+	assert.NoError(t, m.Retain(ptr))
+
+	assert.True(t, m.Is(ptr, mem.CellKindString))
+	assert.Equal(t, 4, m.RefCount(ptr))
+
+	assert.NoError(t, m.Release(ptr))
+	assert.NoError(t, m.Release(ptr))
+	assert.NoError(t, m.Release(ptr))
+
+	assert.True(t, m.Is(ptr, mem.CellKindString))
+	cell, err := m.Get(ptr)
+	assert.NoError(t, err)
+	assert.Equal(t, "an allocated string", cell.(value.StringCell).Get())
+	assert.Equal(t, 1, m.RefCount(ptr))
+
+	assert.NoError(t, m.Release(ptr))
+
+	assert.Equal(t, 0, m.RefCount(ptr))
+	assert.True(t, m.Is(ptr, mem.CellKindEmpty))
+	_, err = m.Get(ptr)
+	assert.Error(t, err)
+}
+
+func TestCantReleaseTooMuch(t *testing.T) {
+	m := mem.New()
+
+	ptr, err := m.Allocate(mem.CellKindString)
+	assert.NoError(t, err)
+	assert.NoError(t, m.Set(ptr, value.StringCell("an allocated string")))
+
+	assert.NoError(t, m.Retain(ptr))
+	assert.Equal(t, 2, m.RefCount(ptr))
+
+	assert.NoError(t, m.Release(ptr))
+	assert.NoError(t, m.Release(ptr))
+	assert.Error(t, m.Release(ptr))
+}
+
+func TestCantRetainNorReleaseUnallocated(t *testing.T) {
+	m := mem.New()
+
+	assert.Error(t, m.Retain(mem.Ptr(123)))
+	assert.Error(t, m.Release(mem.Ptr(123)))
+}
+
+func TestChainDrop(t *testing.T) {
+	m := mem.New()
+
+	arr := []value.Value{
+		unwrap(t, func() (value.Value, error) { return value.NewString(m, "one") }),
+		unwrap(t, func() (value.Value, error) { return value.NewString(m, "two") }),
+
+		unwrap(t, func() (value.Value, error) {
+			return value.NewArray(m, []value.Value{
+				unwrap(t, func() (value.Value, error) { return value.NewString(m, "three") }),
+			})
+		}),
+	}
+
+	val, err := value.NewArray(m, arr)
+	assert.NoError(t, err)
+
+	assert.Equal(t, mem.CellKindString, m.Kind(mem.Ptr(1)))
+	assert.Equal(t, mem.CellKindString, m.Kind(mem.Ptr(2)))
+	assert.Equal(t, mem.CellKindString, m.Kind(mem.Ptr(3)))
+	assert.Equal(t, mem.CellKindArray, m.Kind(mem.Ptr(4)))
+	assert.Equal(t, mem.CellKindArray, m.Kind(mem.Ptr(5)))
+	assert.Equal(t, mem.CellKindEmpty, m.Kind(mem.Ptr(6)))
+
+	val.Drop(m)
+
+	for _, ptr := range []mem.Ptr{1, 2, 3, 4, 5, 6} {
+		assert.Equal(t, mem.CellKindEmpty, m.Kind(ptr))
+	}
+}
+
+func TestFreeCellsGetReused(t *testing.T) {
+	m := mem.New()
+
+	one, err := m.Allocate(mem.CellKindString)
+	assert.NoError(t, err)
+	assert.NoError(t, m.Set(one, value.StringCell("old one")))
+
+	two, err := m.Allocate(mem.CellKindString)
+	assert.NoError(t, err)
+	assert.NoError(t, m.Set(two, value.StringCell("old two")))
+
+	assert.NoError(t, m.Release(one))
+
+	newOne, err := m.Allocate(mem.CellKindString)
+	assert.NoError(t, err)
+	assert.NoError(t, m.Set(newOne, value.StringCell("new one")))
+
+	assert.Equal(t, mem.CellKindString, m.Kind(newOne))
+	newCell, err := m.Get(newOne)
+	assert.NoError(t, err)
+	assert.Equal(t, "new one", newCell.(value.StringCell).Get())
+	assert.Equal(t, 1, m.RefCount(newOne))
+	assert.Equal(t, one, newOne)
+}
+
+func unwrap[X any](t *testing.T, f func() (X, error)) X {
+	x, err := f()
+	assert.NoError(t, err)
+	return x
+}
diff --git a/pkg/lang/vm/value/cells.go b/pkg/lang/vm/value/cells.go
index 1c34762..116fa5a 100644
--- a/pkg/lang/vm/value/cells.go
+++ b/pkg/lang/vm/value/cells.go
@@ -10,6 +10,10 @@ func (a ArrayCell) DropCell(m mem.Mem) {
 	}
 }
 
+func (a ArrayCell) MatchingCellKind() mem.CellKind {
+	return mem.CellKindArray
+}
+
 func (a ArrayCell) Get() []Value {
 	return a
 }
@@ -19,6 +23,10 @@ type StringCell string
 func (s StringCell) DropCell(m mem.Mem) {
 }
 
+func (s StringCell) MatchingCellKind() mem.CellKind {
+	return mem.CellKindString
+}
+
 func (s StringCell) Get() string {
 	return string(s)
 }
@@ -29,6 +37,10 @@ func (o OutletCell) DropCell(m mem.Mem) {
 	Value(o).Drop(m)
 }
 
+func (o OutletCell) MatchingCellKind() mem.CellKind {
+	return mem.CellKindOutlet
+}
+
 func (o OutletCell) Get() Value {
 	return Value(o)
 }
@@ -41,6 +53,10 @@ func (e EnvCell) DropCell(m mem.Mem) {
 	}
 }
 
+func (e EnvCell) MatchingCellKind() mem.CellKind {
+	return mem.CellKindEnv
+}
+
 func (e EnvCell) Get() Env {
 	return Env(e)
 }