1 files changed, 69 insertions, 13 deletions

diff --git a/src/main/scala/firrtl/passes/wiring/WiringUtils.scala b/src/main/scala/firrtl/passes/wiring/WiringUtils.scala
index 5f09bbe0..c220692a 100644
--- a/src/main/scala/firrtl/passes/wiring/WiringUtils.scala
+++ b/src/main/scala/firrtl/passes/wiring/WiringUtils.scala
@@ -7,11 +7,13 @@ import firrtl._
 import firrtl.ir._
 import firrtl.Utils._
 import firrtl.Mappers._
+import firrtl.analyses.InstanceKeyGraph.InstanceKey
 import firrtl.traversals.Foreachers._
+
 import scala.collection.mutable
 import firrtl.annotations._
 import firrtl.annotations.AnnotationUtils._
-import firrtl.analyses.InstanceGraph
+import firrtl.analyses.{InstanceGraph, InstanceKeyGraph}
 
 /** Declaration kind in lineage (e.g. input port, output port, wire)
   */
@@ -112,9 +114,66 @@ object WiringUtils {
     * @return a map of sink instance names to source instance names
     * @throws WiringException if a sink is equidistant to two sources
     */
-  @deprecated("This method can lead to non-determinism in your compiler pass. Use sinksToSourcesSeq instead!", "Firrtl 1.4")
-  def sinksToSources(sinks: Seq[Named], source: String, i: InstanceGraph): Map[Seq[WDefInstance], Seq[WDefInstance]] =
-    sinksToSourcesSeq(sinks, source, i).toMap
+  @deprecated("This method can lead to non-determinism in your compiler pass and exposes internal details." +
+    " Please file an issue with firrtl if you have a use case!", "Firrtl 1.4")
+  def sinksToSources(sinks: Seq[Named], source: String, i: InstanceGraph): Map[Seq[WDefInstance], Seq[WDefInstance]] = {
+    // The order of owners influences the order of the results, it thus needs to be deterministic with a  LinkedHashMap.
+    val owners = new mutable.LinkedHashMap[Seq[WDefInstance], Vector[Seq[WDefInstance]]]
+    val queue = new mutable.Queue[Seq[WDefInstance]]
+    val visited = new mutable.HashMap[Seq[WDefInstance], Boolean].withDefaultValue(false)
+
+    val sourcePaths = i.fullHierarchy.collect { case (k,v) if k.module == source =>  v }
+    sourcePaths.flatten.foreach { l =>
+      queue.enqueue(l)
+      owners(l) = Vector(l)
+    }
+
+    val sinkModuleNames = sinks.map(getModuleName).toSet
+    val sinkPaths = i.fullHierarchy.collect { case (k,v) if sinkModuleNames.contains(k.module) => v }
+    // sinkInsts needs to have unique entries but is also iterated over which is why we use a LinkedHashSet
+    val sinkInsts = mutable.LinkedHashSet() ++ sinkPaths.flatten
+
+    /** If we're lucky and there is only one source, then that source owns
+      * all sinks. If we're unlucky, we need to do a full (slow) BFS
+      * to figure out who owns what. Currently, the BFS is not
+      * performant.
+      *
+      * [todo] The performance of this will need to be improved.
+      * Possible directions are that if we're purely source-under-sink
+      * or sink-under-source, then ownership is trivially a mapping
+      * down/up. Ownership seems to require a BFS if we have
+      * sources/sinks not under sinks/sources.
+      */
+    if (queue.size == 1) {
+      val u = queue.dequeue
+      sinkInsts.foreach { v => owners(v) = Vector(u) }
+    } else {
+      while (queue.nonEmpty) {
+        val u = queue.dequeue
+        visited(u) = true
+
+        val edges = (i.graph.getEdges(u.last).map(u :+ _).toVector :+ u.dropRight(1))
+
+        // [todo] This is the critical section
+        edges
+          .filter( e => !visited(e) && e.nonEmpty )
+          .foreach{ v =>
+            owners(v) = owners.getOrElse(v, Vector()) ++ owners(u)
+            queue.enqueue(v)
+          }
+      }
+
+      // Check that every sink has one unique owner. The only time that
+      // this should fail is if a sink is equidistant to two sources.
+      sinkInsts.foreach { s =>
+        if (!owners.contains(s) || owners(s).size > 1) {
+          throw new WiringException(
+            s"Unable to determine source mapping for sink '${s.map(_.name)}'") }
+      }
+    }
+
+    owners.collect { case (k, v) if sinkInsts.contains(k) => (k, v.flatten) }.toMap
+  }
 
   /** Return a map of sink instances to source instances that minimizes
     * distance
@@ -125,15 +184,12 @@ object WiringUtils {
     * @return a map of sink instance names to source instance names
     * @throws WiringException if a sink is equidistant to two sources
     */
-  def sinksToSourcesSeq(sinks: Seq[Named],
-                     source: String,
-                     i: InstanceGraph):
-      Seq[(Seq[WDefInstance], Seq[WDefInstance])] = {
+  private[firrtl] def sinksToSourcesSeq(sinks: Seq[Named], source: String, i: InstanceKeyGraph):
+      Seq[(Seq[InstanceKey], Seq[InstanceKey])] = {
     // The order of owners influences the order of the results, it thus needs to be deterministic with a  LinkedHashMap.
-    val owners = new mutable.LinkedHashMap[Seq[WDefInstance], Vector[Seq[WDefInstance]]]
-    val queue = new mutable.Queue[Seq[WDefInstance]]
-    val visited = new mutable.HashMap[Seq[WDefInstance], Boolean]
-      .withDefaultValue(false)
+    val owners = new mutable.LinkedHashMap[Seq[InstanceKey], Vector[Seq[InstanceKey]]]
+    val queue = new mutable.Queue[Seq[InstanceKey]]
+    val visited = new mutable.HashMap[Seq[InstanceKey], Boolean].withDefaultValue(false)
 
     val sourcePaths = i.fullHierarchy.collect { case (k,v) if k.module == source =>  v }
     sourcePaths.flatten.foreach { l =>
@@ -165,7 +221,7 @@ object WiringUtils {
         val u = queue.dequeue
         visited(u) = true
 
-        val edges = (i.graph.getEdges(u.last).map(u :+ _).toVector :+ u.dropRight(1))
+        val edges = i.graph.getEdges(u.last).map(u :+ _).toVector :+ u.dropRight(1)
 
         // [todo] This is the critical section
         edges