1 files changed, 48 insertions, 5 deletions

diff --git a/pkg/lang/compiler/compiler.go b/pkg/lang/compiler/compiler.go
index 4167088..e446dd6 100644
--- a/pkg/lang/compiler/compiler.go
+++ b/pkg/lang/compiler/compiler.go
@@ -96,8 +96,8 @@ func (comp *Compiler) compileStmt(t *code.Builder, stmt ast.Stmt) error {
 	case ast.StmtKindThrow:
 		panic("throw statements not implemented")
 	case ast.StmtKindExpr:
-		expr := stmt.Value.(ast.StmtExpr).Value
-		err = comp.compileExpr(t, expr)
+		exprStmt := stmt.Value.(ast.StmtExpr)
+		err = comp.compileExprStmt(t, exprStmt)
 	default:
 		panic(fmt.Errorf("unknown statement kind: %d", stmt.Kind))
 	}
@@ -139,7 +139,8 @@ func (comp *Compiler) compileFnDeclStmt(t *code.Builder, fnDeclStmt ast.StmtFnDe
 		return err
 	}
 
-	comp.scopes.Exit()
+	// Function declaration scopes do not pollute stack
+	_ = comp.scopes.Exit()
 
 	comp.funcs = append(comp.funcs, &functionTarget)
 
@@ -175,6 +176,8 @@ func (comp *Compiler) compileIfStmt(t *code.Builder, ifStmt ast.StmtIf) error {
 	//				   preventing other CondNodes from running. This is missing from the last CondNode.
 	//    Example: `jmp @end`
 
+	comp.scopes.Enter()
+
 	// First we create all the markers we'll need for the if statement
 	parentMarker := comp.scopes.CreateAnonymousFunctionSubUnit()
 
@@ -226,6 +229,8 @@ func (comp *Compiler) compileIfStmt(t *code.Builder, ifStmt ast.StmtIf) error {
 
 	t.PutMarker(endMarker)
 
+	comp.exitScopeAndCleanStack(t)
+
 	return nil
 }
 
@@ -251,18 +256,24 @@ func (comp *Compiler) compileForCondStmt(t *code.Builder, forCondStmt ast.StmtFo
 		t.AppendMarkerReference(endMarker)
 	}
 
+	// Inner scope, dropped on every iteration
+	comp.scopes.Enter()
+
 	// Do block
 	if err := comp.compileBlockNode(t, forCondStmt.Do); err != nil {
 		return err
 	}
 
+	// Drop inner scope
+	comp.exitScopeAndCleanStack(t)
+
 	// Repeat jump
 	t.AppendOp(code.OpJmp)
 	t.AppendMarkerReference(repeatMarker)
 
 	t.PutMarker(endMarker)
 
-	comp.scopes.Exit()
+	comp.exitScopeAndCleanStack(t)
 
 	return nil
 }
@@ -302,6 +313,7 @@ func (comp *Compiler) compileForInStmt(t *code.Builder, forInStmt ast.StmtForIn)
 	// @end:
 	// halt
 
+	// Upper scope houses all internal loop locals, which are dropped when the loop ends.
 	endMarker, repeatMarker := comp.scopes.EnterLoop()
 
 	// Preparation
@@ -364,18 +376,25 @@ func (comp *Compiler) compileForInStmt(t *code.Builder, forInStmt ast.StmtForIn)
 	t.AppendOp(code.OpSetLocal)
 	t.AppendInt(int64(iLocal))
 
+	// Inner scope, dropped every loop iteration.
+	comp.scopes.Enter()
+
 	// Do block
 	if err := comp.compileBlockNode(t, forInStmt.Do); err != nil {
 		return err
 	}
 
+	// Drop inner scope
+	comp.exitScopeAndCleanStack(t)
+
 	// Repeat jump
 	t.AppendOp(code.OpJmp)
 	t.AppendMarkerReference(repeatMarker)
 
 	t.PutMarker(endMarker)
 
-	comp.scopes.Exit()
+	// Drop upper scope
+	comp.exitScopeAndCleanStack(t)
 
 	return nil
 }
@@ -424,6 +443,23 @@ func (comp *Compiler) compileBreakStmt(t *code.Builder, breakStmt ast.StmtBreak)
 	return nil
 }
 
+func (comp *Compiler) compileExprStmt(t *code.Builder, exprStmt ast.StmtExpr) error {
+	if err := comp.compileExpr(t, exprStmt.Value); err != nil {
+		return err
+	}
+
+	isAssignment := exprStmt.Value.Kind == ast.ExprKindBinary &&
+		exprStmt.Value.Value.(ast.ExprBinary).Op == ast.BinOpAssign
+
+	// If the expression is not assignment, we need to drop the junk value.
+	if !isAssignment {
+		t.AppendOp(code.OpDrop)
+		t.AppendInt(1)
+	}
+
+	return nil
+}
+
 func (comp *Compiler) compileExpr(t *code.Builder, expr ast.Expr) error {
 	switch expr.Kind {
 	case ast.ExprKindBinary:
@@ -676,3 +712,10 @@ func (comp *Compiler) compileBlockNode(t *code.Builder, block ast.BlockNode) err
 
 	return nil
 }
+
+func (comp *Compiler) exitScopeAndCleanStack(t *code.Builder) {
+	if stackSpace := comp.scopes.Exit(); stackSpace != 0 {
+		t.AppendOp(code.OpDrop)
+		t.AppendInt(int64(stackSpace))
+	}
+}