1 files changed, 24 insertions, 19 deletions
diff --git a/src/main/scala/firrtl/passes/wiring/WiringTransform.scala b/src/main/scala/firrtl/passes/wiring/WiringTransform.scala
index 7a6fcc65..5c251d6d 100644
--- a/src/main/scala/firrtl/passes/wiring/WiringTransform.scala
+++ b/src/main/scala/firrtl/passes/wiring/WiringTransform.scala
@@ -13,7 +13,7 @@ import WiringUtils._
 
 /** A component, e.g. register etc. Must be declared only once under the TopAnnotation
   */
-case class SourceAnnotation(target: ComponentName) extends Annotation with Loose with Unstable {
+case class SourceAnnotation(target: ComponentName, pin: String) extends Annotation with Loose with Unstable {
   def transform = classOf[WiringTransform]
   def duplicate(n: Named) = n match {
     case n: ComponentName => this.copy(target = n)
@@ -34,7 +34,7 @@ case class SinkAnnotation(target: ModuleName, pin: String) extends Annotation wi
 /** A module under which all sink module must be declared, and there is only
   * one source component
   */
-case class TopAnnotation(target: ModuleName) extends Annotation with Loose with Unstable {
+case class TopAnnotation(target: ModuleName, pin: String) extends Annotation with Loose with Unstable {
   def transform = classOf[WiringTransform]
   def duplicate(n: Named) = n match {
     case n: ModuleName => this.copy(target = n)
@@ -57,30 +57,35 @@ case class TopAnnotation(target: ModuleName) extends Annotation with Loose with
 class WiringTransform extends Transform with SimpleRun {
   def inputForm = MidForm
   def outputForm = MidForm
-  def passSeq(wi: WiringInfo) =
-    Seq(new Wiring(wi),
+  def passSeq(wis: Seq[WiringInfo]) =
+    Seq(new Wiring(wis),
         InferTypes,
         ResolveKinds,
         ResolveGenders)
   def execute(state: CircuitState): CircuitState = getMyAnnotations(state) match {
     case Nil => CircuitState(state.circuit, state.form)
-    case p => 
-      val sinks = mutable.HashMap[String, String]()
-      val sources = mutable.Set[String]()
-      val tops = mutable.Set[String]()
-      val comp = mutable.Set[String]()
-      p.foreach { 
-        case SinkAnnotation(m, pin) => sinks(m.name) = pin
-        case SourceAnnotation(c) =>
-          sources += c.module.name
-          comp += c.name
-        case TopAnnotation(m) => tops += m.name
+    case p =>
+      // Pin to value
+      val sinks = mutable.HashMap[String, Set[String]]()
+      val sources = mutable.HashMap[String, String]()
+      val tops = mutable.HashMap[String, String]()
+      val comp = mutable.HashMap[String, String]()
+      p.foreach { a =>
+        a match {
+          case SinkAnnotation(m, pin) => sinks(pin) = sinks.getOrElse(pin, Set.empty) + m.name
+          case SourceAnnotation(c, pin) =>
+            sources(pin) = c.module.name
+            comp(pin) = c.name
+          case TopAnnotation(m, pin) => tops(pin) = m.name
+        }
       }
       (sources.size, tops.size, sinks.size, comp.size) match {
-        case (0, 0, p, 0) => state
-        case (1, 1, p, 1) if p > 0 =>
-          val winfo = WiringInfo(sources.head, comp.head, sinks.toMap, tops.head)
-          state.copy(circuit = runPasses(state.circuit, passSeq(winfo)))
+        case (0, 0, p, 0) => state.copy(annotations = None)
+        case (s, t, p, c) if (p > 0) & (s == t) & (t == c) =>
+          val wis = tops.foldLeft(Seq[WiringInfo]()) { case (seq, (pin, top)) =>
+            seq :+ WiringInfo(sources(pin), comp(pin), sinks(pin), pin, top)
+          }
+          state.copy(circuit = runPasses(state.circuit, passSeq(wis)), annotations = None)
         case _ => error("Wrong number of sources, tops, or sinks!")
       }
   }