Merge pull request #1134 from lamblin/fix_ordered_set_memleak

theano/gof/python25.py

@@ -290,8 +290,8 @@ if sys.version_info[:2] < (2, 7):
                 return True
             return dict.__eq__(self, other)
 
-    def __ne__(self, other):
-        return not self == other
+        def __ne__(self, other):
+            return not self == other
 
 else:
     try:

theano/misc/ordered_set.py

@@ -19,9 +19,7 @@ def check_deterministic(iterable):
             list, tuple, OrderedSet, types.GeneratorType, basestring))
 
 if MutableSet is not None:
-    # From http://code.activestate.com/recipes/576694/
     # Copyright (C) 2009 Raymond Hettinger
-
     # Permission is hereby granted, free of charge, to any person obtaining a
     # copy of this software and associated documentation files (the
     # "Software"), to deal in the Software without restriction, including
@@ -40,9 +38,22 @@ if MutableSet is not None:
     # CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
     # TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
     # SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
-    KEY, PREV, NEXT = range(3)
-
-    class OrderedSet(MutableSet):
+    ## {{{ http://code.activestate.com/recipes/576696/ (r5)
+    import collections
+    from weakref import proxy
+
+    class Link(object):
+        __slots__ = 'prev', 'next', 'key', '__weakref__'
+
+    class OrderedSet(collections.MutableSet):
+        'Set the remembers the order elements were added'
+        # Big-O running times for all methods are the same as for regular sets.
+        # The internal self.__map dictionary maps keys to links in a doubly linked list.
+        # The circular doubly linked list starts and ends with a sentinel element.
+        # The sentinel element never gets deleted (this simplifies the algorithm).
+        # The prev/next links are weakref proxies (to prevent circular references).
+        # Individual links are kept alive by the hard reference in self.__map.
+        # Those hard references disappear when a key is deleted from an OrderedSet.
 
         # Added by IG-- pre-existing theano code expected sets
         #   to have this method
@@ -53,43 +64,50 @@ if MutableSet is not None:
         def __init__(self, iterable=None):
             # Checks added by IG
             check_deterministic(iterable)
-            self.end = end = []
-            end += [None, end, end]         # sentinel node for doubly linked list
-            self.map = {}                   # key --> [key, prev, next]
+            self.__root = root = Link()         # sentinel node for doubly linked list
+            root.prev = root.next = root
+            self.__map = {}                     # key --> link
             if iterable is not None:
                 self |= iterable
 
         def __len__(self):
-            return len(self.map)
+            return len(self.__map)
 
         def __contains__(self, key):
-            return key in self.map
+            return key in self.__map
 
         def add(self, key):
-            if key not in self.map:
-                end = self.end
-                curr = end[PREV]
-                curr[NEXT] = end[PREV] = self.map[key] = [key, curr, end]
+            # Store new key in a new link at the end of the linked list
+            if key not in self.__map:
+                self.__map[key] = link = Link()
+                root = self.__root
+                last = root.prev
+                link.prev, link.next, link.key = last, root, key
+                last.next = root.prev = proxy(link)
 
         def discard(self, key):
-            if key in self.map:
-                key, prev, next = self.map.pop(key)
-                prev[NEXT] = next
-                next[PREV] = prev
+            # Remove an existing item using self.__map to find the link which is
+            # then removed by updating the links in the predecessor and successors.
+            if key in self.__map:
+                link = self.__map.pop(key)
+                link.prev.next = link.next
+                link.next.prev = link.prev
 
         def __iter__(self):
-            end = self.end
-            curr = end[NEXT]
-            while curr is not end:
-                yield curr[KEY]
-                curr = curr[NEXT]
+            # Traverse the linked list in order.
+            root = self.__root
+            curr = root.next
+            while curr is not root:
+                yield curr.key
+                curr = curr.next
 
         def __reversed__(self):
-            end = self.end
-            curr = end[PREV]
-            while curr is not end:
-                yield curr[KEY]
-                curr = curr[PREV]
+            # Traverse the linked list in reverse order.
+            root = self.__root
+            curr = root.prev
+            while curr is not root:
+                yield curr.key
+                curr = curr.prev
 
         def pop(self, last=True):
             if not self:
@@ -123,8 +141,8 @@ if MutableSet is not None:
             else:
                 return NotImplemented
 
-        def __del__(self):
-            self.clear()                    # remove circular references
+    ## end of http://code.activestate.com/recipes/576696/ }}}
+
 else:
     # Python 2.4
     class OrderedSet(object):