From 310f7449d7191be7233e117fed394c9705cbd0a8 Mon Sep 17 00:00:00 2001 From: Evgenii Stratonikov Date: Tue, 11 Feb 2020 15:26:11 +0300 Subject: [PATCH] compiler: support for-range loops --- pkg/compiler/codegen.go | 39 +++++++++++++++++++ pkg/compiler/for_test.go | 83 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 122 insertions(+) diff --git a/pkg/compiler/codegen.go b/pkg/compiler/codegen.go index 9a9b3aa61..abbdd6b34 100644 --- a/pkg/compiler/codegen.go +++ b/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 diff --git a/pkg/compiler/for_test.go b/pkg/compiler/for_test.go index 83b460517..c95d37a00 100644 --- a/pkg/compiler/for_test.go +++ b/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) +}