Add a few convenience methods

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 {

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,41 @@ 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))
+      }
+
+      dag
+        .allNodes
+        .filter(dag.hasSingleDependent)
+        .foreach { n =>
+          assert(dag.dependentsOf(n).size <= 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)