From 6097b6098a87a42e7180156d9002b2d99723ce9c Mon Sep 17 00:00:00 2001
From: Oscar Boykin <oscar@stripe.com>
Date: Thu, 5 Oct 2017 10:10:10 -1000
Subject: [PATCH 1/2] Add a few convenience methods

---
 .../src/main/scala/com/stripe/dagon/Dag.scala | 36 ++++++++++++++++-
 .../scala/com/stripe/dagon/DataFlowTest.scala | 39 +++++++++++++++++++
 2 files changed, 74 insertions(+), 1 deletion(-)

diff --git a/core/src/main/scala/com/stripe/dagon/Dag.scala b/core/src/main/scala/com/stripe/dagon/Dag.scala
index 5f510da..64dc000 100644
--- a/core/src/main/scala/com/stripe/dagon/Dag.scala
+++ b/core/src/main/scala/com/stripe/dagon/Dag.scala
@@ -406,6 +406,24 @@ sealed abstract class Dag[N[_]] { self =>
   def contains(node: N[_]): Boolean =
     find(node).isDefined
 
+  /**
+   * What nodes do we depend directly on
+   */
+  def dependenciesOf(node: N[_]): List[N[_]] = {
+    toLiteral(node) match {
+      case Literal.Const(_) =>
+        Nil
+      case Literal.Unary(n, _) =>
+        n.evaluate :: Nil
+      case Literal.Binary(n1, n2, _) =>
+        val evalLit = Literal.evaluateMemo[N]
+        evalLit(n1) :: evalLit(n2) :: Nil
+      case Literal.Variadic(inputs, _) =>
+        val evalLit = Literal.evaluateMemo[N]
+        inputs.map(evalLit(_))
+    }
+  }
+
   /**
    * list all the nodes that depend on the given node
    */
@@ -433,7 +451,13 @@ sealed abstract class Dag[N[_]] { self =>
   }
 
   /**
-   * Return all dependendants of a given node.
+   * equivalent to (but maybe faster than) fanOut(n) > 1
+   */
+  def hasSingleDependent(n: N[_]): Boolean =
+    fanOut(n) > 1
+
+  /**
+   * Return all dependents of a given node.
    * Does not include itself
    */
   def transitiveDependentsOf(p: N[_]): Set[N[_]] = {
@@ -442,6 +466,16 @@ sealed abstract class Dag[N[_]] { self =>
 
     Graphs.depthFirstOf(p.asInstanceOf[N[Any]])(nfn _).toSet
   }
+
+  /**
+   * Return the transitive dependencies of a given node
+   */
+  def transitiveDependenciesOf(p: N[_]): Set[N[_]] = {
+    def nfn(n: N[Any]): List[N[Any]] =
+      dependenciesOf(n).toList.asInstanceOf[List[N[Any]]]
+
+    Graphs.depthFirstOf(p.asInstanceOf[N[Any]])(nfn _).toSet
+  }
 }
 
 object Dag {
diff --git a/core/src/test/scala/com/stripe/dagon/DataFlowTest.scala b/core/src/test/scala/com/stripe/dagon/DataFlowTest.scala
index fb4893d..ad8d4a1 100644
--- a/core/src/test/scala/com/stripe/dagon/DataFlowTest.scala
+++ b/core/src/test/scala/com/stripe/dagon/DataFlowTest.scala
@@ -429,6 +429,17 @@ class DataFlowTest extends FunSuite {
     }
   }
 
+  test("transitiveDependenciesOf matches Flow.transitiveDeps") {
+    forAll { (f: Flow[Int], rule: Rule[Flow], max: Int) =>
+      val (dag, id) = Dag(f, Flow.toLiteral)
+
+      val optimizedDag = dag.applyMax(rule, max)
+
+      val optF = optimizedDag.evaluate(id)
+      assert(optimizedDag.transitiveDependenciesOf(optF) == (Flow.transitiveDeps(optF).toSet - optF), s"optimized: $optF $optimizedDag")
+    }
+  }
+
   test("Dag: findAll(n).forall(evaluate(_) == n)") {
     forAll { (f: Flow[Int], rule: Rule[Flow], max: Int) =>
       val (dag, id) = Dag(f, Flow.toLiteral)
@@ -483,6 +494,34 @@ class DataFlowTest extends FunSuite {
     }
   }
 
+  test("dependenciesOf matches toLiteral") {
+    forAll { (f: Flow[Int]) =>
+      val (dag, id) = Dag(f, Flow.toLiteral)
+
+      def contract(n: Flow[_]): List[Flow[_]] =
+        Flow.toLiteral(n) match {
+          case Literal.Const(_) => Nil
+          case Literal.Unary(n, _) => n.evaluate :: Nil
+          case Literal.Binary(n1, n2, _) => n1.evaluate :: n2.evaluate :: Nil
+          case Literal.Variadic(ns, _) => ns.map(_.evaluate)
+        }
+
+      dag.allNodes.foreach { n =>
+        assert(dag.dependenciesOf(n) == contract(n))
+      }
+    }
+  }
+
+  test("hasSingleDependent matches fanOut") {
+    forAll { (f: Flow[Int]) =>
+      val (dag, id) = Dag(f, Flow.toLiteral)
+
+      dag.allNodes.foreach { n =>
+        assert(dag.hasSingleDependent(n) == (dag.fanOut(n) > 1))
+      }
+    }
+  }
+
   test("contains(n) is the same as allNodes.contains(n)") {
     forAll { (f: Flow[Int], rule: Rule[Flow], max: Int, check: List[Flow[Int]]) =>
       val (dag, _) = Dag(f, Flow.toLiteral)

From 81764e97cd60ac67e35aed46c872e2de61a5e2d5 Mon Sep 17 00:00:00 2001
From: Oscar Boykin <oscar@stripe.com>
Date: Thu, 5 Oct 2017 10:20:13 -1000
Subject: [PATCH 2/2] reverse the boolean

---
 core/src/main/scala/com/stripe/dagon/Dag.scala          | 4 ++--
 core/src/test/scala/com/stripe/dagon/DataFlowTest.scala | 9 ++++++++-
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/core/src/main/scala/com/stripe/dagon/Dag.scala b/core/src/main/scala/com/stripe/dagon/Dag.scala
index 64dc000..f325a23 100644
--- a/core/src/main/scala/com/stripe/dagon/Dag.scala
+++ b/core/src/main/scala/com/stripe/dagon/Dag.scala
@@ -451,10 +451,10 @@ sealed abstract class Dag[N[_]] { self =>
   }
 
   /**
-   * equivalent to (but maybe faster than) fanOut(n) > 1
+   * equivalent to (but maybe faster than) fanOut(n) <= 1
    */
   def hasSingleDependent(n: N[_]): Boolean =
-    fanOut(n) > 1
+    fanOut(n) <= 1
 
   /**
    * Return all dependents of a given node.
diff --git a/core/src/test/scala/com/stripe/dagon/DataFlowTest.scala b/core/src/test/scala/com/stripe/dagon/DataFlowTest.scala
index ad8d4a1..47e46db 100644
--- a/core/src/test/scala/com/stripe/dagon/DataFlowTest.scala
+++ b/core/src/test/scala/com/stripe/dagon/DataFlowTest.scala
@@ -517,8 +517,15 @@ class DataFlowTest extends FunSuite {
       val (dag, id) = Dag(f, Flow.toLiteral)
 
       dag.allNodes.foreach { n =>
-        assert(dag.hasSingleDependent(n) == (dag.fanOut(n) > 1))
+        assert(dag.hasSingleDependent(n) == (dag.fanOut(n) <= 1))
       }
+
+      dag
+        .allNodes
+        .filter(dag.hasSingleDependent)
+        .foreach { n =>
+          assert(dag.dependentsOf(n).size <= 1)
+        }
     }
   }