Lower double recursion to simple recurtion

theano/gof/fg.py

@@ -281,7 +281,11 @@ class FunctionGraph(utils.object2):
         """
         Removes all from the clients list of r.
 
-        WRITEME
+        This is the main method to remove variable or apply node from
+        an FunctionGraph.
+
+        If called with an empty list of clients and prune=True, this
+        will remove the owner of the variable (so an apply_node).
 
         Parameters
         ----------
@@ -289,16 +293,60 @@ class FunctionGraph(utils.object2):
             Variable.
         clients_to_remove
             List of (op, i) pairs such that node.inputs[i] is not r anymore.
+        prune
+            If prune is True, it remove r from this fgraph if it don't
+            have clients left.
+
+        Returns
+        -------
+
+        True if r is still in the fgraph and need to be pruned
+        later. This can happen only when prune is False. A second call
+        to this method with an empty list for clients_to_remove and
+        prune=True will remove r.
 
         """
         for entry in clients_to_remove:
             r.clients.remove(entry)
             assert entry not in r.clients  # an op,i pair should be unique
-        if not r.clients:
-            if prune:
-                self.__prune_r__(r, reason)
+        if r.clients:
             return False
+        if not prune:
             return True
+        variable = r
+        if variable.owner:
+            apply_node = variable.owner
+            used_or_output = [output for output in apply_node.outputs
+                              if output.clients or output in self.outputs]
+            # If the apply node is not used and is not an output
+            if not used_or_output:
+                self.apply_nodes.remove(apply_node)
+                self.variables.difference_update(apply_node.outputs)
+                self.execute_callbacks('on_prune', apply_node, reason)
+
+                for i, input in enumerate(apply_node.inputs):
+                    self.__remove_clients__(input, [(apply_node, i)],
+                                            reason=reason)
+        # variable should not have any clients.
+        # assert not variable.clients
+
+        # variable should be in self.variables
+        # Why this assert fail? Making it True could cause opt speed up
+        # I think this is caused as we remove var in self.variables in
+        # another place.
+        # assert variable in self.variables
+
+        if variable in self.variables:
+            # If the owner have other outputs still used,
+            # then we must keep that variable in the graph.
+            if not variable.owner or not any(
+                [var for var in variable.owner.outputs
+                 if var.clients]):
+
+                self.variables.remove(variable)
+                # This allow to quickly know if a var is still in the fgraph
+                # or not.
+                del variable.fgraph
         return False
 
     # import #
@@ -443,52 +491,6 @@ class FunctionGraph(utils.object2):
             assert node.fgraph is self
             self.execute_callbacks('on_import', node, reason)
 
-    # prune #
-    def __prune_r__(self, variable, reason=None):
-        """
-        Should be called for variable that aren't used anymore:
-        len(var.clients) == 0.
-
-        This do not mean we will remove it from fgraph.variables. If
-        the owner stay in the fgraph as other outputs are still used,
-        the variable will stay in fgraph.variables.
-
-        """
-        # Prunes the owners(Apply Node) of the variables.
-        if variable.owner:
-            apply_node = variable.owner
-            used_or_output = [output for output in apply_node.outputs
-                              if output.clients or output in self.outputs]
-            # If the apply node is not used and is not an output
-            if not used_or_output:
-                self.apply_nodes.remove(apply_node)
-                self.variables.difference_update(apply_node.outputs)
-                self.execute_callbacks('on_prune', apply_node, reason)
-
-                for i, input in enumerate(apply_node.inputs):
-                    self.__remove_clients__(input, [(apply_node, i)],
-                                            reason=reason)
-        # variable should not have any clients.
-        # assert not variable.clients
-
-        # variable should be in self.variables
-        # Why this assert fail? Making it True could cause opt speed up
-        # I think this is caused as we remove var in self.variables in
-        # another place.
-        # assert variable in self.variables
-
-        if variable in self.variables:
-            # If the owner have other outputs still used,
-            # then we must keep that variable in the graph.
-            if not variable.owner or not any(
-                [var for var in variable.owner.outputs
-                 if var.clients]):
-
-                self.variables.remove(variable)
-                # This allow to quickly know if a var is still in the fgraph
-                # or not.
-                del variable.fgraph
-
 
     # change input #
     def change_input(self, node, i, new_r, reason=None):
@@ -534,9 +536,7 @@ class FunctionGraph(utils.object2):
         # transaction will be reverted later.
         self.execute_callbacks('on_change_input', node, i,
                                r, new_r, reason=reason)
-
-        if prune:
-            self.__prune_r__(r, reason=reason)
+        prune = self.__remove_clients__(r, [], True)
 
     # replace #
     def replace(self, r, new_r, reason=None, verbose=None):