From d0735257ced31afce1bf8c01ae02d5a5047892ee Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Tue, 19 May 2020 18:07:29 +0300
Subject: [PATCH] compiler: support range loops with value variables

Closes #958.
---
 pkg/compiler/codegen.go    | 12 +++++-------
 pkg/compiler/for_test.go   | 18 +++++++-----------
 pkg/compiler/func_scope.go |  9 +++++++++
 3 files changed, 21 insertions(+), 18 deletions(-)

diff --git a/pkg/compiler/codegen.go b/pkg/compiler/codegen.go
index 68130a89b..7f7c2758b 100644
--- a/pkg/compiler/codegen.go
+++ b/pkg/compiler/codegen.go
@@ -883,13 +883,6 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 		return nil
 
 	case *ast.RangeStmt:
-		// currently only simple for-range loops are supported
-		// for i := range ...
-		if n.Value != nil {
-			c.prog.Err = errors.New("range loops with value variable are not supported")
-			return nil
-		}
-
 		start, label := c.generateLabel(labelStart)
 		end := c.newNamedLabel(labelEnd, label)
 		post := c.newNamedLabel(labelPost, label)
@@ -914,6 +907,11 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 			emit.Syscall(c.prog.BinWriter, "Neo.Iterator.Key")
 			c.emitStoreVar(n.Key.(*ast.Ident).Name)
 		}
+		if n.Value != nil {
+			emit.Opcode(c.prog.BinWriter, opcode.DUP)
+			emit.Syscall(c.prog.BinWriter, "Neo.Enumerator.Value")
+			c.emitStoreVar(n.Value.(*ast.Ident).Name)
+		}
 
 		ast.Walk(c, n.Body)
 
diff --git a/pkg/compiler/for_test.go b/pkg/compiler/for_test.go
index 11d60c507..31515951f 100644
--- a/pkg/compiler/for_test.go
+++ b/pkg/compiler/for_test.go
@@ -3,13 +3,9 @@ package compiler_test
 import (
 	"fmt"
 	"math/big"
-	"strings"
 	"testing"
 
-	"github.com/nspcc-dev/neo-go/pkg/compiler"
-
 	"github.com/nspcc-dev/neo-go/pkg/vm"
-	"github.com/stretchr/testify/require"
 )
 
 func TestEntryPointWithMethod(t *testing.T) {
@@ -707,20 +703,20 @@ func TestForLoopRangeNoVariable(t *testing.T) {
 	eval(t, src, big.NewInt(3))
 }
 
-func TestForLoopRangeCompilerError(t *testing.T) {
+func TestForLoopRangeValue(t *testing.T) {
 	src := `
 	package foo
-	func f(a int) int { return 0 }
+	func f(a int) int { return a }
 	func Main() int {
-		arr := []int{1, 2, 3}
+		var sum int
+		arr := []int{1, 9, 4}
 		for _, v := range arr {
-			f(v)
+			sum += f(v)
 		}
-		return 0
+		return sum
 	}`
 
-	_, err := compiler.Compile(strings.NewReader(src))
-	require.Error(t, err)
+	eval(t, src, big.NewInt(14))
 }
 
 func TestForLoopComplexConditions(t *testing.T) {
diff --git a/pkg/compiler/func_scope.go b/pkg/compiler/func_scope.go
index e5bb6e084..fde6ceadf 100644
--- a/pkg/compiler/func_scope.go
+++ b/pkg/compiler/func_scope.go
@@ -104,6 +104,15 @@ func (c *funcScope) countLocals() int {
 		// This handles the inline GenDecl like "var x = 2"
 		case *ast.ValueSpec:
 			size += len(n.Names)
+		case *ast.RangeStmt:
+			if n.Tok == token.DEFINE {
+				if n.Key != nil {
+					size++
+				}
+				if n.Value != nil {
+					size++
+				}
+			}
 		}
 		return true
 	})