Compile function environments

3 files changed, 80 insertions, 11 deletions

diff --git a/pkg/lang/compiler/scope/scope_chain.go b/pkg/lang/compiler/scope/scope_chain.go
index f4b9f46..1b83c75 100644
--- a/pkg/lang/compiler/scope/scope_chain.go
+++ b/pkg/lang/compiler/scope/scope_chain.go
@@ -91,19 +91,14 @@ func (sc *ScopeChain) Exit() int {
 	id := sc.CurrentScopeID()
 	stackSpace := len(sc.symbolScopes[id].variableSymbols)
 
-	sc.symbolScopes[id] = SymbolScope{}
 	sc.symbolScopes = sc.symbolScopes[:id]
 
 	if sc.CurrentLoop() != nil && sc.CurrentLoop().id == id {
-		lastLoopScope := len(sc.loopScopes) - 1
-		sc.loopScopes[lastLoopScope] = LoopScope{}
-		sc.loopScopes = sc.loopScopes[:lastLoopScope]
+		sc.loopScopes = sc.loopScopes[:len(sc.loopScopes)-1]
 	}
 
 	if sc.CurrentFunction().id == id {
-		lastFunctionScope := len(sc.functionScopes) - 1
-		sc.functionScopes[lastFunctionScope] = FunctionScope{}
-		sc.functionScopes = sc.functionScopes[:lastFunctionScope]
+		sc.functionScopes = sc.functionScopes[:len(sc.functionScopes)-1]
 	}
 
 	return stackSpace
@@ -178,11 +173,23 @@ func (sc *ScopeChain) CreateFunctionSubUnit(subUnitName string) code.Marker {
 }
 
 func (sc *ScopeChain) Lookup(name string) (SymbolID, bool) {
-	if id, ok := sc.nameToSymbol[name]; ok {
-		return id, true
+	id, ok := sc.nameToSymbol[name]
+	if !ok {
+		return SymbolID{}, false
 	}
 
-	return SymbolID{}, false
+	// Check whether the symbol is outside the current function scope.
+	fnScope := sc.CurrentFunction()
+	if id.scopeID < fnScope.id {
+		// Return env symbol instead of a local symbol.
+		return SymbolID{
+			symbolKind:   SymbolKindEnv,
+			scopeID:      id.scopeID,
+			indexInScope: id.indexInScope,
+		}, true
+	}
+
+	return id, true
 }
 
 func (sc *ScopeChain) GetVariable(id SymbolID) Symbol[SymbolVariable] {
@@ -192,3 +199,36 @@ func (sc *ScopeChain) GetVariable(id SymbolID) Symbol[SymbolVariable] {
 
 	return sc.symbolScopes[id.scopeID].variableSymbols[id.indexInScope]
 }
+
+func (sc *ScopeChain) GetEnv(id SymbolID) Symbol[SymbolEnv] {
+	if id.symbolKind != SymbolKindEnv {
+		panic("incorrect symbol id kind given")
+	}
+
+	symbol := sc.symbolScopes[id.scopeID].variableSymbols[id.indexInScope]
+
+	// Add the local to the function scope, if it is not already there.
+	fnScope := sc.CurrentFunction()
+
+	var indexInEnv int
+	alreadyUsed := false
+	for i, env := range fnScope.outsideSymbolsInEnv {
+		if env == id {
+			alreadyUsed = true
+			indexInEnv = i
+			break
+		}
+	}
+
+	if !alreadyUsed {
+		indexInEnv = len(fnScope.outsideSymbolsInEnv)
+		fnScope.outsideSymbolsInEnv = append(fnScope.outsideSymbolsInEnv, id)
+	}
+
+	return Symbol[SymbolEnv]{
+		name: symbol.name,
+		data: SymbolEnv{
+			indexInEnv: indexInEnv,
+		},
+	}
+}
diff --git a/pkg/lang/compiler/scope/scopes.go b/pkg/lang/compiler/scope/scopes.go
index e34b45a..7a1b20c 100644
--- a/pkg/lang/compiler/scope/scopes.go
+++ b/pkg/lang/compiler/scope/scopes.go
@@ -26,6 +26,8 @@ type FunctionScope struct {
 	id           ScopeID
 	unit         code.Marker
 	subUnitCount int
+
+	outsideSymbolsInEnv []SymbolID
 }
 
 func NewFunctionScope(id ScopeID, unit code.Marker) FunctionScope {
@@ -33,6 +35,8 @@ func NewFunctionScope(id ScopeID, unit code.Marker) FunctionScope {
 		id:           id,
 		unit:         unit,
 		subUnitCount: 0,
+
+		outsideSymbolsInEnv: make([]SymbolID, 0),
 	}
 }
 
@@ -44,6 +48,10 @@ func (sf FunctionScope) Unit() code.Marker {
 	return sf.unit
 }
 
+func (sf FunctionScope) OutsideSymbolsInEnv() []SymbolID {
+	return sf.outsideSymbolsInEnv
+}
+
 func (sf FunctionScope) IsRootScope() bool {
 	return sf.ID() == ScopeID(0)
 }
diff --git a/pkg/lang/compiler/scope/symbol.go b/pkg/lang/compiler/scope/symbol.go
index 3b50108..b87d5aa 100644
--- a/pkg/lang/compiler/scope/symbol.go
+++ b/pkg/lang/compiler/scope/symbol.go
@@ -10,17 +10,30 @@ func (id SymbolID) SymbolKind() SymbolKind {
 	return id.symbolKind
 }
 
+func (id SymbolID) ScopeID() ScopeID {
+	return id.scopeID
+}
+
+func (id SymbolID) IndexInScope() int {
+	return id.indexInScope
+}
+
 type SymbolKind int
 
 const (
 	// A variable symbol is bound to a local on the stack.
 	SymbolKindVariable SymbolKind = iota
+	// An env symbol is bound to a local on the stack, outside of the function's scope.
+	// Emitted at lookup time, so the SymbolScope has no array for them.
+	SymbolKindEnv SymbolKind = iota
 )
 
 func (s SymbolKind) String() string {
 	switch s {
 	case SymbolKindVariable:
 		return "variable"
+	case SymbolKindEnv:
+		return "env"
 	default:
 		panic("unknown symbol kind")
 	}
@@ -36,7 +49,7 @@ func (s Symbol[D]) Data() D {
 }
 
 type SymbolData interface {
-	SymbolVariable
+	SymbolVariable | SymbolEnv
 }
 
 type SymbolVariable struct {
@@ -46,3 +59,11 @@ type SymbolVariable struct {
 func (sv SymbolVariable) LocalIndex() int {
 	return sv.localIndex
 }
+
+type SymbolEnv struct {
+	indexInEnv int
+}
+
+func (se SymbolEnv) IndexInEnv() int {
+	return se.indexInEnv
+}