1 files changed, 66 insertions, 40 deletions

diff --git a/pkg/lang/vm/stack.go b/pkg/lang/vm/stack.go
index 9cf7db8..937be32 100644
--- a/pkg/lang/vm/stack.go
+++ b/pkg/lang/vm/stack.go
@@ -4,70 +4,96 @@ import (
 	"jinx/pkg/lang/vm/value"
 )
 
-type CallStack []*LocalStack
-
-func NewCallStack() CallStack {
-	return []*LocalStack{{}}
+type Stack struct {
+	data  []value.Value
+	calls []callFrame
 }
 
-func (cs *CallStack) Push() error {
-	if len(*cs) > 1000 {
-		return ErrCallStackOverflow
+func NewStack() Stack {
+	data := make([]value.Value, 0, 64)
+	calls := make([]callFrame, 0, 8)
+
+	calls = append(calls, callFrame{
+		pc:       0,
+		returnPc: 0,
+		base:     0,
+	})
+
+	return Stack{
+		data:  data,
+		calls: calls,
 	}
+}
 
-	*cs = append(*cs, &LocalStack{})
-	return nil
+func (stack *Stack) Push(value value.Value) {
+	stack.data = append(stack.data, value)
 }
 
-func (cs *CallStack) Pop() error {
-	if len(*cs) <= 1 {
-		return ErrCantPopRootFrame
+func (stack *Stack) Pop() (value.Value, error) {
+	if stack.IsEmpty() || stack.ReachedBaseOfCall() {
+		return value.Value{}, ErrStackUnderflow
 	}
 
-	*cs = (*cs)[:len(*cs)-1]
-	return nil
+	v, err := stack.Top()
+	if err != nil {
+		return value.Value{}, err
+	}
+
+	stack.data = stack.data[:stack.Len()-1]
+	return v, nil
 }
 
-func (cs *CallStack) Top() *LocalStack {
-	return (*cs)[len(*cs)-1]
+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
 }
 
-func (cs *CallStack) Prev() (*LocalStack, error) {
-	if len(*cs) <= 1 {
-		return nil, ErrNoPreviousCallFrame
+func (stack *Stack) Top() (value.Value, error) {
+	if stack.IsEmpty() || stack.ReachedBaseOfCall() {
+		return value.Value{}, ErrStackUnderflow
 	}
 
-	return (*cs)[len(*cs)-2], nil
+	return stack.data[stack.Len()-1], nil
 }
 
-type LocalStack []value.Value
+func (stack *Stack) ShiftTopCallBase(by int) error {
+	call := stack.TopCall()
+	newBase := call.base - by
 
-func (ls *LocalStack) Push(v value.Value) error {
-	if len(*ls) > 1000 {
-		return ErrLocalStackOverflow
+	if newBase < 0 {
+		return ErrCallBaseCantBeNegative
 	}
 
-	*ls = append(*ls, v)
+	call.base = newBase
 	return nil
 }
 
-func (ls *LocalStack) Pop() (value.Value, error) {
-	if len(*ls) == 0 {
-		return value.Value{}, ErrCallFrameEmpty
-	}
+func (stack *Stack) Len() int {
+	return len(stack.data)
+}
 
-	v := (*ls)[len(*ls)-1]
-	*ls = (*ls)[:len(*ls)-1]
-	return v, nil
+func (stack *Stack) IsEmpty() bool {
+	return len(stack.data) == 0
 }
 
-func (ls *LocalStack) At(at int) (value.Value, error) {
-	if at >= len(*ls) {
-		return value.Value{}, ErrLocalIndexOutOfBounds{
-			Index: at,
-			Len:   len(*ls),
-		}
-	}
+func (stack *Stack) TopCall() *callFrame {
+	return &stack.calls[len(stack.calls)-1]
+}
+
+func (stack *Stack) ReachedBaseOfCall() bool {
+	return stack.TopCall().base == stack.Len()
+}
 
-	return (*ls)[at], nil
+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.
 }