resolve: retain globals across REPL chunks (#247)

Previously, in the REPL, the globals bound by one chunk would become
the "predeclared" bindings of the next chunk. However, this produced
the wrong result for [x = x + 1] because both occurrences of x would
resolve to the same global x, which was undefined when the right-hand
side expression was evaluated, leading to a dynamic error. It would
fail for [x += 1] for a similar reason.

This change extends the resolver to allow clients to provide it with a
non-empty set of module globals. Thus the globals of one chunk remain
the globals of the next chunk, and the predeclared set is always empty.

The new API functions is intended for use only by the REPL.

Fixes #233

4 files changed, 104 insertions, 26 deletions

diff --git a/repl/repl.go b/repl/repl.go
index d766eff..5f2ee36 100644
--- a/repl/repl.go
+++ b/repl/repl.go
@@ -133,26 +133,9 @@ func rep(rl *readline.Instance, thread *starlark.Thread, globals starlark.String
 		if v != starlark.None {
 			fmt.Println(v)
 		}
-	} else {
-		// compile
-		prog, err := starlark.FileProgram(f, globals.Has)
-		if err != nil {
-			PrintError(err)
-			return nil
-		}
-
-		// execute (but do not freeze)
-		res, err := prog.Init(thread, globals)
-		if err != nil {
-			PrintError(err)
-		}
-
-		// The global names from the previous call become
-		// the predeclared names of this call.
-		// If execution failed, some globals may be undefined.
-		for k, v := range res {
-			globals[k] = v
-		}
+	} else if err := starlark.ExecREPLChunk(f, thread, globals); err != nil {
+		PrintError(err)
+		return nil
 	}
 
 	return nil
diff --git a/resolve/resolve.go b/resolve/resolve.go
index 13ddfe6..440bcf0 100644
--- a/resolve/resolve.go
+++ b/resolve/resolve.go
@@ -122,7 +122,13 @@ var (
 // dependency upon starlark.Universe, not because users should ever need
 // to redefine it.
 func File(file *syntax.File, isPredeclared, isUniversal func(name string) bool) error {
-	r := newResolver(isPredeclared, isUniversal)
+	return REPLChunk(file, nil, isPredeclared, isUniversal)
+}
+
+// REPLChunk is a generalization of the File function that supports a
+// non-empty initial global block, as occurs in a REPL.
+func REPLChunk(file *syntax.File, isGlobal, isPredeclared, isUniversal func(name string) bool) error {
+	r := newResolver(isGlobal, isPredeclared, isUniversal)
 	r.stmts(file.Stmts)
 
 	r.env.resolveLocalUses()
@@ -148,7 +154,7 @@ func File(file *syntax.File, isPredeclared, isUniversal func(name string) bool)
 //
 // The isPredeclared and isUniversal predicates behave as for the File function.
 func Expr(expr syntax.Expr, isPredeclared, isUniversal func(name string) bool) ([]*Binding, error) {
-	r := newResolver(isPredeclared, isUniversal)
+	r := newResolver(nil, isPredeclared, isUniversal)
 	r.expr(expr)
 	r.env.resolveLocalUses()
 	r.resolveNonLocalUses(r.env) // globals & universals
@@ -171,11 +177,12 @@ type Error struct {
 
 func (e Error) Error() string { return e.Pos.String() + ": " + e.Msg }
 
-func newResolver(isPredeclared, isUniversal func(name string) bool) *resolver {
+func newResolver(isGlobal, isPredeclared, isUniversal func(name string) bool) *resolver {
 	file := new(block)
 	return &resolver{
 		file:          file,
 		env:           file,
+		isGlobal:      isGlobal,
 		isPredeclared: isPredeclared,
 		isUniversal:   isUniversal,
 		globals:       make(map[string]*Binding),
@@ -202,8 +209,10 @@ type resolver struct {
 	predeclared map[string]*Binding
 
 	// These predicates report whether a name is
-	// pre-declared, either in this module or universally.
-	isPredeclared, isUniversal func(name string) bool
+	// pre-declared, either in this module or universally,
+	// or already declared in the module globals (as in a REPL).
+	// isGlobal may be nil.
+	isGlobal, isPredeclared, isUniversal func(name string) bool
 
 	loops int // number of enclosing for loops
 
@@ -390,6 +399,15 @@ func (r *resolver) useToplevel(use use) (bind *Binding) {
 	} else if prev, ok := r.globals[id.Name]; ok {
 		// use of global declared by module
 		bind = prev
+	} else if r.isGlobal != nil && r.isGlobal(id.Name) {
+		// use of global defined in a previous REPL chunk
+		bind = &Binding{
+			First: id, // wrong: this is not even a binding use
+			Scope: Global,
+			Index: len(r.moduleGlobals),
+		}
+		r.globals[id.Name] = bind
+		r.moduleGlobals = append(r.moduleGlobals, bind)
 	} else if prev, ok := r.predeclared[id.Name]; ok {
 		// repeated use of predeclared or universal
 		bind = prev
@@ -433,7 +451,8 @@ func (r *resolver) spellcheck(use use) string {
 
 	// globals
 	//
-	// We have no way to enumerate predeclared/universe,
+	// We have no way to enumerate the sets whose membership
+	// tests are isPredeclared, isUniverse, and isGlobal,
 	// which includes prior names in the REPL session.
 	for _, bind := range r.moduleGlobals {
 		names = append(names, bind.First.Name)
diff --git a/starlark/eval.go b/starlark/eval.go
index be70c0f..de492ca 100644
--- a/starlark/eval.go
+++ b/starlark/eval.go
@@ -366,6 +366,57 @@ func (prog *Program) Init(thread *Thread, predeclared StringDict) (StringDict, e
 	return toplevel.Globals(), err
 }
 
+// ExecREPLChunk compiles and executes file f in the specified thread
+// and global environment. This is a variant of ExecFile specialized to
+// the needs of a REPL, in which a sequence of input chunks, each
+// syntactically a File, manipulates the same set of module globals,
+// which are not frozen after execution.
+//
+// This function is intended to support only go.starlark.net/repl.
+// Its API stability is not guaranteed.
+func ExecREPLChunk(f *syntax.File, thread *Thread, globals StringDict) error {
+	var predeclared StringDict
+
+	// -- variant of FileProgram --
+
+	if err := resolve.REPLChunk(f, globals.Has, predeclared.Has, Universe.Has); err != nil {
+		return err
+	}
+
+	var pos syntax.Position
+	if len(f.Stmts) > 0 {
+		pos = syntax.Start(f.Stmts[0])
+	} else {
+		pos = syntax.MakePosition(&f.Path, 1, 1)
+	}
+
+	module := f.Module.(*resolve.Module)
+	compiled := compile.File(f.Stmts, pos, "<toplevel>", module.Locals, module.Globals)
+	prog := &Program{compiled}
+
+	// -- variant of Program.Init --
+
+	toplevel := makeToplevelFunction(prog.compiled, predeclared)
+
+	// Initialize module globals from parameter.
+	for i, id := range prog.compiled.Globals {
+		if v := globals[id.Name]; v != nil {
+			toplevel.module.globals[i] = v
+		}
+	}
+
+	_, err := Call(thread, toplevel, nil, nil)
+
+	// Reflect changes to globals back to parameter, even after an error.
+	for i, id := range prog.compiled.Globals {
+		if v := toplevel.module.globals[i]; v != nil {
+			globals[id.Name] = v
+		}
+	}
+
+	return err
+}
+
 func makeToplevelFunction(prog *compile.Program, predeclared StringDict) *Function {
 	// Create the Starlark value denoted by each program constant c.
 	constants := make([]Value, len(prog.Constants))
diff --git a/starlark/eval_test.go b/starlark/eval_test.go
index b3ec02e..f5a667a 100644
--- a/starlark/eval_test.go
+++ b/starlark/eval_test.go
@@ -763,3 +763,28 @@ func TestUnpackErrorBadType(t *testing.T) {
 		}
 	}
 }
+
+// Regression test for github.com/google/starlark-go/issues/233.
+func TestREPLChunk(t *testing.T) {
+	thread := new(starlark.Thread)
+	globals := make(starlark.StringDict)
+	exec := func(src string) {
+		f, err := syntax.Parse("<repl>", src, 0)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if err := starlark.ExecREPLChunk(f, thread, globals); err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	exec("x = 0; y = 0")
+	if got, want := fmt.Sprintf("%v %v", globals["x"], globals["y"]), "0 0"; got != want {
+		t.Fatalf("chunk1: got %s, want %s", got, want)
+	}
+
+	exec("x += 1; y = y + 1")
+	if got, want := fmt.Sprintf("%v %v", globals["x"], globals["y"]), "1 1"; got != want {
+		t.Fatalf("chunk2: got %s, want %s", got, want)
+	}
+}