Move ClosureTempAllocator to Code.py

Cython/Compiler/Code.py

@@ -1393,3 +1393,40 @@ class PyrexCodeWriter(object):
     def dedent(self):
         self.level -= 1
 
+
+class ClosureTempAllocator(object):
+    def __init__(self, klass=None):
+        self.klass = klass
+        self.temps_allocated = {}
+        self.temps_free = {}
+        self.temps_count = 0
+
+    def reset(self):
+        for type, cnames in self.temps_allocated.items():
+            self.temps_free[type] = list(cnames)
+
+    def allocate_temp(self, type):
+        if not type in self.temps_allocated:
+            self.temps_allocated[type] = []
+            self.temps_free[type] = []
+        elif self.temps_free[type]:
+            return self.temps_free[type].pop(0)
+        cname = '%s%d' % (Naming.codewriter_temp_prefix, self.temps_count)
+        self.klass.declare_var(pos=None, name=cname, cname=cname, type=type, is_cdef=True)
+        self.temps_allocated[type].append(cname)
+        self.temps_count += 1
+        return cname
+
+    def put_gotref(self, code):
+        for entry in self.klass.entries.values():
+            if entry.cname == Naming.outer_scope_cname: # XXX
+                continue
+            if entry.type.is_pyobject:
+                code.put_xgotref('%s->%s' % (Naming.cur_scope_cname, entry.cname))
+
+    def put_giveref(self, code):
+        for entry in self.klass.entries.values():
+            if entry.cname == Naming.outer_scope_cname: # XXX
+                continue
+            if entry.type.is_pyobject:
+                code.put_xgiveref('%s->%s' % (Naming.cur_scope_cname, entry.cname))

Cython/Compiler/ExprNodes.py

@@ -5005,15 +5005,14 @@ class YieldExprNode(ExprNode):
         else:
             code.put_init_to_py_none(Naming.retval_cname, py_object_type)
         saved = []
-        self.temp_allocator.reset()
+        code.temp_allocator.reset()
         code.putln('/* Save temporary variables */')
         for cname, type, manage_ref in code.funcstate.temps_in_use():
-            save_cname = self.temp_allocator.allocate_temp(type)
+            save_cname = code.temp_allocator.allocate_temp(type)
             saved.append((cname, save_cname, type))
             code.putln('%s->%s = %s;' % (Naming.cur_scope_cname, save_cname, cname))
 
-        # XXX: safe here as all used temps are handled but not clean
-        self.temp_allocator.put_giveref(code)
+        code.temp_allocator.put_giveref(code)
         code.put_xgiveref(Naming.retval_cname)
         code.put_finish_refcount_context()
         code.putln("/* return from function, yielding value */")

Cython/Compiler/Nodes.py

@@ -23,7 +23,7 @@ from PyrexTypes import py_object_type, error_type, CFuncType
 from Symtab import ModuleScope, LocalScope, ClosureScope, \
     StructOrUnionScope, PyClassScope, CClassScope, CppClassScope
 from Cython.Utils import open_new_file, replace_suffix
-from Code import UtilityCode
+from Code import UtilityCode, ClosureTempAllocator
 from StringEncoding import EncodedString, escape_byte_string, split_string_literal
 import Options
 import ControlFlow
@@ -1361,6 +1361,7 @@ class FuncDefNode(StatNode, BlockNode):
         if not self.is_generator:
             self.generate_preamble(env, code)
         if self.is_generator:
+            code.temp_allocator = ClosureTempAllocator(lenv.scope_class.type.scope)
             resume_code = code.insertion_point()
             first_run_label = code.new_label('first_run')
             code.use_label(first_run_label)
@@ -1524,7 +1525,7 @@ class FuncDefNode(StatNode, BlockNode):
 
         if self.is_generator:
             gotref_code.putln('/* Make refnanny happy */')
-            self.temp_allocator.put_gotref(gotref_code)
+            code.temp_allocator.put_gotref(gotref_code)
             self.generator.generate_function_body(self.local_scope, code)
 
     def generate_preamble(self, env, code):

Cython/Compiler/ParseTreeTransforms.py

@@ -1303,43 +1303,6 @@ class AlignFunctionDefinitions(CythonTransform):
         return node
 
 
-class ClosureTempAllocator(object):
-    def __init__(self, klass=None):
-        self.klass = klass
-        self.temps_allocated = {}
-        self.temps_free = {}
-        self.temps_count = 0
-
-    def reset(self):
-        for type, cnames in self.temps_allocated.items():
-            self.temps_free[type] = list(cnames)
-
-    def allocate_temp(self, type):
-        if not type in self.temps_allocated:
-            self.temps_allocated[type] = []
-            self.temps_free[type] = []
-        elif self.temps_free[type]:
-            return self.temps_free[type].pop(0)
-        cname = '%s%d' % (Naming.codewriter_temp_prefix, self.temps_count)
-        self.klass.declare_var(pos=None, name=cname, cname=cname, type=type, is_cdef=True)
-        self.temps_allocated[type].append(cname)
-        self.temps_count += 1
-        return cname
-
-    def put_gotref(self, code):
-        for entry in self.klass.entries.values():
-            if entry.cname == Naming.outer_scope_cname: # XXX
-                continue
-            if entry.type.is_pyobject:
-                code.put_xgotref('%s->%s' % (Naming.cur_scope_cname, entry.cname))
-
-    def put_giveref(self, code):
-        for entry in self.klass.entries.values():
-            if entry.cname == Naming.outer_scope_cname: # XXX
-                continue
-            if entry.type.is_pyobject:
-                code.put_xgiveref('%s->%s' % (Naming.cur_scope_cname, entry.cname))
-
 class YieldNodeCollector(TreeVisitor):
 
     def __init__(self):
@@ -1385,11 +1348,6 @@ class MarkClosureVisitor(CythonTransform):
         collector.visitchildren(node)
 
         if collector.yields:
-            allocator = ClosureTempAllocator()
-            # XXX: move allocator inside local scope
-            for y in collector.yields:
-                y.temp_allocator = allocator
-            node.temp_allocator = allocator
             node.is_generator = True
             node.needs_closure = True
             node.yields = collector.yields
@@ -1534,8 +1492,6 @@ class CreateClosureClasses(CythonTransform):
         class_scope.is_internal = True
         class_scope.directives = {'final': True}
 
-        if node.is_generator:
-            node.temp_allocator.klass = class_scope
         if from_closure:
             assert cscope.is_closure_scope
             class_scope.declare_var(pos=node.pos,