feat: make Next iterator instruction faster and more efficient

Branching control is delegated directly to the Next instruction,
which now takes the location of loop exit as first argument.
It avoids the use of JumpFalse, plus the stack storage for the
condition.

5 files changed, 28 insertions, 15 deletions

diff --git a/comp/compiler.go b/comp/compiler.go
index 05a4b3f..ef379ea 100644
--- a/comp/compiler.go
+++ b/comp/compiler.go
@@ -408,8 +408,15 @@ func (c *Compiler) Generate(tokens parser.Tokens) (err error) {
 			}
 
 		case lang.Next:
+			var i int
+			if s, ok := c.Symbols[t.Str]; !ok {
+				t.Arg = []any{len(c.Code)} // current code location
+				fixList = append(fixList, t)
+			} else {
+				i = int(s.Value.Int()) - len(c.Code)
+			}
 			k := stack[len(stack)-2]
-			c.emit(t, vm.Next, k.Index)
+			c.emit(t, vm.Next, i, k.Index)
 
 		case lang.Range:
 			// FIXME: handle all iterator types.
diff --git a/interp/interpreter_test.go b/interp/interpreter_test.go
index 89ab4ae..ea4da53 100644
--- a/interp/interpreter_test.go
+++ b/interp/interpreter_test.go
@@ -128,7 +128,7 @@ func TestFor(t *testing.T) {
 		{src: "func f() int {a := 0; for {a = a+1; if a < 3 {continue}; break}; return a}; f()", res: "3"},       // #04
 		{src: "a := []int{1,2,3,4}; b := 0; for i := range a {b = b+i}; b", res: "6"},                            // #05
 		{src: "func f() int {a := []int{1,2,3,4}; b := 0; for i := range a {b = b+i}; return b}; f()", res: "6"}, // #06
-		// {src: "a := []int{1,2,3,4}; b := 0; for i, e := range a {b = b+i+e}; b", res: "6"},                       // #07
+		// {src: "a := []int{1,2,3,4}; b := 0; for i, e := range a {b = b+i+e}; b", res: "16"},                       // #07
 	})
 }
 
diff --git a/parser/parse.go b/parser/parse.go
index daf30ca..e72e16e 100644
--- a/parser/parse.go
+++ b/parser/parse.go
@@ -201,6 +201,7 @@ func (p *Parser) parseGoto(in Tokens) (out Tokens, err error) {
 func (p *Parser) parseFor(in Tokens) (out Tokens, err error) {
 	// TODO: detect invalid code.
 	var init, cond, post, body, final Tokens
+	hasRange := in.Index(lang.Range) >= 0
 	fc := strconv.Itoa(p.labelCount[p.scope])
 	p.labelCount[p.scope]++
 	breakLabel, continueLabel := p.breakLabel, p.continueLabel
@@ -213,10 +214,10 @@ func (p *Parser) parseFor(in Tokens) (out Tokens, err error) {
 	pre := in[1 : len(in)-1].Split(lang.Semicolon)
 	switch len(pre) {
 	case 1:
-		if in.Index(lang.Range) >= 0 {
+		if hasRange {
 			init = pre[0]
-			cond = Tokens{{Token: scanner.Token{Tok: lang.Next}}}
-			final = Tokens{{Token: scanner.Token{Tok: lang.Stop}}}
+			cond = Tokens{newNext(p.breakLabel, in[1].Pos)}
+			final = Tokens{newStop(in[1].Pos)}
 		} else {
 			cond = pre[0]
 		}
@@ -231,13 +232,15 @@ func (p *Parser) parseFor(in Tokens) (out Tokens, err error) {
 		}
 		out = init
 	}
-	out = append(out, newLabel(p.scope+"b", in[0].Pos))
+	out = append(out, newLabel(p.continueLabel, in[0].Pos))
 	if len(cond) > 0 {
 		if cond, err = p.parseExpr(cond, ""); err != nil {
 			return nil, err
 		}
 		out = append(out, cond...)
-		out = append(out, newJumpFalse(p.scope+"e", in[0].Pos))
+		if !hasRange {
+			out = append(out, newJumpFalse(p.breakLabel, in[0].Pos))
+		}
 	}
 	if body, err = p.Parse(in[len(in)-1].Block()); err != nil {
 		return nil, err
@@ -250,8 +253,8 @@ func (p *Parser) parseFor(in Tokens) (out Tokens, err error) {
 		out = append(out, post...)
 	}
 	out = append(out,
-		newGoto(p.scope+"b", in[0].Pos),
-		newLabel(p.scope+"e", in[0].Pos),
+		newGoto(p.continueLabel, in[0].Pos),
+		newLabel(p.breakLabel, in[0].Pos),
 	)
 	out = append(out, final...)
 	return out, err
diff --git a/parser/tokens.go b/parser/tokens.go
index a0a93fa..7a241f3 100644
--- a/parser/tokens.go
+++ b/parser/tokens.go
@@ -79,6 +79,8 @@ func newCall(pos int, arg ...any) Token               { return newToken(lang.Cal
 func newGoto(label string, pos int) Token             { return newToken(lang.Goto, label, pos) }
 func newLabel(label string, pos int) Token            { return newToken(lang.Label, label, pos) }
 func newJumpFalse(label string, pos int) Token        { return newToken(lang.JumpFalse, label, pos) }
+func newNext(label string, pos int) Token             { return newToken(lang.Next, label, pos) }
+func newStop(pos int) Token                           { return newToken(lang.Stop, "", pos) }
 func newGrow(size, pos int) Token                     { return newToken(lang.Grow, "", pos, size) }
 func newSemicolon(pos int) Token                      { return newToken(lang.Semicolon, "", pos) }
 func newEqualSet(pos int) Token                       { return newToken(lang.EqualSet, "", pos) }
diff --git a/vm/vm.go b/vm/vm.go
index f6dc6fc..7ab5a25 100644
--- a/vm/vm.go
+++ b/vm/vm.go
@@ -108,8 +108,8 @@ func (m *Machine) Run() (err error) {
 	defer func() { m.mem, m.ip, m.fp, m.ic = mem, ip, fp, ic }()
 
 	for {
-		sp = len(mem) // stack pointer
-		c := m.code[ip]
+		sp = len(mem)   // stack pointer
+		c := m.code[ip] // current instruction
 		if debug {
 			log.Printf("ip:%-3d sp:%-3d fp:%-3d op:[%-20v] mem:%v\n", ip, sp, fp, c, Vstring(mem))
 		}
@@ -247,11 +247,12 @@ func (m *Machine) Run() (err error) {
 		case Negate:
 			mem[sp-1] = ValueOf(-mem[sp-1].Int())
 		case Next:
-			v, ok := mem[sp-2].Interface().(func() (reflect.Value, bool))()
-			if ok {
-				mem[c.Arg[0]].Set(v)
+			if v, ok := mem[sp-2].Interface().(func() (reflect.Value, bool))(); ok {
+				mem[c.Arg[1]].Set(v)
+			} else {
+				ip += c.Arg[0]
+				continue
 			}
-			mem = append(mem, ValueOf(ok))
 		case Not:
 			mem[sp-1] = ValueOf(!mem[sp-1].Bool())
 		case Pop: