Implement stack hygiene

1 files changed, 20 insertions, 4 deletions

diff --git a/pkg/lang/compiler/scope/scope_chain.go b/pkg/lang/compiler/scope/scope_chain.go
index 4d5bbcf..d0108ee 100644
--- a/pkg/lang/compiler/scope/scope_chain.go
+++ b/pkg/lang/compiler/scope/scope_chain.go
@@ -84,13 +84,29 @@ func (sc *ScopeChain) EnterLoop() (code.Marker, code.Marker) {
 	return breakMarker, continueMarker
 }
 
-func (sc *ScopeChain) Exit() {
+func (sc *ScopeChain) Exit() int {
 	if sc.CurrentScopeID() == 0 {
-		return
+		return 0
 	}
+	id := sc.CurrentScopeID()
+	stackSpace := len(sc.symbolScopes[id].variableSymbols)
 
-	sc.symbolScopes[sc.CurrentScopeID()] = SymbolScope{}
-	sc.symbolScopes = sc.symbolScopes[:sc.CurrentScopeID()]
+	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]
+	}
+
+	if sc.CurrentFunction().id == id {
+		lastFunctionScope := len(sc.functionScopes) - 1
+		sc.functionScopes[lastFunctionScope] = FunctionScope{}
+		sc.functionScopes = sc.functionScopes[:lastFunctionScope]
+	}
+
+	return stackSpace
 }
 
 func (sc *ScopeChain) Declare(name string) (int, bool) {