Wiring (#348)

Added wiring pass and simple test

3 files changed, 406 insertions, 0 deletions

diff --git a/src/main/scala/firrtl/passes/wiring/Wiring.scala b/src/main/scala/firrtl/passes/wiring/Wiring.scala
new file mode 100644
index 00000000..d3f6f3dd
--- /dev/null
+++ b/src/main/scala/firrtl/passes/wiring/Wiring.scala
@@ -0,0 +1,162 @@
+package firrtl.passes
+package wiring
+
+import firrtl._
+import firrtl.ir._
+import firrtl.Utils._
+import firrtl.Mappers._
+import scala.collection.mutable
+import firrtl.Annotations._
+import WiringUtils._
+
+case class WiringException(msg: String) extends PassException(msg)
+
+case class WiringInfo(source: String, comp: String, sinks: Map[String, String], top: String)
+
+class Wiring(wi: WiringInfo) extends Pass {
+  def name = this.getClass.getSimpleName
+
+  /** Add pins to modules and wires a signal to them, under the scope of a specified top module
+    * Description:
+    *   Adds a pin to each sink module
+    *   Punches ports up from the source signal to the specified top module
+    *   Punches ports down to each sink module
+    *   Wires the source up and down, connecting to all sink modules
+    * Restrictions:
+    *   - Can only have one source module instance under scope of the specified top
+    *   - All instances of each sink module must be under the scope of the specified top
+    * Notes:
+    *   - No module uniquification occurs (due to imposed restrictions)
+    */
+  def run(c: Circuit): Circuit = {
+    // Split out WiringInfo
+    val source = wi.source
+    val sinks = wi.sinks.keys.toSet
+    val compName = wi.comp
+
+    // Maps modules to children instances, i.e. (instance, module)
+    val childrenMap = getChildrenMap(c)
+
+    // Check restrictions
+    val nSources = countInstances(childrenMap, wi.top, source)
+    if(nSources != 1)
+      throw new WiringException(s"Cannot have $nSources instance of $source under ${wi.top}")
+    sinks.foreach { m =>
+      val total = countInstances(childrenMap, c.main, m)
+      val nTop = countInstances(childrenMap, c.main, wi.top)
+      val perTop = countInstances(childrenMap, wi.top, m)
+      if(total != nTop * perTop)
+        throw new WiringException(s"Module ${wi.top} does not contain all instances of $m.")
+    }
+
+    // Create valid port names for wiring that have no name conflicts
+    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 "_")
+        }
+      })
+    }
+
+    // Create a lineage tree from children map
+    val lineages = getLineage(childrenMap, wi.top)
+
+    // Populate lineage tree with relationship information, i.e. who is source,
+    //   sink, parent of source, etc.
+    val withFields = setSharedParent(wi.top)(setFields(sinks, source)(lineages))
+
+    // Populate lineage tree with what to instantiate, connect to/from, etc.
+    val withThings = setThings(portNames, compName)(withFields)
+
+    // Create a map from module name to lineage information
+    val map = pointToLineage(withThings)
+
+    // Obtain the source component type
+    val sourceComponentType = getType(c, source, compName)
+
+    // Return new circuit with correct wiring
+    val cx = c.copy(modules = c.modules map onModule(map, sourceComponentType))
+
+    // Replace inserted IR nodes with WIR nodes
+    ToWorkingIR.run(cx)
+  }
+
+  /** Return new module with correct wiring
+    */
+  private def onModule(map: Map[String, Lineage], t: Type)(m: DefModule) = {
+    map.get(m.name) match {
+      case None => m
+      case Some(l) =>
+        val stmts = mutable.ArrayBuffer[Statement]()
+        val ports = mutable.ArrayBuffer[Port]()
+        l.addPort match {
+          case None =>
+          case Some((s, dt)) => dt match {
+            case DecInput => ports += Port(NoInfo, s, Input, t)
+            case DecOutput => ports += Port(NoInfo, s, Output, t)
+            case DecWire =>
+              stmts += DefWire(NoInfo, s, t)
+          }
+        }
+        stmts ++= (l.cons map { case ((l, r)) =>
+          Connect(NoInfo, toExp(l), toExp(r))
+        })
+        def onStmt(s: Statement): Statement = Block(Seq(s) ++ stmts)
+        m match {
+          case Module(i, n, ps, s) => Module(i, n, ps ++ ports, Block(Seq(s) ++ stmts))
+          case ExtModule(i, n, ps, dn, p) => ExtModule(i, n, ps ++ ports, dn, p)
+        }
+    }
+  }
+
+  /** Returns the type of the component specified
+    */
+  private def getType(c: Circuit, module: String, comp: String) = {
+    def getRoot(e: Expression): String = e match {
+      case r: Reference => r.name
+      case i: SubIndex => getRoot(i.expr)
+      case a: SubAccess => getRoot(a.expr)
+      case f: SubField => getRoot(f.expr)
+    }
+    val eComp = toExp(comp)
+    val root = getRoot(eComp)
+    var tpe: Option[Type] = None
+    def getType(s: Statement): Statement = s match {
+      case DefRegister(_, n, t, _, _, _) if n == root =>
+        tpe = Some(t)
+        s
+      case DefWire(_, n, t) if n == root =>
+        tpe = Some(t)
+        s
+      case WDefInstance(_, n, m, t) if n == root => 
+        tpe = Some(t)
+        s
+      case DefNode(_, n, e) if n == root =>
+        tpe = Some(e.tpe)
+        s
+      case sx: DefMemory if sx.name == root =>
+        tpe = Some(MemPortUtils.memType(sx))
+        sx
+      case sx => sx map getType
+    }
+    val m = c.modules find (_.name == module) getOrElse error(s"Must have a module named $module")
+    tpe = m.ports find (_.name == root) map (_.tpe)
+    m match {
+      case Module(i, n, ps, b) => getType(b)
+      case e: ExtModule =>
+    }
+    tpe match {
+      case None => error(s"Didn't find $comp in $module!")
+      case Some(t) => 
+        def setType(e: Expression): Expression = e map setType match {
+          case ex: Reference => ex.copy(tpe = t)
+          case ex: SubField => ex.copy(tpe = field_type(ex.expr.tpe, ex.name))
+          case ex: SubIndex => ex.copy(tpe = sub_type(ex.expr.tpe))
+          case ex: SubAccess => ex.copy(tpe = sub_type(ex.expr.tpe))
+        }
+        setType(eComp).tpe
+    }
+  }
+}
diff --git a/src/main/scala/firrtl/passes/wiring/WiringTransform.scala b/src/main/scala/firrtl/passes/wiring/WiringTransform.scala
new file mode 100644
index 00000000..919948b6
--- /dev/null
+++ b/src/main/scala/firrtl/passes/wiring/WiringTransform.scala
@@ -0,0 +1,80 @@
+package firrtl.passes
+package wiring
+
+import firrtl._
+import firrtl.ir._
+import firrtl.Utils._
+import firrtl.Mappers._
+import scala.collection.mutable
+import firrtl.Annotations._
+import WiringUtils._
+
+/** A component, e.g. register etc. Must be declared only once under the TopAnnotation
+  */
+case class SourceAnnotation(target: ComponentName, tID: TransID) extends Annotation with Loose with Unstable {
+  def duplicate(n: Named) = n match {
+    case n: ComponentName => this.copy(target = n)
+    case _ => throwInternalError
+  }
+}
+
+/** A module, e.g. ExtModule etc., that should add the input pin
+  */
+case class SinkAnnotation(target: ModuleName, tID: TransID, pin: String) extends Annotation with Loose with Unstable {
+  def duplicate(n: Named) = n match {
+    case n: ModuleName => this.copy(target = n)
+    case _ => throwInternalError
+  }
+}
+
+/** A module under which all sink module must be declared, and there is only
+  * one source component
+  */
+case class TopAnnotation(target: ModuleName, tID: TransID) extends Annotation with Loose with Unstable {
+  def duplicate(n: Named) = n match {
+    case n: ModuleName => this.copy(target = n)
+    case _ => throwInternalError
+  }
+}
+
+/** Add pins to modules and wires a signal to them, under the scope of a specified top module
+  * Description:
+  *   Adds a pin to each sink module
+  *   Punches ports up from the source signal to the specified top module
+  *   Punches ports down to each sink module
+  *   Wires the source up and down, connecting to all sink modules
+  * Restrictions:
+  *   - Can only have one source module instance under scope of the specified top
+  *   - All instances of each sink module must be under the scope of the specified top
+  * Notes:
+  *   - No module uniquification occurs (due to imposed restrictions)
+  */
+class WiringTransform(transID: TransID) extends Transform with SimpleRun {
+  def passSeq(wi: WiringInfo) =
+    Seq(new Wiring(wi),
+        InferTypes,
+        ResolveKinds,
+        ResolveGenders)
+  def execute(c: Circuit, map: AnnotationMap) = map get transID match {
+    case Some(p) => 
+      val sinks = mutable.HashMap[String, String]()
+      val sources = mutable.Set[String]()
+      val tops = mutable.Set[String]()
+      val comp = mutable.Set[String]()
+      p.values.foreach { a =>
+        a match {
+          case SinkAnnotation(m, _, pin) => sinks(m.name) = pin
+          case SourceAnnotation(c, _) =>
+            sources += c.module.name
+            comp += c.name
+          case TopAnnotation(m, _) => tops += m.name
+        }
+      }
+      (sources.size, tops.size, sinks.size, comp.size) match {
+        case (0, 0, p, 0) => TransformResult(c)
+        case (1, 1, p, 1) if p > 0 => run(c, passSeq(WiringInfo(sources.head, comp.head, sinks.toMap, tops.head)))
+        case _ => error("Wrong number of sources, tops, or sinks!")
+      }
+    case None => TransformResult(c)
+  }
+}
diff --git a/src/main/scala/firrtl/passes/wiring/WiringUtils.scala b/src/main/scala/firrtl/passes/wiring/WiringUtils.scala
new file mode 100644
index 00000000..bfa94a81
--- /dev/null
+++ b/src/main/scala/firrtl/passes/wiring/WiringUtils.scala
@@ -0,0 +1,164 @@
+package firrtl.passes
+package wiring
+
+import firrtl._
+import firrtl.ir._
+import firrtl.Utils._
+import firrtl.Mappers._
+import scala.collection.mutable
+import firrtl.Annotations._
+import WiringUtils._
+
+/** Declaration kind in lineage (e.g. input port, output port, wire)
+  */
+sealed trait DecKind
+case object DecInput extends DecKind
+case object DecOutput extends DecKind
+case object DecWire extends DecKind
+
+/** A lineage tree representing the instance hierarchy in a design
+  */
+case class Lineage(
+    name: String,
+    children: Seq[(String, Lineage)] = Seq.empty,
+    source: Boolean = false,
+    sink: Boolean = false,
+    sourceParent: Boolean = false,
+    sinkParent: Boolean = false,
+    sharedParent: Boolean = false,
+    addPort: Option[(String, DecKind)] = None,
+    cons: Seq[(String, String)] = Seq.empty) {
+
+  def map(f: Lineage => Lineage): Lineage =
+    this.copy(children = children.map{ case (i, m) => (i, f(m)) })
+
+  override def toString: String = shortSerialize("")
+
+  def shortSerialize(tab: String): String = s"""
+    |$tab name: $name,
+    |$tab children: ${children.map(c => tab + "   " + c._2.shortSerialize(tab + "    "))}
+    |""".stripMargin
+
+  def serialize(tab: String): String = s"""
+    |$tab name: $name,
+    |$tab source: $source,
+    |$tab sink: $sink,
+    |$tab sourceParent: $sourceParent,
+    |$tab sinkParent: $sinkParent,
+    |$tab sharedParent: $sharedParent,
+    |$tab addPort: $addPort
+    |$tab cons: $cons
+    |$tab children: ${children.map(c => tab + "   " + c._2.serialize(tab + "    "))}
+    |""".stripMargin
+}
+
+
+
+
+object WiringUtils {
+  type ChildrenMap = mutable.HashMap[String, Seq[(String, String)]]
+
+  /** Given a circuit, returns a map from module name to children
+    * instance/module names
+    */
+  def getChildrenMap(c: Circuit): ChildrenMap = {
+    val childrenMap = new ChildrenMap()
+    def getChildren(mname: String)(s: Statement): Statement = s match {
+      case s: WDefInstance => 
+        childrenMap(mname) = childrenMap(mname) :+ (s.name, s.module)
+        s
+      case s: DefInstance => 
+        childrenMap(mname) = childrenMap(mname) :+ (s.name, s.module)
+        s
+      case s => s map getChildren(mname)
+    }
+    c.modules.foreach{ m =>
+      childrenMap(m.name) = Nil
+      m map getChildren(m.name)
+    }
+    childrenMap
+  }
+
+  /** Counts the number of instances of a module declared under a top module
+    */
+  def countInstances(childrenMap: ChildrenMap, top: String, module: String): Int = {
+    if(top == module) 1
+    else childrenMap(top).foldLeft(0) { case (count, (i, child)) =>
+      count + countInstances(childrenMap, child, module)
+    }
+  }
+
+  /** Returns a module's lineage, containing all children lineages as well
+    */
+  def getLineage(childrenMap: ChildrenMap, module: String): Lineage =
+    Lineage(module, childrenMap(module) map { case (i, m) => (i, getLineage(childrenMap, m)) } )
+
+  /** Sets the sink, sinkParent, source, and sourceParent fields of every
+    *  Lineage in tree
+    */
+  def setFields(sinks: Set[String], source: String)(lin: Lineage): Lineage = lin map setFields(sinks, source) match {
+    case l if sinks.contains(l.name) => l.copy(sink = true)
+    case l => 
+      val src = l.name == source
+      val sinkParent = l.children.foldLeft(false) { case (b, (i, m)) => b || m.sink || m.sinkParent }
+      val sourceParent = if(src) true else l.children.foldLeft(false) { case (b, (i, m)) => b || m.source || m.sourceParent }
+      l.copy(sinkParent=sinkParent, sourceParent=sourceParent, source=src)
+  }
+
+  /** Sets the sharedParent of lineage top
+    */
+  def setSharedParent(top: String)(lin: Lineage): Lineage = lin map setSharedParent(top) match {
+    case l if l.name == top => l.copy(sharedParent = true)
+    case l => l
+  }
+
+  /** Sets the addPort and cons fields of the lineage tree
+    */
+  def setThings(portNames:Map[String, String], compName: String)(lin: Lineage): Lineage = {
+    val funs = Seq(
+      ((l: Lineage) => l map setThings(portNames, compName)),
+      ((l: Lineage) => l match {
+        case Lineage(name, _, _, _, _, _, true, _, _) => //SharedParent
+          l.copy(addPort=Some((portNames(name), DecWire)))
+        case Lineage(name, _, _, _, true, _, _, _, _) => //SourceParent
+          l.copy(addPort=Some((portNames(name), DecOutput)))
+        case Lineage(name, _, _, _, _, true, _, _, _) => //SinkParent
+          l.copy(addPort=Some((portNames(name), DecInput)))
+        case Lineage(name, _, _, true, _, _, _, _, _) => //Sink
+          l.copy(addPort=Some((portNames(name), DecInput)))
+        case l => l
+      }),
+      ((l: Lineage) => l match {
+        case Lineage(name, _, true, _, _, _, _, _, _) => //Source
+          val tos = Seq(s"${portNames(name)}")
+          val from = compName
+          l.copy(cons = l.cons ++ tos.map(t => (t, from)))
+        case Lineage(name, _, _, _, true, _, _, _, _) => //SourceParent
+          val tos = Seq(s"${portNames(name)}")
+          val from = l.children.filter { case (i, c) => c.sourceParent }.map { case (i, c) => s"$i.${portNames(c.name)}" }.head
+          l.copy(cons = l.cons ++ tos.map(t => (t, from)))
+        case l => l
+      }),
+      ((l: Lineage) => l match {
+        case Lineage(name, _, _, _, _, true, _, _, _) => //SinkParent
+          val tos = l.children.filter { case (i, c) => (c.sinkParent || c.sink) && !c.sourceParent } map { case (i, c) => s"$i.${portNames(c.name)}" }
+          val from = s"${portNames(name)}"
+          l.copy(cons = l.cons ++ tos.map(t => (t, from)))
+        case l => l
+      })
+    )
+    funs.foldLeft(lin)((l, fun) => fun(l))
+  }
+
+  /** Return a map from module to its lineage in the tree
+    */
+  def pointToLineage(lin: Lineage): Map[String, Lineage] = {
+    val map = mutable.HashMap[String, Lineage]()
+    def onLineage(l: Lineage): Lineage = {
+      map(l.name) = l
+      l map onLineage
+    }
+    onLineage(lin)
+    map.toMap
+  }
+}