- Fixed iterator stack leak when returning from the loop body. Fixes #148.

- Fixed stack leak when 'if' statement with value is last.

compiler.go

@@ -10,6 +10,7 @@ import (
 
 const (
 	blockLoop = iota
+	blockLoopEnum
 	blockTry
 	blockBranch
 	blockSwitch
@@ -85,7 +86,7 @@ func (c *compiler) leaveBlock() {
 	for _, item := range c.block.breaks {
 		c.p.code[item] = jump(lbl - item)
 	}
-	if c.block.typ == blockLoop {
+	if t := c.block.typ; t == blockLoop || t == blockLoopEnum {
 		for _, item := range c.block.conts {
 			c.p.code[item] = jump(c.block.cont - item)
 		}

compiler_stmt.go

@@ -304,7 +304,7 @@ func (c *compiler) compileForInStatement(v *ast.ForInStatement, needResult bool)
 
 func (c *compiler) compileLabeledForInStatement(v *ast.ForInStatement, needResult bool, label string) {
 	c.block = &block{
-		typ:        blockLoop,
+		typ:        blockLoopEnum,
 		outer:      c.block,
 		label:      label,
 		needResult: needResult,
@@ -421,7 +421,7 @@ func (c *compiler) findBranchBlock(st *ast.BranchStatement) *block {
 func (c *compiler) findContinueBlock(label *ast.Identifier) (block *block) {
 	if label != nil {
 		for b := c.block; b != nil; b = b.outer {
-			if b.typ == blockLoop && b.label == label.Name {
+			if (b.typ == blockLoop || b.typ == blockLoopEnum) && b.label == label.Name {
 				block = b
 				break
 			}
@@ -429,7 +429,7 @@ func (c *compiler) findContinueBlock(label *ast.Identifier) (block *block) {
 	} else {
 		// find the nearest loop
 		for b := c.block; b != nil; b = b.outer {
-			if b.typ == blockLoop {
+			if b.typ == blockLoop || b.typ == blockLoopEnum {
 				block = b
 				break
 			}
@@ -452,7 +452,7 @@ func (c *compiler) findBreakBlock(label *ast.Identifier) (block *block) {
 	L:
 		for b := c.block; b != nil; b = b.outer {
 			switch b.typ {
-			case blockLoop, blockSwitch:
+			case blockLoop, blockLoopEnum, blockSwitch:
 				block = b
 				break L
 			}
@@ -486,7 +486,7 @@ func (c *compiler) compileBreak(label *ast.Identifier, idx file.Idx) {
 				c.emit(halt)
 			case blockWith:
 				c.emit(leaveWith)
-			case blockLoop, blockSwitch:
+			case blockLoop, blockLoopEnum, blockSwitch:
 				block = b
 				break L
 			}
@@ -510,7 +510,7 @@ func (c *compiler) compileContinue(label *ast.Identifier, idx file.Idx) {
 		for b := c.block; b != nil; b = b.outer {
 			if b.typ == blockTry {
 				c.emit(halt)
-			} else if b.typ == blockLoop && b.label == label.Name {
+			} else if (b.typ == blockLoop || b.typ == blockLoopEnum) && b.label == label.Name {
 				block = b
 				break
 			}
@@ -520,7 +520,7 @@ func (c *compiler) compileContinue(label *ast.Identifier, idx file.Idx) {
 		for b := c.block; b != nil; b = b.outer {
 			if b.typ == blockTry {
 				c.emit(halt)
-			} else if b.typ == blockLoop {
+			} else if b.typ == blockLoop || b.typ == blockLoopEnum {
 				block = b
 				break
 			}
@@ -587,10 +587,14 @@ func (c *compiler) compileIfStatement(v *ast.IfStatement, needResult bool) {
 		c.p.code[jmp1] = jump(len(c.p.code) - jmp1)
 		c.markBlockStart()
 	} else {
-		c.p.code[jmp] = jne(len(c.p.code) - jmp)
-		c.markBlockStart()
 		if needResult {
+			c.emit(jump(2))
+			c.p.code[jmp] = jne(len(c.p.code) - jmp)
 			c.emit(loadUndef)
+			c.markBlockStart()
+		} else {
+			c.p.code[jmp] = jne(len(c.p.code) - jmp)
+			c.markBlockStart()
 		}
 	}
 }
@@ -603,8 +607,11 @@ func (c *compiler) compileReturnStatement(v *ast.ReturnStatement) {
 		c.emit(loadUndef)
 	}
 	for b := c.block; b != nil; b = b.outer {
-		if b.typ == blockTry {
+		switch b.typ {
+		case blockTry:
 			c.emit(halt)
+		case blockLoopEnum:
+			c.emit(enumPop)
 		}
 	}
 	c.emit(ret)

compiler_test.go

@@ -67,6 +67,10 @@ func testScript1(script string, expectedResult Value, t *testing.T) {
 	if vm.sp != 0 {
 		t.Fatalf("sp: %d", vm.sp)
 	}
+
+	if l := len(vm.iterStack); l > 0 {
+		t.Fatalf("iter stack is not empty: %d", l)
+	}
 }
 
 func TestEmptyProgram(t *testing.T) {
@@ -1968,6 +1972,27 @@ func TestEmptyCodeError(t *testing.T) {
 	}
 }
 
+func TestReturnFromForInLoop(t *testing.T) {
+	const SCRIPT = `
+	(function f() {
+		for (var i in {a: 1}) {
+			return true;
+		}
+	})();
+	`
+	testScript1(SCRIPT, valueTrue, t)
+}
+
+func TestIfStackLeaks(t *testing.T) {
+	const SCRIPT = `
+	var t = 0;
+	if (t === 0) {
+		t;
+	}
+	`
+	testScript1(SCRIPT, _positiveZero, t)
+}
+
 // FIXME
 /*
 func TestDummyCompile(t *testing.T) {

tc39_test.go

@@ -177,6 +177,14 @@ func (ctx *tc39TestCtx) runTC39Test(name, src string, meta *tc39Meta, t testing.
 			t.Fatalf("%s: Expected error: %v", name, err)
 		}
 	}
+
+	if vm.vm.sp != 0 {
+		t.Fatalf("sp: %d", vm.vm.sp)
+	}
+
+	if l := len(vm.vm.iterStack); l > 0 {
+		t.Fatalf("iter stack is not empty: %d", l)
+	}
 }
 
 func (ctx *tc39TestCtx) runTC39File(name string, t testing.TB) {