1 files changed, 110 insertions, 2 deletions

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
 	}