compiler: support range loops with value variables

Closes #958.
This commit is contained in:
Evgenii Stratonikov 2020-05-19 18:07:29 +03:00
parent b126056f04
commit d0735257ce
3 changed files with 21 additions and 18 deletions

View file

@ -883,13 +883,6 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
return nil return nil
case *ast.RangeStmt: 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) start, label := c.generateLabel(labelStart)
end := c.newNamedLabel(labelEnd, label) end := c.newNamedLabel(labelEnd, label)
post := c.newNamedLabel(labelPost, 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") emit.Syscall(c.prog.BinWriter, "Neo.Iterator.Key")
c.emitStoreVar(n.Key.(*ast.Ident).Name) 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) ast.Walk(c, n.Body)

View file

@ -3,13 +3,9 @@ package compiler_test
import ( import (
"fmt" "fmt"
"math/big" "math/big"
"strings"
"testing" "testing"
"github.com/nspcc-dev/neo-go/pkg/compiler"
"github.com/nspcc-dev/neo-go/pkg/vm" "github.com/nspcc-dev/neo-go/pkg/vm"
"github.com/stretchr/testify/require"
) )
func TestEntryPointWithMethod(t *testing.T) { func TestEntryPointWithMethod(t *testing.T) {
@ -707,20 +703,20 @@ func TestForLoopRangeNoVariable(t *testing.T) {
eval(t, src, big.NewInt(3)) eval(t, src, big.NewInt(3))
} }
func TestForLoopRangeCompilerError(t *testing.T) { func TestForLoopRangeValue(t *testing.T) {
src := ` src := `
package foo package foo
func f(a int) int { return 0 } func f(a int) int { return a }
func Main() int { func Main() int {
arr := []int{1, 2, 3} var sum int
arr := []int{1, 9, 4}
for _, v := range arr { for _, v := range arr {
f(v) sum += f(v)
} }
return 0 return sum
}` }`
_, err := compiler.Compile(strings.NewReader(src)) eval(t, src, big.NewInt(14))
require.Error(t, err)
} }
func TestForLoopComplexConditions(t *testing.T) { func TestForLoopComplexConditions(t *testing.T) {

View file

@ -104,6 +104,15 @@ func (c *funcScope) countLocals() int {
// This handles the inline GenDecl like "var x = 2" // This handles the inline GenDecl like "var x = 2"
case *ast.ValueSpec: case *ast.ValueSpec:
size += len(n.Names) size += len(n.Names)
case *ast.RangeStmt:
if n.Tok == token.DEFINE {
if n.Key != nil {
size++
}
if n.Value != nil {
size++
}
}
} }
return true return true
}) })