diff options
| -rw-r--r-- | internal/compile/compile.go | 2 | ||||
| -rw-r--r-- | starlark/eval_test.go | 41 | ||||
| -rw-r--r-- | starlark/testdata/recursion.star | 43 |
3 files changed, 55 insertions, 31 deletions
diff --git a/internal/compile/compile.go b/internal/compile/compile.go index 7ba17eb..0f3f264 100644 --- a/internal/compile/compile.go +++ b/internal/compile/compile.go @@ -1086,7 +1086,9 @@ func (fcomp *fcomp) stmt(stmt syntax.Stmt) { fcomp.ifelse(stmt.Cond, body, done) fcomp.block = body + fcomp.loops = append(fcomp.loops, loop{break_: done, continue_: head}) fcomp.stmts(stmt.Body) + fcomp.loops = fcomp.loops[:len(fcomp.loops)-1] fcomp.jump(head) fcomp.block = done diff --git a/starlark/eval_test.go b/starlark/eval_test.go index 09c6bfa..5c26e56 100644 --- a/starlark/eval_test.go +++ b/starlark/eval_test.go @@ -9,6 +9,7 @@ import ( "fmt" "math" "path/filepath" + "strings" "testing" "go.starlark.net/internal/chunkedfile" @@ -109,6 +110,7 @@ func TestExecFile(t *testing.T) { "testdata/set.star", "testdata/string.star", "testdata/tuple.star", + "testdata/recursion.star", } { filename := filepath.Join(testdata, file) for _, chunk := range chunkedfile.Read(filename, t) { @@ -116,6 +118,9 @@ func TestExecFile(t *testing.T) { "hasfields": starlark.NewBuiltin("hasfields", newHasFields), "fibonacci": fib{}, } + + resolve.AllowRecursion = option(chunk.Source, "recursion") + _, err := starlark.ExecFile(thread, filename, chunk.Source, predeclared) switch err := err.(type) { case *starlark.EvalError: @@ -136,11 +141,16 @@ func TestExecFile(t *testing.T) { default: t.Error(err) } + resolve.AllowRecursion = false chunk.Done() } } } +func option(chunk, name string) bool { + return strings.Contains(chunk, "option:"+name) +} + // A fib is an iterable value representing the infinite Fibonacci sequence. type fib struct{} @@ -473,34 +483,3 @@ def somefunc(): t.Fatal("docstring not found") } } - -func TestRecursion(t *testing.T) { - resolve.AllowRecursion = true - defer func() { resolve.AllowRecursion = false }() - - globals, err := starlark.ExecFile(&starlark.Thread{}, "rec.star", ` -def sum(n): - r = 0 - while n > 0: - r += n - n -= 1 - return r - -def fib(n): - if n <= 1: - return 1 - return fib(n-1) + fib(n-2) - -fib5 = fib(5) -sum5 = sum(5) -`, nil) - if err != nil { - t.Fatal(err) - } - if got, _ := globals["fib5"].(starlark.Int).Int64(); int(got) != 8 { - t.Fatalf("wrong value for fib5, got %d expected 8\n", got) - } - if got, _ := globals["sum5"].(starlark.Int).Int64(); int(got) != 5+4+3+2+1 { - t.Fatalf("wrong value for sum5, got %d expected %d\n", got, 5+4+3+2+1) - } -} diff --git a/starlark/testdata/recursion.star b/starlark/testdata/recursion.star new file mode 100644 index 0000000..3368614 --- /dev/null +++ b/starlark/testdata/recursion.star @@ -0,0 +1,43 @@ +# Tests of Starlark recursion and while statement. + +# This is a "chunked" file: each "---" effectively starts a new file. + +# option:recursion + +load("assert.star", "assert") + +def sum(n): + r = 0 + while n > 0: + r += n + n -= 1 + return r + +def fib(n): + if n <= 1: + return 1 + return fib(n-1) + fib(n-2) + +def while_break(n): + r = 0 + while n > 0: + if n == 5: + break + r += n + n -= 1 + return r + +def while_continue(n): + r = 0 + while n > 0: + if n % 2 == 0: + n -= 1 + continue + r += n + n -= 1 + return r + +assert.eq(fib(5), 8) +assert.eq(sum(5), 5+4+3+2+1) +assert.eq(while_break(10), 40) +assert.eq(while_continue(10), 25) |