Merge pull request #658 from nspcc-dev/feature/range

compiler: support for-range loops

pkg/compiler/codegen.go

@@ -660,6 +660,45 @@ 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 := c.newLabel()
+		end := c.newLabel()
+
+		ast.Walk(c, n.X)
+
+		emit.Opcode(c.prog.BinWriter, opcode.ARRAYSIZE)
+		emit.Opcode(c.prog.BinWriter, opcode.PUSH0)
+
+		c.setLabel(start)
+
+		emit.Opcode(c.prog.BinWriter, opcode.OVER)
+		emit.Opcode(c.prog.BinWriter, opcode.OVER)
+		emit.Opcode(c.prog.BinWriter, opcode.LTE) // finish if len <= i
+		emit.Jmp(c.prog.BinWriter, opcode.JMPIF, int16(end))
+
+		if n.Key != nil {
+			emit.Opcode(c.prog.BinWriter, opcode.DUP)
+
+			pos := c.scope.loadLocal(n.Key.(*ast.Ident).Name)
+			c.emitStoreLocal(pos)
+		}
+
+		ast.Walk(c, n.Body)
+
+		emit.Opcode(c.prog.BinWriter, opcode.INC)
+		emit.Jmp(c.prog.BinWriter, opcode.JMP, int16(start))
+
+		c.setLabel(end)
+
+		return nil
+
 	// We dont really care about assertions for the core logic.
 	// The only thing we need is to please the compiler type checking.
 	// For this to work properly, we only need to walk the expression

pkg/compiler/for_test.go

@@ -2,9 +2,13 @@ package compiler_test
 
 import (
 	"math/big"
+	"strings"
 	"testing"
 
+	"github.com/CityOfZion/neo-go/pkg/compiler"
+
 	"github.com/CityOfZion/neo-go/pkg/vm"
+	"github.com/stretchr/testify/require"
 )
 
 func TestEntryPointWithMethod(t *testing.T) {
@@ -397,3 +401,82 @@ func TestForLoopNoPost(t *testing.T) {
 	`
 	eval(t, src, big.NewInt(10))
 }
+
+func TestForLoopRange(t *testing.T) {
+	src := `
+	package foo
+	func Main() int {
+		sum := 0
+		arr := []int{1, 2, 3}
+		for i := range arr {
+			sum += arr[i] 
+		}
+		return sum
+	}`
+
+	eval(t, src, big.NewInt(6))
+}
+
+func TestForLoopRangeGlobalIndex(t *testing.T) {
+	src := `
+	package foo
+	func Main() int {
+		sum := 0
+		i := 0
+		arr := []int{1, 2, 3}
+		for i = range arr {
+			sum += arr[i] 
+		}
+		return sum + i
+	}`
+
+	eval(t, src, big.NewInt(8))
+}
+
+func TestForLoopRangeChangeVariable(t *testing.T) {
+	src := `
+	package foo
+	func Main() int {
+		sum := 0
+		arr := []int{1, 2, 3}
+		for i := range arr {
+			sum += arr[i]
+			i++
+			sum += i
+		}
+		return sum
+	}`
+
+	eval(t, src, big.NewInt(12))
+}
+
+func TestForLoopRangeNoVariable(t *testing.T) {
+	src := `
+	package foo
+	func Main() int {
+		sum := 0
+		arr := []int{1, 2, 3}
+		for range arr {
+			sum += 1
+		}
+		return sum
+	}`
+
+	eval(t, src, big.NewInt(3))
+}
+
+func TestForLoopRangeCompilerError(t *testing.T) {
+	src := `
+	package foo
+	func f(a int) int { return 0 }
+	func Main() int {
+		arr := []int{1, 2, 3}
+		for _, v := range arr {
+			f(v)
+		}
+		return 0
+	}`
+
+	_, err := compiler.Compile(strings.NewReader(src))
+	require.Error(t, err)
+}