Fix Free and Cell variables to make closures work properly

Author: ncw

py/frame.go

@@ -12,12 +12,14 @@ type TryBlock struct {
 // A python Frame object
 type Frame struct {
 	// Back       *Frame        // previous frame, or nil
-	Code     *Code      // code segment
-	Builtins StringDict // builtin symbol table
-	Globals  StringDict // global symbol table
-	Locals   StringDict // local symbol table
-	Stack    []Object   // Valuestack
-	Closure  Tuple      // Tuple of Cells that this function is closed over
+	Code            *Code      // code segment
+	Builtins        StringDict // builtin symbol table
+	Globals         StringDict // global symbol table
+	Locals          StringDict // local symbol table
+	Stack           []Object   // Valuestack
+	LocalVars       Tuple      // Fast access local vars
+	CellAndFreeVars Tuple      // Cellvars then Freevars Cell objects in one Tuple
+
 	// Next free slot in f_valuestack.  Frame creation sets to f_valuestack.
 	// Frame evaluation usually NULLs it, but a frame that yields sets it
 	// to the current stack top.
@@ -62,13 +64,26 @@ func (o *Frame) Type() *Type {
 
 // Make a new frame for a code object
 func NewFrame(globals, locals StringDict, code *Code, closure Tuple) *Frame {
+	nlocals := int(code.Nlocals)
+	ncells := len(code.Cellvars)
+	nfrees := len(code.Freevars)
+	varsize := nlocals + ncells + nfrees
+	size := varsize + int(code.Stacksize)
+	// Allocate the stack, locals, cells and frees in a contigious block of memory
+	allocation := make([]Object, varsize, size)
+	localVars := allocation[0:nlocals]
+	//cellVars := allocation[nlocals : nlocals+ncells]
+	//freeVars := allocation[nlocals+ncells : varsize]
+	cellAndFreeVars := allocation[nlocals:varsize]
+
 	return &Frame{
-		Globals:  globals,
-		Locals:   locals,
-		Code:     code,
-		Closure:  closure,
-		Builtins: Builtins.Globals,
-		Stack:    make([]Object, 0, code.Stacksize),
+		Globals:         globals,
+		Locals:          locals,
+		Code:            code,
+		LocalVars:       localVars,
+		CellAndFreeVars: cellAndFreeVars,
+		Builtins:        Builtins.Globals,
+		Stack:           make([]Object, 0, code.Stacksize),
 	}
 }
 

vm/eval.go

@@ -1,6 +1,9 @@
 // Evaluate opcodes
 package vm
 
+// FIXME use LocalVars instead of storing everything in the Locals dict
+// see frameobject.c dict_to_map and LocalsToFast
+
 /* FIXME
 
 cpython has one stack per frame, not one stack in total
@@ -1055,18 +1058,15 @@ func _var_name(vm *Vm, i int32) (string, bool) {
 // co_freevars[i - len(co_cellvars)].
 func do_LOAD_CLOSURE(vm *Vm, i int32) {
 	defer vm.CheckException()
-	varname, _ := _var_name(vm, i)
-	// FIXME this is making a new cell each time rather than
-	// returning a reference to an old one
-	vm.PUSH(py.NewCell(vm.frame.Locals[varname]))
+	vm.PUSH(vm.frame.CellAndFreeVars[i])
 }
 
 // Loads the cell contained in slot i of the cell and free variable
 // storage. Pushes a reference to the object the cell contains on the
 // stack.
 func do_LOAD_DEREF(vm *Vm, i int32) {
 	defer vm.CheckException()
-	res := vm.frame.Closure[i].(*py.Cell).Get()
+	res := vm.frame.CellAndFreeVars[i].(*py.Cell).Get()
 	if res == nil {
 		varname, free := _var_name(vm, i)
 		if free {
@@ -1470,6 +1470,32 @@ func RunFrame(frame *py.Frame) (res py.Object, err error) {
 func Run(globals, locals py.StringDict, code *py.Code, closure py.Tuple) (res py.Object, err error) {
 	frame := py.NewFrame(globals, locals, code, closure)
 
+	// Allocate and initialize storage for cell vars, and copy free
+	// vars into frame.
+	for i := range code.Cellvars {
+		var c py.Object
+		// Possibly account for the cell variable being an argument.
+		if code.Cell2arg != nil && code.Cell2arg[i] != py.CO_CELL_NOT_AN_ARG {
+			arg := code.Cell2arg[i]
+			// FIXME if use localvars need to change this
+			// c = NewCell(GETLOCAL(arg))
+			// // Clear the local copy.
+			// SETLOCAL(arg, nil)
+			varname := code.Varnames[arg]
+			c = py.NewCell(locals[varname])
+		} else {
+			c = py.NewCell(nil)
+		}
+		//SETLOCAL(code.nlocals+i, c)
+		frame.CellAndFreeVars[i] = c
+	}
+	ncells := len(code.Cellvars)
+	for i := range code.Freevars {
+		//PyObject * o = PyTuple_GET_ITEM(closure, i)
+		//freevars[PyTuple_GET_SIZE(code.cellvars)+i] = o
+		frame.CellAndFreeVars[i+ncells] = closure[i]
+	}
+
 	// If this is a generator then make a generator object from
 	// the frame and return that instead
 	if code.Flags&py.CO_GENERATOR != 0 {