Fixed multi wiring (#368)

* Fixed multi wiring

* Minor style changes

2 files changed, 34 insertions, 27 deletions

diff --git a/src/main/scala/firrtl/passes/wiring/Wiring.scala b/src/main/scala/firrtl/passes/wiring/Wiring.scala
index bc1a2524..1ced07eb 100644
--- a/src/main/scala/firrtl/passes/wiring/Wiring.scala
+++ b/src/main/scala/firrtl/passes/wiring/Wiring.scala
@@ -13,10 +13,13 @@ import WiringUtils._
 
 case class WiringException(msg: String) extends PassException(msg)
 
-case class WiringInfo(source: String, comp: String, sinks: Map[String, String], top: String)
+case class WiringInfo(source: String, comp: String, sinks: Set[String], pin: String, top: String)
 
-class Wiring(wi: WiringInfo) extends Pass {
+class Wiring(wiSeq: Seq[WiringInfo]) extends Pass {
   def name = this.getClass.getSimpleName
+  def run(c: Circuit): Circuit = {
+    wiSeq.foldLeft(c) { (circuit, wi) => wire(circuit, wi) }
+  }
 
   /** Add pins to modules and wires a signal to them, under the scope of a specified top module
     * Description:
@@ -30,11 +33,12 @@ class Wiring(wi: WiringInfo) extends Pass {
     * Notes:
     *   - No module uniquification occurs (due to imposed restrictions)
     */
-  def run(c: Circuit): Circuit = {
+  def wire(c: Circuit, wi: WiringInfo): Circuit = {
     // Split out WiringInfo
     val source = wi.source
-    val sinks = wi.sinks.keys.toSet
+    val sinks = wi.sinks
     val compName = wi.comp
+    val pin = wi.pin
 
     // Maps modules to children instances, i.e. (instance, module)
     val childrenMap = getChildrenMap(c)
@@ -55,10 +59,8 @@ class Wiring(wi: WiringInfo) extends Pass {
     val portNames = c.modules.foldLeft(Map.empty[String, String]) { (map, m) =>
       map + (m.name -> {
         val ns = Namespace(m)
-        wi.sinks.get(m.name) match {
-          case Some(pin) => ns.newName(pin)
-          case None => ns.newName(tokenize(compName) filterNot ("[]." contains _) mkString "_")
-        }
+        if(sinks.contains(m.name)) ns.newName(pin)
+        else ns.newName(tokenize(compName) filterNot ("[]." contains _) mkString "_")
       })
     }
 
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!")
       }
   }