TopWiring Transform (#798)

* top wiring transform

* fixup comments

* TopWiring cosmetics

* move prefix into TopWiringAnnotation

* remove test function from transform file

* add ChildrenMap to InstanceGraph API

* use namespaces

* remove wiringUtils from TopWiring pass

* enable multiple output functions

* TopWiring cosmetics, tests and lowform

2 files changed, 278 insertions, 1 deletions

diff --git a/src/main/scala/firrtl/analyses/InstanceGraph.scala b/src/main/scala/firrtl/analyses/InstanceGraph.scala
index c6e590af..91cf8d44 100644
--- a/src/main/scala/firrtl/analyses/InstanceGraph.scala
+++ b/src/main/scala/firrtl/analyses/InstanceGraph.scala
@@ -18,7 +18,7 @@ import firrtl.Mappers._
   */
 class InstanceGraph(c: Circuit) {
 
-  private val moduleMap = c.modules.map({m => (m.name,m) }).toMap
+  val moduleMap = c.modules.map({m => (m.name,m) }).toMap
   private val instantiated = new mutable.HashSet[String]
   private val childInstances =
     new mutable.HashMap[String,mutable.Set[WDefInstance]]
@@ -93,6 +93,14 @@ class InstanceGraph(c: Circuit) {
   def moduleOrder: Seq[DefModule] = {
     graph.transformNodes(_.module).linearize.map(moduleMap(_))
   }
+
+  
+  /** Given a circuit, returns a map from module name to children
+     * instance/module definitions
+     */
+  def getChildrenInstances: scala.collection.Map[String,mutable.Set[WDefInstance]] = childInstances 
+
+
 }
 
 object InstanceGraph {
diff --git a/src/main/scala/firrtl/transforms/TopWiring.scala b/src/main/scala/firrtl/transforms/TopWiring.scala
new file mode 100644
index 00000000..b698cda6
--- /dev/null
+++ b/src/main/scala/firrtl/transforms/TopWiring.scala
@@ -0,0 +1,269 @@
+// See LICENSE for license details.
+package firrtl.transform
+package TopWiring
+
+import firrtl._
+import firrtl.ir._
+import firrtl.passes.{Pass,
+      InferTypes,
+      ResolveKinds,
+      ResolveGenders
+      }
+import firrtl.annotations._
+import firrtl.Mappers._
+import firrtl.graph._
+
+import java.io._
+import scala.io.Source
+import collection.mutable
+
+/** Annotation for optional output files, and what directory to put those files in (absolute path) **/
+case class TopWiringOutputFilesAnnotation(dirName: String,
+                                          outputFunction: (String,Seq[((ComponentName, Type, Boolean,
+                                                                         Seq[String],String), Int)],
+                                                           CircuitState) => CircuitState) extends NoTargetAnnotation
+
+/** Annotation for indicating component to be wired, and what prefix to add to the ports that are generated */
+case class TopWiringAnnotation(target: ComponentName, prefix: String) extends
+    SingleTargetAnnotation[ComponentName] {
+  def duplicate(n: ComponentName) = this.copy(target = n)
+}
+
+
+/** Punch out annotated ports out to the toplevel of the circuit.
+    This also has an option to pass a function as a parmeter to generate
+    custom output files as a result of the additional ports
+  * @note This *does* work for deduped modules
+  */
+class TopWiringTransform extends Transform {
+  def inputForm: CircuitForm = LowForm
+  def outputForm: CircuitForm = LowForm
+
+  type InstPath = Seq[String]
+
+  /** Get the names of the targets that need to be wired */
+  private def getSourceNames(state: CircuitState): Map[ComponentName, String] = {
+    state.annotations.collect { case TopWiringAnnotation(srcname,prefix) =>
+                                  (srcname -> prefix) }.toMap.withDefaultValue("")
+  }
+
+
+  /** Get the names of the modules which include the  targets that need to be wired */
+  private def getSourceModNames(state: CircuitState): Seq[String] = {
+    state.annotations.collect { case TopWiringAnnotation(ComponentName(_,ModuleName(srcmodname, _)),_) => srcmodname }
+  }
+
+
+
+  /** Get the Type of each wire to be connected
+    *
+    * Find the definition of each wire in sourceList, and get the type and whether or not it's a port
+    * Update the results in sourceMap
+    */
+  private def getSourceTypes(sourceList: Map[ComponentName, String],
+                     sourceMap: mutable.Map[String, Seq[(ComponentName, Type, Boolean, InstPath, String)]],
+                     currentmodule: ModuleName, state: CircuitState)(s: Statement): Statement = s match {
+    // If target wire, add name and size to to sourceMap
+    case w: IsDeclaration =>
+      if (sourceList.keys.toSeq.contains(ComponentName(w.name, currentmodule))) {
+          val (isport, tpe, prefix) = w match {
+            case d: DefWire => (false, d.tpe, sourceList(ComponentName(w.name,currentmodule)))
+            case d: DefNode => (false, d.value.tpe, sourceList(ComponentName(w.name,currentmodule)))
+            case d: DefRegister => (false, d.tpe, sourceList(ComponentName(w.name,currentmodule)))
+            case d: Port => (true, d.tpe, sourceList(ComponentName(w.name,currentmodule)))
+            case _ => throw new Exception(s"Cannot wire this type of declaration! ${w.serialize}")
+          }
+          val name = w.name
+          sourceMap.get(currentmodule.name) match {
+            case Some(xs:Seq[(ComponentName, Type, Boolean, InstPath, String)]) =>
+              sourceMap.update(currentmodule.name, xs :+
+                 (ComponentName(w.name,currentmodule), tpe, isport ,Seq[String](w.name), prefix))
+            case None =>
+              sourceMap(currentmodule.name) = Seq((ComponentName(w.name,currentmodule),
+                                                   tpe, isport ,Seq[String](w.name), prefix))
+          }
+      }
+      w // Return argument unchanged (ok because DefWire has no Statement children)
+    // If not, apply to all children Statement
+    case _ => s map getSourceTypes(sourceList, sourceMap, currentmodule, state)
+  }
+
+
+
+  /** Get the Type of each port to be connected
+    *
+    * Similar to getSourceTypes, but specifically for ports since they are not found in statements.
+    * Find the definition of each port in sourceList, and get the type and whether or not it's a port
+    * Update the results in sourceMap
+    */
+  private def getSourceTypesPorts(sourceList: Map[ComponentName, String], sourceMap: mutable.Map[String,
+                          Seq[(ComponentName, Type, Boolean, InstPath, String)]],
+                          currentmodule: ModuleName, state: CircuitState)(s: Port): CircuitState = s match {
+    // If target port, add name and size to to sourceMap
+    case w: IsDeclaration =>
+      if (sourceList.keys.toSeq.contains(ComponentName(w.name, currentmodule))) {
+          val (isport, tpe, prefix) = w match {
+            case d: Port => (true, d.tpe, sourceList(ComponentName(w.name,currentmodule)))
+            case _ => throw new Exception(s"Cannot wire this type of declaration! ${w.serialize}")
+          }
+          val name = w.name
+          sourceMap.get(currentmodule.name) match {
+            case Some(xs:Seq[(ComponentName, Type, Boolean, InstPath, String)]) =>
+                sourceMap.update(currentmodule.name, xs :+
+                  (ComponentName(w.name,currentmodule), tpe, isport ,Seq[String](w.name), prefix))
+            case None =>
+                sourceMap(currentmodule.name) = Seq((ComponentName(w.name,currentmodule),
+                                                     tpe, isport ,Seq[String](w.name), prefix))
+          }
+      }
+      state // Return argument unchanged (ok because DefWire has no Statement children)
+    // If not, apply to all children Statement
+    case _ => state
+  }
+
+
+  /** Create a map of Module name to target wires under this module
+    *
+    * These paths are relative but cross module (they refer down through instance hierarchy)
+    */
+  private def getSourcesMap(state: CircuitState): Map[String,Seq[(ComponentName, Type, Boolean, InstPath, String)]] = {
+    val sSourcesModNames = getSourceModNames(state)
+    val sSourcesNames = getSourceNames(state)
+    val instGraph = new firrtl.analyses.InstanceGraph(state.circuit)
+    val cMap = instGraph.getChildrenInstances.map{ case (m, wdis) =>
+        (m -> wdis.map{ case wdi => (wdi.name, wdi.module) }.toSeq) }.toMap
+    val topSort = instGraph.moduleOrder.reverse
+
+    // Map of component name to relative instance paths that result in a debug wire
+    val sourcemods: mutable.Map[String, Seq[(ComponentName, Type, Boolean, InstPath, String)]] =
+      mutable.Map(sSourcesModNames.map(_ -> Seq()): _*)
+
+    state.circuit.modules.foreach { m => m map
+      getSourceTypes(sSourcesNames, sourcemods, ModuleName(m.name, CircuitName(state.circuit.main)) , state) }
+    state.circuit.modules.foreach { m => m.ports.foreach {
+       p => Seq(p) map
+          getSourceTypesPorts(sSourcesNames, sourcemods, ModuleName(m.name, CircuitName(state.circuit.main)) , state) }}
+
+    for (mod <- topSort) {
+      val seqChildren: Seq[(ComponentName,Type,Boolean,InstPath,String)] = cMap(mod.name).flatMap {
+        case (inst, module) =>
+          sourcemods.get(module).map( _.map { case (a,b,c,path,p) => (a,b,c, inst +: path, p)})
+      }.flatten
+      if (seqChildren.nonEmpty) {
+        sourcemods(mod.name) = seqChildren
+      }
+    }
+
+    sourcemods.toMap
+  }
+
+
+
+  /** Process a given DefModule
+    *
+    * For Modules that contain or are in the parent hierarchy to modules containing target wires
+    * 1. Add ports for each target wire this module is parent to
+    * 2. Connect these ports to ports of instances that are parents to some number of target wires
+    */
+  private def onModule(sources: Map[String, Seq[(ComponentName, Type, Boolean, InstPath, String)]],
+                       portnamesmap : mutable.Map[String,String],
+                       instgraph : firrtl.analyses.InstanceGraph,
+                       namespacemap : Map[String, Namespace])
+                      (module: DefModule): DefModule = {
+    val namespace = namespacemap(module.name)
+    sources.get(module.name) match {
+      case Some(p) =>
+        val newPorts = p.map{ case (ComponentName(cname,_), tpe, _ , path, prefix) => {
+              val newportname = portnamesmap.get(prefix + path.mkString("_")) match {
+                case Some(pn) => pn
+                 case None => {
+                    val npn = namespace.newName(prefix + path.mkString("_"))
+                    portnamesmap(prefix + path.mkString("_")) = npn
+                    npn
+                 }
+              }
+              Port(NoInfo, newportname, Output, tpe)
+        } }
+
+        // Add connections to Module
+        module match {
+          case m: Module =>
+            val connections: Seq[Connect] = p.map { case (ComponentName(cname,_), _, _ , path, prefix) =>
+                val modRef = portnamesmap.get(prefix + path.mkString("_")) match {
+                       case Some(pn) => WRef(pn)
+                       case None => {
+                          portnamesmap(prefix + path.mkString("_")) = namespace.newName(prefix + path.mkString("_"))
+                          WRef(portnamesmap(prefix + path.mkString("_")))
+                       }
+                }
+                path.size match {
+                   case 1 => {
+                       val leafRef = WRef(path.head.mkString("_"))
+                       Connect(NoInfo, modRef, leafRef)
+                   }
+                   case _ =>  {
+                       val instportname =  portnamesmap.get(prefix + path.tail.mkString("_")) match {
+                           case Some(ipn) => ipn
+                           case None => {
+                             val instmod = instgraph.getChildrenInstances(module.name).collectFirst {
+                                 case wdi if wdi.name == path.head => wdi.module}.get
+                             val instnamespace = namespacemap(instmod)
+                             portnamesmap(prefix + path.tail.mkString("_")) =
+                               instnamespace.newName(prefix + path.tail.mkString("_"))
+                             portnamesmap(prefix + path.tail.mkString("_"))
+                           }
+                       }
+                       val instRef = WSubField(WRef(path.head), instportname)
+                       Connect(NoInfo, modRef, instRef)
+                  }
+                }
+            }
+            m.copy(ports = m.ports ++ newPorts, body = Block(Seq(m.body) ++ connections ))
+          case e: ExtModule =>
+            e.copy(ports = e.ports ++ newPorts)
+      }
+      case None => module // unchanged if no paths
+    }
+  }
+
+  /** Run passes to fix up the circuit of making the new connections  */
+  private def fixupCircuit(circuit: Circuit): Circuit = {
+    val passes = Seq(
+      InferTypes,
+      ResolveKinds,
+      ResolveGenders
+    )
+    passes.foldLeft(circuit) { case (c: Circuit, p: Pass) => p.run(c) }
+  }
+
+
+  /** Dummy function that is currently unused. Can be used to fill an outputFunction requirment in the future  */
+  def topWiringDummyOutputFilesFunction(dir: String,
+                                        mapping: Seq[((ComponentName, Type, Boolean, InstPath, String), Int)],
+                                        state: CircuitState): CircuitState = {
+     state
+  }
+
+
+  def execute(state: CircuitState): CircuitState = {
+
+    val outputTuples: Seq[(String,
+                          (String,Seq[((ComponentName, Type, Boolean, InstPath, String), Int)],
+                                        CircuitState) => CircuitState)] = state.annotations.collect {
+         case TopWiringOutputFilesAnnotation(td,of) => (td, of) }
+
+    // Do actual work of this transform
+    val sources = getSourcesMap(state)
+    val portnamesmap : mutable.Map[String,String] = mutable.Map()
+    val instgraph = new firrtl.analyses.InstanceGraph(state.circuit)
+    val namespacemap = state.circuit.modules.map{ case m => (m.name -> Namespace(m)) }.toMap
+    val modulesx = state.circuit.modules map onModule(sources, portnamesmap, instgraph, namespacemap)
+    val newCircuit = state.circuit.copy(modules = modulesx)
+    val fixedCircuit = fixupCircuit(newCircuit)
+    val mappings = sources(state.circuit.main).zipWithIndex
+    //Generate output files based on the mapping.
+    outputTuples.map { case (dir, outputfunction) => outputfunction(dir, mappings, state) }
+    // fin.
+    state.copy(circuit = fixedCircuit)
+  }
+}