From 31436cdeddd554d9685ad19479ca5521cf2bf909 Mon Sep 17 00:00:00 2001
From: Scott Sanderson <ssanderson@quantopian.com>
Date: Wed, 14 Sep 2016 11:16:40 -0400
Subject: [PATCH] MAINT: Move refcount management into TermGraph.

---
 zipline/pipeline/engine.py | 12 ++++--------
 zipline/pipeline/graph.py  | 32 +++++++++++++++++++++++++++++++-
 2 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/zipline/pipeline/engine.py b/zipline/pipeline/engine.py
index 93d296ed..e17331a3 100644
--- a/zipline/pipeline/engine.py
+++ b/zipline/pipeline/engine.py
@@ -344,7 +344,7 @@ class SimplePipelineEngine(object):
         loader_group_key = juxt(get_loader, getitem(graph.extra_rows))
         loader_groups = groupby(loader_group_key, graph.loadable_terms)
 
-        refcounts = graph.initial_refcounts()
+        refcounts = graph.initial_refcounts(workspace)
 
         for term in graph.ordered():
             # `term` may have been supplied in `initial_workspace`, and in the
@@ -382,13 +382,9 @@ class SimplePipelineEngine(object):
                 else:
                     assert workspace[term].shape == (mask.shape[0], 1)
 
-                # Decref any term we depended on.
-                for (parent, _) in graph.in_edges(term):
-                    refcounts[parent] -= 1
-                    # No one else depends on this term. Remove it from the
-                    # workspace to conserve memory.
-                    if refcounts[parent] == 0:
-                        del workspace[parent]
+                # Decref dependencies of ``term``, and clear any terms whose refcounts hit 0.
+                for garbage_term in graph.decref_dependencies(term, refcounts):
+                    del workspace[garbage_term]
 
         out = {}
         graph_extra_rows = graph.extra_rows
diff --git a/zipline/pipeline/graph.py b/zipline/pipeline/graph.py
index 9078c7ea..78dffc14 100644
--- a/zipline/pipeline/graph.py
+++ b/zipline/pipeline/graph.py
@@ -119,7 +119,7 @@ class TermGraph(DiGraph):
     def _repr_png_(self):
         return self.png.data
 
-    def initial_refcounts(self):
+    def initial_refcounts(self, initial_workspace):
         """
         Calculate initial refcounts for execution of this graph.
 
@@ -130,8 +130,38 @@ class TermGraph(DiGraph):
         refcounts = self.out_degree()
         for t in self.outputs.values():
             refcounts[t] += 1
+
+        for t in initial_workspace:
+            self.decref_dependencies(t, refcounts)
+
         return refcounts
 
+    def decref_dependencies(self, term, refcounts):
+        """
+        Decrement in-edges for ``term`` after computation.
+
+        Parameters
+        ----------
+        term : zipline.pipeline.Term
+            The term whose parents should be decref'ed.
+        refcounts : dict[Term -> int]
+            Dictionary of refcounts.
+
+        Return
+        ------
+        garbage : set[Term]
+            Terms whose refcounts hit zero after decrefing.
+        """
+        garbage = set()
+        # Edges are tuple of (from, to).
+        for parent, _ in self.in_edges([term]):
+            refcounts[parent] -= 1
+            # No one else depends on this term. Remove it from the
+            # workspace to conserve memory.
+            if refcounts[parent] == 0:
+                garbage.add(parent)
+        return garbage
+
 
 class ExecutionPlan(TermGraph):
     """