Function envs and value escaping

13 files changed, 361 insertions, 19 deletions

diff --git a/pkg/lang/vm/code/op.go b/pkg/lang/vm/code/op.go
index 84d468f..a0214ac 100644
--- a/pkg/lang/vm/code/op.go
+++ b/pkg/lang/vm/code/op.go
@@ -23,7 +23,10 @@ const (
 	OpGetLocal
 	OpGetMember
 	OpGetMethod
+
 	OpGetEnv
+	OpSetEnv
+	OpAddToEnv
 
 	OpAdd
 	OpSub
diff --git a/pkg/lang/vm/errors.go b/pkg/lang/vm/errors.go
index 37f3542..264dd3a 100644
--- a/pkg/lang/vm/errors.go
+++ b/pkg/lang/vm/errors.go
@@ -42,6 +42,15 @@ func (e ErrLocalIndexOutOfBounds) Error() string {
 	return fmt.Sprintf("local index out of bounds: %d (len: %d)", e.Index, e.Len)
 }
 
+type ErrStackIndexOutOfBounds struct {
+	Index int
+	Len   int
+}
+
+func (e ErrStackIndexOutOfBounds) Error() string {
+	return fmt.Sprintf("stack index out of bounds: %d (len: %d)", e.Index, e.Len)
+}
+
 type ErrInvalidOp struct {
 	Op uint8
 }
diff --git a/pkg/lang/vm/exec.go b/pkg/lang/vm/exec.go
index 90b79a9..ee977b7 100644
--- a/pkg/lang/vm/exec.go
+++ b/pkg/lang/vm/exec.go
@@ -2,6 +2,7 @@ package vm
 
 import (
 	"jinx/pkg/lang/vm/code"
+	"jinx/pkg/lang/vm/mem"
 	"jinx/pkg/lang/vm/value"
 )
 
@@ -45,6 +46,113 @@ func (vm *VM) execGetLocal(offset int) error {
 	return nil
 }
 
+func (vm *VM) execGetEnv(envIndex int) error {
+	envPtr := vm.stack.CurrentCallEnv()
+	env := vm.memory.Get(envPtr).(value.Env)
+
+	// First check outlet.
+	outletPtr := env.GetOutlet(envIndex)
+	outlet := vm.memory.Get(outletPtr)
+	if outlet != nil {
+		// Outlet is not null, so value escaped.
+		val := outlet.(value.Value).Clone(&vm.memory)
+		vm.stack.Push(val)
+		return nil
+	}
+
+	// Outlet is null, so value has not escaped.
+	stackIndex := env.GetStackIndex(envIndex)
+	val, err := vm.stack.Get(stackIndex)
+	if err != nil {
+		return err
+	}
+
+	val = val.Clone(&vm.memory)
+	vm.stack.Push(val)
+
+	return nil
+}
+
+func (vm *VM) execSetEnv(envIndex int) error {
+	new, err := vm.stack.Pop()
+	if err != nil {
+		return err
+	}
+
+	envPtr := vm.stack.CurrentCallEnv()
+	env := vm.memory.Get(envPtr).(value.Env)
+
+	outletPtr := env.GetOutlet(envIndex)
+	outlet := vm.memory.Get(outletPtr)
+	if outlet != nil {
+		outlet.(value.Value).Drop(&vm.memory)
+		vm.memory.Set(outletPtr, new)
+		return nil
+	}
+
+	stackIndex := env.GetStackIndex(envIndex)
+	old, err := vm.stack.Get(stackIndex)
+	if err != nil {
+		return err
+	}
+
+	old.Drop(&vm.memory)
+
+	return vm.stack.Set(stackIndex, new)
+}
+
+func (vm *VM) execAddToEnv(localIndex int) error {
+	f, err := vm.stack.Pop()
+	if err != nil {
+		return err
+	}
+
+	if f.Type().Kind != value.FunctionType {
+		return ErrInvalidOperandType{
+			Op: code.OpAddToEnv,
+			X:  f.Type(),
+		}
+	}
+
+	fn := f.Data().(value.FunctionData)
+
+	var envPtr mem.Ptr
+	if fn.Env().IsNull() {
+		// Allocate new Env.
+		envPtr = vm.memory.Allocate(mem.CellKindEnv)
+		vm.memory.Set(envPtr, value.NewEnv())
+		fn = fn.WithEnv(envPtr)
+	} else {
+		envPtr = fn.Env()
+	}
+
+	// Create outlet for referenced local, or use existing one.
+	local, err := vm.stack.Local(int(localIndex))
+	if err != nil {
+		return err
+	}
+
+	if local.Outlet().IsNull() {
+		outlet := vm.memory.Allocate(mem.CellKindOutlet)
+		local = local.WithOutlet(outlet)
+	}
+
+	// Add local to env.
+	stackIndex := vm.stack.LocalToStackIndex(localIndex)
+
+	env := vm.memory.Get(envPtr).(value.Env)
+	env.Add(stackIndex, local.Outlet())
+	vm.memory.Set(envPtr, env)
+
+	// Push function back onto stack.
+	f = f.WithData(fn)
+	vm.stack.Push(f)
+
+	vm.stack.Set(stackIndex, local)
+
+	return nil
+}
+
 func (vm *VM) execAdd() error {
 	x, err := vm.popAndDrop()
 	if err != nil {
@@ -273,7 +381,7 @@ func (vm *VM) execLte() error {
 }
 
 func (vm *VM) execCall() error {
-	f, err := vm.popAndDrop()
+	f, err := vm.stack.Pop()
 	if err != nil {
 		return err
 	}
@@ -287,7 +395,7 @@ func (vm *VM) execCall() error {
 
 	fn := f.Data().(value.FunctionData)
 
-	if err = vm.stack.PushCall(fn.Pc(), vm.pc); err != nil {
+	if err = vm.stack.PushCall(fn.Pc(), vm.pc, fn.Env()); err != nil {
 		return err
 	}
 
diff --git a/pkg/lang/vm/mem/cell.go b/pkg/lang/vm/mem/cell.go
index 044a45c..3f052d5 100644
--- a/pkg/lang/vm/mem/cell.go
+++ b/pkg/lang/vm/mem/cell.go
@@ -6,7 +6,11 @@ const (
 	CellKindEmpty CellKind = iota
 	CellKindString
 	CellKindArray
-	CellKindEscaped
+
+	CellKindEnv
+	CellKindOutlet
+
+	CellKindForbidden
 )
 
 type cell struct {
diff --git a/pkg/lang/vm/mem/mem.go b/pkg/lang/vm/mem/mem.go
index 4335347..4cb6fc4 100644
--- a/pkg/lang/vm/mem/mem.go
+++ b/pkg/lang/vm/mem/mem.go
@@ -6,8 +6,13 @@ type Mem struct {
 }
 
 func New() Mem {
+	cells := make([]cell, 1)
+
+	// Reserve NullPtr
+	cells[NullPtr].kind = CellKindForbidden
+
 	return Mem{
-		cells: make([]cell, 0),
+		cells: cells,
 		free:  make([]int, 0),
 	}
 }
@@ -25,7 +30,7 @@ func (m *Mem) Allocate(kind CellKind) Ptr {
 		return Ptr(idx)
 	} else {
 		idx := len(m.cells)
-		m.cells = append(m.cells, cell{kind: kind})
+		m.cells = append(m.cells, cell{kind: kind, refs: 1})
 		return Ptr(idx)
 	}
 
@@ -47,6 +52,14 @@ func (m *Mem) Get(ptr Ptr) any {
 	return m.cells[ptr].data
 }
 
+func (m *Mem) Is(ptr Ptr, kind CellKind) bool {
+	if ptr >= Ptr(len(m.cells)) {
+		panic("out of bounds")
+	}
+
+	return m.cells[ptr].kind == kind
+}
+
 func (m *Mem) Retain(ptr Ptr) {
 	if ptr >= Ptr(len(m.cells)) {
 		panic("out of bounds")
diff --git a/pkg/lang/vm/mem/ptr.go b/pkg/lang/vm/mem/ptr.go
index 5e3bf3d..f6e3e64 100644
--- a/pkg/lang/vm/mem/ptr.go
+++ b/pkg/lang/vm/mem/ptr.go
@@ -4,7 +4,15 @@ import "strconv"
 
 type Ptr uint64
 
+const (
+	NullPtr Ptr = 0
+)
+
 func (p Ptr) String() string {
 	val := strconv.FormatUint(uint64(p), 10)
 	return "@" + val
 }
+
+func (p Ptr) IsNull() bool {
+	return p == NullPtr
+}
diff --git a/pkg/lang/vm/stack.go b/pkg/lang/vm/stack.go
index 51e43c8..a8965ec 100644
--- a/pkg/lang/vm/stack.go
+++ b/pkg/lang/vm/stack.go
@@ -1,6 +1,7 @@
 package vm
 
 import (
+	"jinx/pkg/lang/vm/mem"
 	"jinx/pkg/lang/vm/value"
 )
 
@@ -43,17 +44,29 @@ func (stack *Stack) Pop() (value.Value, error) {
 	return v, nil
 }
 
+func (stack *Stack) Get(index int) (value.Value, error) {
+	if index < 0 || index >= stack.Len() {
+		return value.Value{}, ErrStackIndexOutOfBounds{Index: index, Len: stack.Len()}
+	}
+
+	return stack.data[index], nil
+}
+
+func (stack *Stack) Set(index int, value value.Value) error {
+	if index < 0 || index >= stack.Len() {
+		return ErrStackIndexOutOfBounds{Index: index, Len: stack.Len()}
+	}
+
+	stack.data[index] = value
+	return nil
+}
+
 func (stack *Stack) Local(offset int) (value.Value, error) {
 	if stack.ReachedBaseOfCall() {
 		return value.Value{}, ErrStackUnderflow
 	}
 
-	if offset < 0 || offset >= stack.Len() {
-		return value.Value{}, ErrLocalIndexOutOfBounds{Index: offset, Len: stack.Len()}
-	}
-
-	base := stack.TopCall().base
-	return stack.data[base+offset], nil
+	return stack.Get(stack.LocalToStackIndex(offset))
 }
 
 func (stack *Stack) Top() (value.Value, error) {
@@ -72,7 +85,7 @@ func (stack *Stack) IsEmpty() bool {
 	return len(stack.data) == 0
 }
 
-func (stack *Stack) PushCall(newPc, returnPc int) error {
+func (stack *Stack) PushCall(newPc, returnPc int, env mem.Ptr) error {
 	if stack.CallDepth() == 1000 {
 		return ErrReachedMaxCallDepth
 	}
@@ -81,6 +94,7 @@ func (stack *Stack) PushCall(newPc, returnPc int) error {
 		pc:       newPc,
 		returnPc: returnPc,
 		base:     stack.Len(),
+		env:      env,
 	})
 
 	return nil
@@ -97,6 +111,14 @@ func (stack *Stack) PopCall() (int, error) {
 	return call.returnPc, nil
 }
 
+func (stack *Stack) CurrentCallEnv() mem.Ptr {
+	if stack.CallDepth() == 0 || stack.TopCall().env.IsNull() {
+		return mem.NullPtr
+	}
+
+	return stack.TopCall().env
+}
+
 func (stack *Stack) ShiftTopCallBase(by int) error {
 	call := stack.TopCall()
 	newBase := call.base - by
@@ -121,8 +143,16 @@ func (stack *Stack) ReachedBaseOfCall() bool {
 	return stack.TopCall().base == stack.Len()
 }
 
+func (stack *Stack) LocalToStackIndex(local int) int {
+	if stack.CallDepth() == 0 {
+		return local
+	}
+	return stack.TopCall().base + local
+}
+
 type callFrame struct {
-	pc       int // Beginning of the called function.
-	returnPc int // Where to return to after the called function returns.
-	base     int // Base of the local variables on the data stack.
+	pc       int     // Beginning of the called function.
+	returnPc int     // Where to return to after the called function returns.
+	base     int     // Base of the local variables on the data stack.
+	env      mem.Ptr // Environment of the called function.
 }
diff --git a/pkg/lang/vm/text/op.go b/pkg/lang/vm/text/op.go
index df40d8f..dc6ac1e 100644
--- a/pkg/lang/vm/text/op.go
+++ b/pkg/lang/vm/text/op.go
@@ -22,6 +22,8 @@ var (
 		code.OpGetMember:    "get_member",
 		code.OpGetMethod:    "get_method",
 		code.OpGetEnv:       "get_env",
+		code.OpSetEnv:       "set_env",
+		code.OpAddToEnv:     "add_to_env",
 		code.OpAdd:          "add",
 		code.OpSub:          "sub",
 		code.OpIndex:        "index",
diff --git a/pkg/lang/vm/value/data.go b/pkg/lang/vm/value/data.go
index fecf8cf..a49753e 100644
--- a/pkg/lang/vm/value/data.go
+++ b/pkg/lang/vm/value/data.go
@@ -91,13 +91,22 @@ func (a ArrayData) Push(m *mem.Mem, v Value) {
 type NullData struct{}
 
 type FunctionData struct {
-	pc int
+	pc  int
+	env mem.Ptr
 }
 
 func (f FunctionData) Pc() int {
 	return f.pc
 }
 
+func (f FunctionData) Env() mem.Ptr {
+	return f.env
+}
+
+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)
 }
diff --git a/pkg/lang/vm/value/env.go b/pkg/lang/vm/value/env.go
new file mode 100644
index 0000000..b7bb3ab
--- /dev/null
+++ b/pkg/lang/vm/value/env.go
@@ -0,0 +1,33 @@
+package value
+
+import "jinx/pkg/lang/vm/mem"
+
+type Env struct {
+	references []reference
+}
+
+func NewEnv() Env {
+	return Env{
+		references: make([]reference, 0),
+	}
+}
+
+func (e *Env) Add(stackIndex int, outlet mem.Ptr) {
+	e.references = append(e.references, reference{
+		stackIndex: stackIndex,
+		outlet:     outlet,
+	})
+}
+
+func (e *Env) GetOutlet(envIndex int) mem.Ptr {
+	return e.references[envIndex].outlet
+}
+
+func (e *Env) GetStackIndex(envIndex int) int {
+	return e.references[envIndex].stackIndex
+}
+
+type reference struct {
+	stackIndex int
+	outlet     mem.Ptr
+}
diff --git a/pkg/lang/vm/value/value.go b/pkg/lang/vm/value/value.go
index 06306df..d19bbd6 100644
--- a/pkg/lang/vm/value/value.go
+++ b/pkg/lang/vm/value/value.go
@@ -5,8 +5,9 @@ import (
 )
 
 type Value struct {
-	t Type
-	d Data
+	t      Type
+	d      Data
+	outlet mem.Ptr
 }
 
 func NewInt(x int64) Value {
@@ -64,6 +65,18 @@ func (v Value) Data() Data {
 	return v.d
 }
 
+func (v Value) WithData(d Data) Value {
+	return Value{t: v.t, d: d, outlet: v.outlet}
+}
+
+func (v Value) Outlet() mem.Ptr {
+	return v.outlet
+}
+
+func (v Value) WithOutlet(outlet mem.Ptr) Value {
+	return Value{t: v.t, d: v.d, outlet: outlet}
+}
+
 func (v Value) Clone(m *mem.Mem) Value {
 	if v.t.Kind == StringType {
 		str := v.d.(StringData)
@@ -75,10 +88,21 @@ func (v Value) Clone(m *mem.Mem) Value {
 		m.Retain(arr.data)
 	}
 
+	if v.t.Kind == FunctionType {
+		fn := v.d.(FunctionData)
+		m.Retain(fn.env)
+	}
+
 	return v
 }
 
 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)
+		return
+	}
+
 	if v.t.Kind == StringType {
 		str := v.d.(StringData)
 		m.Release(str.data)
@@ -88,4 +112,9 @@ func (v Value) Drop(m *mem.Mem) {
 		arr := v.d.(ArrayData)
 		m.Release(arr.data)
 	}
+
+	if v.t.Kind == FunctionType {
+		f := v.d.(FunctionData)
+		m.Release(f.env)
+	}
 }
diff --git a/pkg/lang/vm/vm.go b/pkg/lang/vm/vm.go
index 7284257..c933af4 100644
--- a/pkg/lang/vm/vm.go
+++ b/pkg/lang/vm/vm.go
@@ -116,8 +116,22 @@ func (vm *VM) step(op code.Op) (stepDecision, error) {
 		panic("not implemented")
 	case code.OpGetMethod:
 		panic("not implemented")
+
 	case code.OpGetEnv:
-		panic("not implemented")
+		envIndex, advance := vm.code.GetUint(vm.pc)
+		vm.pc += advance
+
+		err = vm.execGetEnv(int(envIndex))
+	case code.OpSetEnv:
+		envIndex, advance := vm.code.GetUint(vm.pc)
+		vm.pc += advance
+
+		err = vm.execSetEnv(int(envIndex))
+	case code.OpAddToEnv:
+		local, advance := vm.code.GetUint(vm.pc)
+		vm.pc += advance
+
+		err = vm.execAddToEnv(int(local))
 
 	case code.OpAdd:
 		err = vm.execAdd()
@@ -166,6 +180,9 @@ func (vm *VM) popAndDrop() (value.Value, error) {
 }
 
 func (vm *VM) popCallAndDrop() (int, error) {
+	envPtr := vm.stack.CurrentCallEnv()
+	vm.memory.Release(envPtr)
+
 	for !vm.stack.ReachedBaseOfCall() {
 		_, err := vm.popAndDrop()
 		if err != nil {
diff --git a/pkg/lang/vm/vm_test.go b/pkg/lang/vm/vm_test.go
index 74e64d0..fd4f961 100644
--- a/pkg/lang/vm/vm_test.go
+++ b/pkg/lang/vm/vm_test.go
@@ -108,6 +108,83 @@ func TestFunction(t *testing.T) {
 	test(t, src, "42")
 }
 
+func TestSimpleEnv(t *testing.T) {
+	src := `
+    push_int 404
+    push_int 1
+    push_int 405
+
+    push_function @test
+    # Add the local 1 to the environment.
+    add_to_env 1
+    call
+    halt
+
+    @test:
+        push_int 2
+        get_env 0
+        add
+        ret
+    `
+
+	test(t, src, "3")
+}
+
+func TestEscapedEnv(t *testing.T) {
+	/*
+	   fn create() {
+	       var x = 0
+	       return fn () {
+	           x = x + 1
+	           return x
+	       }
+	   }
+
+	   var f = create()
+	   var res = [f(), f(), f()]
+	*/
+
+	src := `    
+    push_function @create
+    call
+
+    push_array
+
+    get_local 0
+    call
+    get_local 1
+    temp_arr_push
+
+    get_local 0
+    call
+    get_local 1
+    temp_arr_push
+    
+    get_local 0
+    call
+    get_local 1
+    temp_arr_push
+    halt
+
+    @create:
+        push_int 0
+        push_int 404
+        push_int 405
+        push_function @create:anon_0
+        add_to_env 0
+        ret
+    @create:anon_0:
+        push_int 1
+        get_env 0
+        add
+        set_env 0
+        get_env 0
+        ret
+    `
+
+	test(t, src, "[1, 2, 3]")
+}
+
 func test(t *testing.T, src string, expected string) {
 	bc := compile(t, src)
 	vm := vm.New(&bc)