Add optionalPrerequisites to Dependency API

Signed-off-by: Schuyler Eldridge <schuyler.eldridge@ibm.com>

3 files changed, 64 insertions, 2 deletions

diff --git a/src/main/scala/firrtl/options/DependencyManager.scala b/src/main/scala/firrtl/options/DependencyManager.scala
index dc23d3ce..4880ab8f 100644
--- a/src/main/scala/firrtl/options/DependencyManager.scala
+++ b/src/main/scala/firrtl/options/DependencyManager.scala
@@ -126,6 +126,14 @@ trait DependencyManager[A, B <: TransformLike[A] with DependencyAPI[B]] extends
     DiGraph(new LinkedHashMap() ++ v.map(vv => vv -> (v & (vv.dependents.toSet).map(dToO)))).reverse
   }
 
+  /** A directed graph of *optional* prerequisites. Each optional prerequisite is promoted to a full prerequisite if the
+    * optional prerequisite is already a node in the prerequisite graph.
+    */
+  private lazy val optionalPrerequisitesGraph: DiGraph[B] = {
+    val v = new LinkedHashSet() ++ prerequisiteGraph.getVertices
+    DiGraph(new LinkedHashMap() ++ v.map(vv => vv -> (v & (vv.optionalPrerequisites.toSet).map(dToO))))
+  }
+
   /** A directed graph consisting of prerequisites derived from ALL targets. This is necessary for defining targets for
     * [[DependencyManager]] sub-problems.
     */
@@ -143,7 +151,7 @@ trait DependencyManager[A, B <: TransformLike[A] with DependencyAPI[B]] extends
   }
 
   /** A directed graph consisting of all prerequisites, including prerequisites derived from dependents */
-  lazy val dependencyGraph: DiGraph[B] = prerequisiteGraph + dependentsGraph
+  lazy val dependencyGraph: DiGraph[B] = prerequisiteGraph + dependentsGraph + optionalPrerequisitesGraph
 
   /** A directed graph consisting of invalidation edges */
   lazy val invalidateGraph: DiGraph[B] = {
@@ -282,7 +290,8 @@ trait DependencyManager[A, B <: TransformLike[A] with DependencyAPI[B]] extends
     val connections =
       Seq( (prerequisiteGraph, "edge []"),
            (dependentsGraph,   """edge [style=bold color="#4292c6"]"""),
-           (invalidateGraph,   """edge [minlen=2 style=dashed constraint=false color="#fb6a4a"]""") )
+           (invalidateGraph,   """edge [minlen=2 style=dashed constraint=false color="#fb6a4a"]"""),
+           (optionalPrerequisitesGraph, """edge [style=dotted color="#a1d99b"]""") )
         .flatMap{ case (a, b) => toGraphviz(a, b) }
         .mkString("\n")
 
diff --git a/src/main/scala/firrtl/options/Phase.scala b/src/main/scala/firrtl/options/Phase.scala
index 0f29d165..0e534ec8 100644
--- a/src/main/scala/firrtl/options/Phase.scala
+++ b/src/main/scala/firrtl/options/Phase.scala
@@ -103,6 +103,13 @@ trait DependencyAPI[A <: DependencyAPI[A]] { this: TransformLike[_] =>
   def prerequisites: Seq[Dependency[A]] = Seq.empty
   private[options] lazy val _prerequisites: LinkedHashSet[Dependency[A]] = new LinkedHashSet() ++ prerequisites.toSet
 
+  /** All transforms that, if a prerequisite of *another* transform, will run before this transform.
+    * $seqNote
+    */
+  def optionalPrerequisites: Seq[Dependency[A]] = Seq.empty
+  private[options] lazy val _optionalPrerquisites: LinkedHashSet[Dependency[A]] =
+    new LinkedHashSet() ++ optionalPrerequisites.toSet
+
   /** All transforms that must run ''after'' this transform
     *
     * ''This is a means of prerequisite injection into some other transform.'' Normally a transform will define its own
diff --git a/src/test/scala/firrtlTests/options/PhaseManagerSpec.scala b/src/test/scala/firrtlTests/options/PhaseManagerSpec.scala
index b13ee215..744d621e 100644
--- a/src/test/scala/firrtlTests/options/PhaseManagerSpec.scala
+++ b/src/test/scala/firrtlTests/options/PhaseManagerSpec.scala
@@ -328,6 +328,34 @@ object CustomAfterOptimizationFixture {
 
 }
 
+object OptionalPrerequisitesFixture {
+
+  class Root extends IdentityPhase
+
+  class OptMinimum extends IdentityPhase with PreservesAll[Phase] {
+    override val prerequisites = Seq(Dependency[Root])
+  }
+
+  class OptFull extends IdentityPhase with PreservesAll[Phase] {
+    override val prerequisites = Seq(Dependency[Root], Dependency[OptMinimum])
+  }
+
+  class DoneMinimum extends IdentityPhase with PreservesAll[Phase] {
+    override val prerequisites = Seq(Dependency[OptMinimum])
+  }
+
+  class DoneFull extends IdentityPhase with PreservesAll[Phase] {
+    override val prerequisites = Seq(Dependency[OptFull])
+  }
+
+  class Custom extends IdentityPhase with PreservesAll[Phase] {
+    override val prerequisites = Seq(Dependency[Root])
+    override val optionalPrerequisites = Seq(Dependency[OptMinimum], Dependency[OptFull])
+    override val dependents = Seq(Dependency[DoneMinimum], Dependency[DoneFull])
+  }
+
+}
+
 object OrderingFixture {
 
   class A extends IdentityPhase with PreservesAll[Phase]
@@ -618,6 +646,24 @@ class PhaseManagerSpec extends FlatSpec with Matchers {
     pmFull.flattenedTransformOrder.map(_.getClass) should be (expectedFull)
   }
 
+  it should "support optional prerequisites" in {
+    val f = OptionalPrerequisitesFixture
+
+    val targetsMinimum = Seq(Dependency[f.Custom], Dependency[f.DoneMinimum])
+    val pmMinimum = new PhaseManager(targetsMinimum)
+
+    val targetsFull = Seq(Dependency[f.Custom], Dependency[f.DoneFull])
+    val pmFull = new PhaseManager(targetsFull)
+
+    val expectedMinimum = Seq(classOf[f.Root], classOf[f.OptMinimum], classOf[f.Custom], classOf[f.DoneMinimum])
+    writeGraphviz(pmMinimum, "test_run_dir/PhaseManagerSpec/CustomAfterOptimization/minimum")
+    pmMinimum.flattenedTransformOrder.map(_.getClass) should be (expectedMinimum)
+
+    val expectedFull = Seq(classOf[f.Root], classOf[f.OptMinimum], classOf[f.OptFull], classOf[f.Custom], classOf[f.DoneFull])
+    writeGraphviz(pmFull, "test_run_dir/PhaseManagerSpec/CustomAfterOptimization/full")
+    pmFull.flattenedTransformOrder.map(_.getClass) should be (expectedFull)
+  }
+
   /** This tests a situation the ordering of edges matters. Namely, this test is dependent on the ordering in which
     * DiGraph.linearize walks the edges of each node.
     */