1 files changed, 40 insertions, 10 deletions

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.
 }