compile,py: fix closure and decorator

Author: wetor

compile/compile.go

@@ -213,13 +213,17 @@ func (c *compiler) compileAst(Ast ast.Ast, filename string, futureFlags int, don
 	case *ast.Suite:
 		panic("suite should not be possible")
 	case *ast.Lambda:
+		code.Argcount = int32(len(node.Args.Args))
+		code.Kwonlyargcount = int32(len(node.Args.Kwonlyargs))
 		// Make None the first constant as lambda can't have a docstring
 		c.Const(py.None)
 		code.Name = "<lambda>"
 		c.setQualname()
 		c.Expr(node.Body)
 		valueOnStack = true
 	case *ast.FunctionDef:
+		code.Argcount = int32(len(node.Args.Args))
+		code.Kwonlyargcount = int32(len(node.Args.Kwonlyargs))
 		code.Name = string(node.Name)
 		c.setQualname()
 		c.Stmts(c.docString(node.Body, true))
@@ -299,6 +303,7 @@ func (c *compiler) compileAst(Ast ast.Ast, filename string, futureFlags int, don
 	code.Stacksize = int32(c.OpCodes.StackDepth())
 	code.Nlocals = int32(len(code.Varnames))
 	code.Lnotab = string(c.OpCodes.Lnotab())
+	code.InitCell2arg()
 	return nil
 }
 
@@ -479,7 +484,8 @@ func (c *compiler) makeClosure(code *py.Code, args uint32, child *compiler, qual
 		if reftype == symtable.ScopeCell {
 			arg = c.FindId(name, c.Code.Cellvars)
 		} else { /* (reftype == FREE) */
-			arg = c.FindId(name, c.Code.Freevars)
+			// using CellAndFreeVars in closures requires skipping Cellvars
+			arg = len(c.Code.Cellvars) + c.FindId(name, c.Code.Freevars)
 		}
 		if arg < 0 {
 			panic(fmt.Sprintf("compile: makeClosure: lookup %q in %q %v %v\nfreevars of %q: %v\n", name, c.SymTable.Name, reftype, arg, code.Name, code.Freevars))
@@ -1363,7 +1369,12 @@ func (c *compiler) NameOp(name string, ctx ast.ExprContext) {
 	if op == 0 {
 		panic("NameOp: Op not set")
 	}
-	c.OpArg(op, c.Index(mangled, dict))
+	i := c.Index(mangled, dict)
+	// using CellAndFreeVars in closures requires skipping Cellvars
+	if scope == symtable.ScopeFree {
+		i += uint32(len(c.Code.Cellvars))
+	}
+	c.OpArg(op, i)
 }
 
 // Call a function which is already on the stack with n arguments already on the stack

py/code.go

@@ -112,8 +112,6 @@ func NewCode(argcount int32, kwonlyargcount int32,
 	filename_ Object, name_ Object, firstlineno int32,
 	lnotab_ Object) *Code {
 
-	var cell2arg []byte
-
 	// Type assert the objects
 	consts := consts_.(Tuple)
 	namesTuple := names_.(Tuple)
@@ -154,7 +152,6 @@ func NewCode(argcount int32, kwonlyargcount int32,
 	// 	return nil;
 	// }
 
-	n_cellvars := len(cellvars)
 	intern_strings(namesTuple)
 	intern_strings(varnamesTuple)
 	intern_strings(freevarsTuple)
@@ -167,13 +164,40 @@ func NewCode(argcount int32, kwonlyargcount int32,
 			}
 		}
 	}
+
+	co := &Code{
+		Argcount:       argcount,
+		Kwonlyargcount: kwonlyargcount,
+		Nlocals:        nlocals,
+		Stacksize:      stacksize,
+		Flags:          flags,
+		Code:           code,
+		Consts:         consts,
+		Names:          names,
+		Varnames:       varnames,
+		Freevars:       freevars,
+		Cellvars:       cellvars,
+		Filename:       filename,
+		Name:           name,
+		Firstlineno:    firstlineno,
+		Lnotab:         lnotab,
+		Weakreflist:    nil,
+	}
+	co.InitCell2arg()
+	return co
+}
+
+// Create mapping between cells and arguments if needed.
+func (co *Code) InitCell2arg() {
+	var cell2arg []byte
+	n_cellvars := len(co.Cellvars)
 	/* Create mapping between cells and arguments if needed. */
 	if n_cellvars != 0 {
-		total_args := argcount + kwonlyargcount
-		if flags&CO_VARARGS != 0 {
+		total_args := co.Argcount + co.Kwonlyargcount
+		if co.Flags&CO_VARARGS != 0 {
 			total_args++
 		}
-		if flags&CO_VARKEYWORDS != 0 {
+		if co.Flags&CO_VARKEYWORDS != 0 {
 			total_args++
 		}
 		used_cell2arg := false
@@ -182,9 +206,9 @@ func NewCode(argcount int32, kwonlyargcount int32,
 			cell2arg[i] = CO_CELL_NOT_AN_ARG
 		}
 		// Find cells which are also arguments.
-		for i, cell := range cellvars {
+		for i, cell := range co.Cellvars {
 			for j := int32(0); j < total_args; j++ {
-				arg := varnames[j]
+				arg := co.Varnames[j]
 				if cell == arg {
 					cell2arg[i] = byte(j)
 					used_cell2arg = true
@@ -196,26 +220,7 @@ func NewCode(argcount int32, kwonlyargcount int32,
 			cell2arg = nil
 		}
 	}
-
-	return &Code{
-		Argcount:       argcount,
-		Kwonlyargcount: kwonlyargcount,
-		Nlocals:        nlocals,
-		Stacksize:      stacksize,
-		Flags:          flags,
-		Code:           code,
-		Consts:         consts,
-		Names:          names,
-		Varnames:       varnames,
-		Freevars:       freevars,
-		Cellvars:       cellvars,
-		Cell2arg:       cell2arg,
-		Filename:       filename,
-		Name:           name,
-		Firstlineno:    firstlineno,
-		Lnotab:         lnotab,
-		Weakreflist:    nil,
-	}
+	co.Cell2arg = cell2arg
 }
 
 // Return number of free variables

vm/tests/class.py

@@ -47,17 +47,16 @@ def method1(self, x):
 c = x()
 assert c.method1(1) == 2
 
-# FIXME doesn't work
-# doc="CLASS_DEREF2"
-# def classderef2(x):
-#     class DeRefTest:
-#         VAR = x
-#         def method1(self, x):
-#             "method1"
-#             return self.VAR+x
-#     return DeRefTest
-# x = classderef2(1)
-# c = x()
-# assert c.method1(1) == 2
+doc="CLASS_DEREF2"
+def classderef2(x):
+    class DeRefTest:
+        VAR = x
+        def method1(self, x):
+            "method1"
+            return self.VAR+x
+    return DeRefTest
+x = classderef2(1)
+c = x()
+assert c.method1(1) == 2
 
 doc="finished"