2 files changed, 40 insertions, 20 deletions

diff --git a/pkg/lang/vm/stack/errors.go b/pkg/lang/vm/stack/errors.go
index 55ef2ea..55145d5 100644
--- a/pkg/lang/vm/stack/errors.go
+++ b/pkg/lang/vm/stack/errors.go
@@ -11,6 +11,7 @@ var (
 
 	ErrReachedMaxCallDepth  = errors.New("reached max call depth (max depth: 1000)")
 	ErrReachedRootCallFrame = errors.New("reached root call frame")
+	ErrNotAtRootCallFrame   = errors.New("not at root call frame")
 
 	ErrCallBaseCantBeNegative = errors.New("call base cannot be negative")
 )
diff --git a/pkg/lang/vm/stack/stack.go b/pkg/lang/vm/stack/stack.go
index 34ca896..74d4ab8 100644
--- a/pkg/lang/vm/stack/stack.go
+++ b/pkg/lang/vm/stack/stack.go
@@ -1,6 +1,7 @@
 package stack
 
 import (
+	"jinx/pkg/lang/vm/code"
 	"jinx/pkg/lang/vm/mem"
 	"jinx/pkg/lang/vm/value"
 )
@@ -17,8 +18,9 @@ type Stack interface {
 	Len() int
 	IsEmpty() bool
 
-	PushCall(newPc, returnPc int, env mem.Ptr) error
-	PopCall() (int, error)
+	PutRootCall(basePos code.Pos) error
+	PushCall(newPos, returnPos code.Pos, env mem.Ptr) error
+	PopCall() (code.Pos, error)
 
 	CurrentCallEnv() mem.Ptr
 	CallDepth() int
@@ -35,12 +37,6 @@ func New() Stack {
 	data := make([]value.Value, 0, 64)
 	calls := make([]callFrame, 0, 8)
 
-	calls = append(calls, callFrame{
-		pc:       0,
-		returnPc: 0,
-		base:     0,
-	})
-
 	return &stackImpl{
 		data:  data,
 		calls: calls,
@@ -106,30 +102,53 @@ func (stack *stackImpl) IsEmpty() bool {
 	return len(stack.data) == 0
 }
 
-func (stack *stackImpl) PushCall(newPc, returnPc int, env mem.Ptr) error {
+func (stack *stackImpl) PutRootCall(basePos code.Pos) error {
+	frame := callFrame{
+		pos:       basePos,
+		returnPos: basePos,
+		base:      0,
+		env:       mem.NullPtr,
+	}
+
+	if stack.CallDepth() == 0 {
+		stack.calls = append(stack.calls, frame)
+
+		return nil
+	}
+
+	if stack.CallDepth() != 1 {
+		return ErrNotAtRootCallFrame
+	}
+
+	stack.calls[0] = frame
+
+	return nil
+}
+
+func (stack *stackImpl) PushCall(newPos, returnPos code.Pos, env mem.Ptr) error {
 	if stack.CallDepth() == 1000 {
 		return ErrReachedMaxCallDepth
 	}
 
 	stack.calls = append(stack.calls, callFrame{
-		pc:       newPc,
-		returnPc: returnPc,
-		base:     stack.Len(),
-		env:      env,
+		pos:       newPos,
+		returnPos: returnPos,
+		base:      stack.Len(),
+		env:       env,
 	})
 
 	return nil
 }
 
-func (stack *stackImpl) PopCall() (int, error) {
+func (stack *stackImpl) PopCall() (code.Pos, error) {
 	if stack.CallDepth() == 0 {
-		return 0, ErrReachedRootCallFrame
+		return code.Pos{}, ErrReachedRootCallFrame
 	}
 
 	call := stack.topCall()
 	stack.calls = stack.calls[:stack.CallDepth()-1]
 
-	return call.returnPc, nil
+	return call.returnPos, nil
 }
 
 func (stack *stackImpl) CurrentCallEnv() mem.Ptr {
@@ -160,8 +179,8 @@ func (stack *stackImpl) topCall() *callFrame {
 }
 
 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.
-	env      mem.Ptr // Environment of the called function.
+	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.
 }