3 files changed, 32 insertions, 22 deletions

diff --git a/pkg/lang/vm/exec.go b/pkg/lang/vm/exec.go
index 8cdcdcc..c2d847f 100644
--- a/pkg/lang/vm/exec.go
+++ b/pkg/lang/vm/exec.go
@@ -987,7 +987,7 @@ func (vm *VM) execCall(argCount uint) error {
 		}
 	}
 
-	if err = vm.stack.PushCall(fn.Pos(), vm.pos, fn.Env()); err != nil {
+	if err = vm.stack.PushCall(fn, vm.pos); err != nil {
 		return err
 	}
 
@@ -997,7 +997,7 @@ func (vm *VM) execCall(argCount uint) error {
 			return err
 		}
 
-		if _, err := vm.stack.PopCall(); err != nil {
+		if _, _, err := vm.stack.PopCall(); err != nil {
 			return err
 		}
 
@@ -1007,6 +1007,11 @@ func (vm *VM) execCall(argCount uint) error {
 			vm.stack.Push(value.NewNull())
 		}
 
+		// If the function was native we can drop it immidiately, after the call.
+		if err := f.Drop(vm.memory); err != nil {
+			return err
+		}
+
 		return nil
 	}
 
@@ -1016,6 +1021,8 @@ func (vm *VM) execCall(argCount uint) error {
 
 	vm.setPos(fn.Pos())
 
+	// Note that value 'f' is still not dropped, as it's enviroment has
+	// been stored in the call stack. It will be dropped on the next return.
 	return nil
 }
 
diff --git a/pkg/lang/vm/stack/stack.go b/pkg/lang/vm/stack/stack.go
index 74d4ab8..c710878 100644
--- a/pkg/lang/vm/stack/stack.go
+++ b/pkg/lang/vm/stack/stack.go
@@ -19,8 +19,8 @@ type Stack interface {
 	IsEmpty() bool
 
 	PutRootCall(basePos code.Pos) error
-	PushCall(newPos, returnPos code.Pos, env mem.Ptr) error
-	PopCall() (code.Pos, error)
+	PushCall(fn value.FunctionData, returnPos code.Pos) error
+	PopCall() (value.FunctionData, code.Pos, error)
 
 	CurrentCallEnv() mem.Ptr
 	CallDepth() int
@@ -104,10 +104,9 @@ func (stack *stackImpl) IsEmpty() bool {
 
 func (stack *stackImpl) PutRootCall(basePos code.Pos) error {
 	frame := callFrame{
-		pos:       basePos,
+		fn:        value.FunctionData{},
 		returnPos: basePos,
 		base:      0,
-		env:       mem.NullPtr,
 	}
 
 	if stack.CallDepth() == 0 {
@@ -125,38 +124,37 @@ func (stack *stackImpl) PutRootCall(basePos code.Pos) error {
 	return nil
 }
 
-func (stack *stackImpl) PushCall(newPos, returnPos code.Pos, env mem.Ptr) error {
+func (stack *stackImpl) PushCall(fn value.FunctionData, returnPos code.Pos) error {
 	if stack.CallDepth() == 1000 {
 		return ErrReachedMaxCallDepth
 	}
 
 	stack.calls = append(stack.calls, callFrame{
-		pos:       newPos,
+		fn:        fn,
 		returnPos: returnPos,
 		base:      stack.Len(),
-		env:       env,
 	})
 
 	return nil
 }
 
-func (stack *stackImpl) PopCall() (code.Pos, error) {
+func (stack *stackImpl) PopCall() (value.FunctionData, code.Pos, error) {
 	if stack.CallDepth() == 0 {
-		return code.Pos{}, ErrReachedRootCallFrame
+		return value.FunctionData{}, code.Pos{}, ErrReachedRootCallFrame
 	}
 
 	call := stack.topCall()
 	stack.calls = stack.calls[:stack.CallDepth()-1]
 
-	return call.returnPos, nil
+	return call.fn, call.returnPos, nil
 }
 
 func (stack *stackImpl) CurrentCallEnv() mem.Ptr {
-	if stack.CallDepth() == 0 || stack.topCall().env.IsNull() {
+	if stack.CallDepth() == 0 || stack.topCall().fn.Env().IsNull() {
 		return mem.NullPtr
 	}
 
-	return stack.topCall().env
+	return stack.topCall().fn.Env()
 }
 
 func (stack *stackImpl) CallDepth() int {
@@ -179,8 +177,7 @@ func (stack *stackImpl) topCall() *callFrame {
 }
 
 type callFrame struct {
-	pos       code.Pos // Beginning of the called function.
-	returnPos code.Pos // 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.
+	fn        value.FunctionData // The called function.
+	returnPos code.Pos           // Where to return to after the called function returns.
+	base      int                // Base of the local variables on the data stack.
 }
diff --git a/pkg/lang/vm/utils.go b/pkg/lang/vm/utils.go
index eddcc11..166d404 100644
--- a/pkg/lang/vm/utils.go
+++ b/pkg/lang/vm/utils.go
@@ -19,9 +19,6 @@ func (vm *VM) popAndDrop() (value.Value, error) {
 }
 
 func (vm *VM) popCallAndDrop() (code.Pos, error) {
-	envPtr := vm.stack.CurrentCallEnv()
-	vm.memory.Release(envPtr)
-
 	for !vm.stack.ReachedBaseOfCall() {
 		_, err := vm.popAndDrop()
 		if err != nil {
@@ -29,7 +26,16 @@ func (vm *VM) popCallAndDrop() (code.Pos, error) {
 		}
 	}
 
-	return vm.stack.PopCall()
+	fn, returnPos, err := vm.stack.PopCall()
+	if err != nil {
+		return code.Pos{}, err
+	}
+
+	if err := value.NewFunction(code.Pos{}, 0).WithData(fn).Drop(vm.memory); err != nil {
+		return code.Pos{}, err
+	}
+
+	return returnPos, nil
 }
 
 func (vm *VM) getMemCell(ptr mem.Ptr, kind mem.CellKind, allowNil bool) (mem.CellData, error) {