Instance Annotations (#865)

Added Target, which now supports Instance Annotations. See #865 for details.

5 files changed, 200 insertions, 94 deletions

diff --git a/src/main/scala/firrtl/transforms/CheckCombLoops.scala b/src/main/scala/firrtl/transforms/CheckCombLoops.scala
index 98033a2f..44785c62 100644
--- a/src/main/scala/firrtl/transforms/CheckCombLoops.scala
+++ b/src/main/scala/firrtl/transforms/CheckCombLoops.scala
@@ -26,9 +26,9 @@ case object DontCheckCombLoopsAnnotation extends NoTargetAnnotation
 
 case class CombinationalPath(sink: ComponentName, sources: Seq[ComponentName]) extends Annotation {
   override def update(renames: RenameMap): Seq[Annotation] = {
-    val newSources = sources.flatMap { s => renames.get(s).getOrElse(Seq(s)) }
-    val newSinks = renames.get(sink).getOrElse(Seq(sink))
-    newSinks.map(snk => CombinationalPath(snk, newSources))
+    val newSources: Seq[IsComponent] = sources.flatMap { s => renames.get(s).getOrElse(Seq(s.toTarget)) }.collect {case x: IsComponent if x.isLocal => x}
+    val newSinks = renames.get(sink).getOrElse(Seq(sink.toTarget)).collect { case x: IsComponent if x.isLocal => x}
+    newSinks.map(snk => CombinationalPath(snk.toNamed, newSources.map(_.toNamed)))
   }
 }
 
diff --git a/src/main/scala/firrtl/transforms/ConstantPropagation.scala b/src/main/scala/firrtl/transforms/ConstantPropagation.scala
index 0d30446c..da7f1a46 100644
--- a/src/main/scala/firrtl/transforms/ConstantPropagation.scala
+++ b/src/main/scala/firrtl/transforms/ConstantPropagation.scala
@@ -12,6 +12,7 @@ import firrtl.PrimOps._
 import firrtl.graph.DiGraph
 import firrtl.WrappedExpression.weq
 import firrtl.analyses.InstanceGraph
+import firrtl.annotations.TargetToken.Ref
 
 import annotation.tailrec
 import collection.mutable
@@ -46,11 +47,13 @@ object ConstantPropagation {
   }
 }
 
-class ConstantPropagation extends Transform {
+class ConstantPropagation extends Transform with ResolvedAnnotationPaths {
   import ConstantPropagation._
   def inputForm = LowForm
   def outputForm = LowForm
 
+  override val annotationClasses: Traversable[Class[_]] = Seq(classOf[DontTouchAnnotation])
+
   trait FoldCommutativeOp {
     def fold(c1: Literal, c2: Literal): Expression
     def simplify(e: Expression, lhs: Literal, rhs: Expression): Expression
@@ -520,7 +523,7 @@ class ConstantPropagation extends Transform {
 
   def execute(state: CircuitState): CircuitState = {
     val dontTouches: Seq[(String, String)] = state.annotations.collect {
-      case DontTouchAnnotation(ComponentName(c, ModuleName(m, _))) => m -> c
+      case DontTouchAnnotation(Target(_, Some(m), Seq(Ref(c)))) => m -> c
     }
     // Map from module name to component names
     val dontTouchMap: Map[String, Set[String]] =
diff --git a/src/main/scala/firrtl/transforms/DeadCodeElimination.scala b/src/main/scala/firrtl/transforms/DeadCodeElimination.scala
index c98b892c..523c997b 100644
--- a/src/main/scala/firrtl/transforms/DeadCodeElimination.scala
+++ b/src/main/scala/firrtl/transforms/DeadCodeElimination.scala
@@ -30,7 +30,7 @@ import java.io.{File, FileWriter}
   * circumstances of their instantiation in their parent module, they will still not be removed. To
   * remove such modules, use the [[NoDedupAnnotation]] to prevent deduplication.
   */
-class DeadCodeElimination extends Transform {
+class DeadCodeElimination extends Transform with ResolvedAnnotationPaths {
   def inputForm = LowForm
   def outputForm = LowForm
 
@@ -321,9 +321,12 @@ class DeadCodeElimination extends Transform {
     state.copy(circuit = newCircuit, renames = Some(renames))
   }
 
+  override val annotationClasses: Traversable[Class[_]] =
+    Seq(classOf[DontTouchAnnotation], classOf[OptimizableExtModuleAnnotation])
+
   def execute(state: CircuitState): CircuitState = {
     val dontTouches: Seq[LogicNode] = state.annotations.collect {
-      case DontTouchAnnotation(component) => LogicNode(component)
+      case DontTouchAnnotation(component: ReferenceTarget) if component.isLocal => LogicNode(component)
     }
     val doTouchExtMods: Seq[String] = state.annotations.collect {
       case OptimizableExtModuleAnnotation(ModuleName(name, _)) => name
diff --git a/src/main/scala/firrtl/transforms/Dedup.scala b/src/main/scala/firrtl/transforms/Dedup.scala
index 5630cecf..1c20b448 100644
--- a/src/main/scala/firrtl/transforms/Dedup.scala
+++ b/src/main/scala/firrtl/transforms/Dedup.scala
@@ -6,17 +6,18 @@ package transforms
 import firrtl.ir._
 import firrtl.Mappers._
 import firrtl.analyses.InstanceGraph
+import firrtl.annotations.TargetToken.{Instance, OfModule, Ref}
 import firrtl.annotations._
 import firrtl.passes.{InferTypes, MemPortUtils}
+import firrtl.Utils.throwInternalError
 
 // Datastructures
 import scala.collection.mutable
 
 
-/** A component, e.g. register etc. Must be declared only once under the TopAnnotation
-  */
+/** A component, e.g. register etc. Must be declared only once under the TopAnnotation */
 case class NoDedupAnnotation(target: ModuleName) extends SingleTargetAnnotation[ModuleName] {
-  def duplicate(n: ModuleName) = NoDedupAnnotation(n)
+  def duplicate(n: ModuleName): NoDedupAnnotation = NoDedupAnnotation(n)
 }
 
 /** Only use on legal Firrtl.
@@ -28,62 +29,63 @@ class DedupModules extends Transform {
   def inputForm: CircuitForm = HighForm
   def outputForm: CircuitForm = HighForm
 
-  /**
-    * Deduplicate a Circuit
+  /** Deduplicate a Circuit
     * @param state Input Firrtl AST
     * @return A transformed Firrtl AST
     */
   def execute(state: CircuitState): CircuitState = {
     val noDedups = state.annotations.collect { case NoDedupAnnotation(ModuleName(m, c)) => m }
-    val (newC, renameMap) = run(state.circuit, noDedups)
+    val (newC, renameMap) = run(state.circuit, noDedups, state.annotations)
     state.copy(circuit = newC, renames = Some(renameMap))
   }
 
-  /**
-    * Deduplicates a circuit, and records renaming
+  /** Deduplicates a circuit, and records renaming
     * @param c Circuit to dedup
     * @param noDedups Modules not to dedup
     * @return Deduped Circuit and corresponding RenameMap
     */
-  def run(c: Circuit, noDedups: Seq[String]): (Circuit, RenameMap) = {
+  def run(c: Circuit, noDedups: Seq[String], annos: Seq[Annotation]): (Circuit, RenameMap) = {
 
     // RenameMap
     val renameMap = RenameMap()
     renameMap.setCircuit(c.main)
 
     // Maps module name to corresponding dedup module
-    val dedupMap = DedupModules.deduplicate(c, noDedups.toSet, renameMap)
+    val dedupMap = DedupModules.deduplicate(c, noDedups.toSet, annos, renameMap)
 
     // Use old module list to preserve ordering
     val dedupedModules = c.modules.map(m => dedupMap(m.name)).distinct
 
     val cname = CircuitName(c.main)
-    renameMap.addMap(dedupMap.map { case (from, to) =>
+    val map = dedupMap.map { case (from, to) =>
       logger.debug(s"[Dedup] $from -> ${to.name}")
       ModuleName(from, cname) -> List(ModuleName(to.name, cname))
-    })
+    }
+    renameMap.recordAll(
+      map.map {
+        case (k: ModuleName, v: List[ModuleName]) => Target.convertNamed2Target(k) -> v.map(Target.convertNamed2Target)
+      }
+    )
 
     (InferTypes.run(c.copy(modules = dedupedModules)), renameMap)
   }
 }
 
-/**
-  * Utility functions for [[DedupModules]]
-  */
+/** Utility functions for [[DedupModules]] */
 object DedupModules {
-  /**
-    * Change's a module's internal signal names, types, infos, and modules.
+
+  /** Change's a module's internal signal names, types, infos, and modules.
     * @param rename Function to rename a signal. Called on declaration and references.
     * @param retype Function to retype a signal. Called on declaration, references, and subfields
     * @param reinfo Function to re-info a statement
-    * @param renameModule Function to rename an instance's module
+    * @param renameOfModule Function to rename an instance's module
     * @param module Module to change internals
     * @return Changed Module
     */
   def changeInternals(rename: String=>String,
                       retype: String=>Type=>Type,
                       reinfo: Info=>Info,
-                      renameModule: String=>String
+                      renameOfModule: (String, String)=>String
                      )(module: DefModule): DefModule = {
     def onPort(p: Port): Port = Port(reinfo(p.info), rename(p.name), p.direction, retype(p.name)(p.tpe))
     def onExp(e: Expression): Expression = e match {
@@ -99,9 +101,9 @@ object DedupModules {
     }
     def onStmt(s: Statement): Statement = s match {
       case WDefInstance(i, n, m, t) =>
-        val newmod = renameModule(m)
+        val newmod = renameOfModule(n, m)
         WDefInstance(reinfo(i), rename(n), newmod, retype(n)(t))
-      case DefInstance(i, n, m) => DefInstance(reinfo(i), rename(n), renameModule(m))
+      case DefInstance(i, n, m) => DefInstance(reinfo(i), rename(n), renameOfModule(n, m))
       case d: DefMemory =>
         val oldType = MemPortUtils.memType(d)
         val newType = retype(d.name)(oldType)
@@ -129,49 +131,79 @@ object DedupModules {
     module map onPort map onStmt
   }
 
-  /**
-    * Turns a module into a name-agnostic module
+  def uniquifyField(ref: String, depth: Int, field: String): String = ref + depth + field
+
+  /** Turns a module into a name-agnostic module
     * @param module module to change
     * @return name-agnostic module
     */
-  def agnostify(module: DefModule, name2tag: mutable.HashMap[String, String], tag2name: mutable.HashMap[String, String]): DefModule = {
+  def agnostify(top: CircuitTarget,
+                module: DefModule,
+                renameMap: RenameMap
+               ): DefModule = {
+
+
     val namespace = Namespace()
-    val nameMap = mutable.HashMap[String, String]()
     val typeMap = mutable.HashMap[String, Type]()
+
+    renameMap.setCircuit(top.circuitOpt.get)
+    renameMap.setModule(module.name)
+
     def rename(name: String): String = {
-      if (nameMap.contains(name)) nameMap(name) else {
-        val newName = namespace.newTemp
-        nameMap(name) = newName
-        newName
+      val ret = renameMap.get(top.module(module.name).ref(name))
+      ret match {
+        case Some(Seq(Target(_, _, Seq(Ref(x))))) => x
+        case None =>
+          val newName = namespace.newTemp
+          renameMap.rename(name, newName)
+          newName
+        case other => throwInternalError(other.toString)
       }
     }
+
     def retype(name: String)(tpe: Type): Type = {
       if (typeMap.contains(name)) typeMap(name) else {
-        def onType(tpe: Type): Type = tpe map onType match {
-          case BundleType(fields) => BundleType(fields.map(f => Field(rename(f.name), f.flip, f.tpe)))
+        def onType(depth: Int)(tpe: Type): Type = tpe map onType(depth + 1) match {
+          //TODO bugfix: ref.data.data and ref.datax.data will not rename to the right tags, even if they should be
+          case BundleType(fields) =>
+            BundleType(fields.map(f => Field(rename(uniquifyField(name, depth, f.name)), f.flip, f.tpe)))
           case other => other
         }
-        val newType = onType(tpe)
+        val newType = onType(0)(tpe)
         typeMap(name) = newType
         newType
       }
     }
-    def remodule(name: String): String = tag2name(name2tag(name))
-    changeInternals(rename, retype, {i: Info => NoInfo}, remodule)(module)
+
+    def reOfModule(instance: String, ofModule: String): String = {
+      renameMap.get(top.module(ofModule)) match {
+        case Some(Seq(Target(_, Some(ofModuleTag), Nil))) => ofModuleTag
+        case None => ofModule
+        case other => throwInternalError(other.toString)
+      }
+    }
+
+    val renamedModule = changeInternals(rename, retype, {i: Info => NoInfo}, reOfModule)(module)
+    renamedModule
   }
 
   /** Dedup a module's instances based on dedup map
     *
     * Will fixes up module if deduped instance's ports are differently named
     *
-    * @param moduleName Module name who's instances will be deduped
+    * @param top CircuitTarget of circuit
+    * @param originalModule Module name who's instances will be deduped
     * @param moduleMap Map of module name to its original module
     * @param name2name Map of module name to the module deduping it. Not mutated in this function.
     * @param renameMap Will be modified to keep track of renames in this function
     * @return fixed up module deduped instances
     */
-  def dedupInstances(moduleName: String, moduleMap: Map[String, DefModule], name2name: mutable.Map[String, String], renameMap: RenameMap): DefModule = {
-    val module = moduleMap(moduleName)
+  def dedupInstances(top: CircuitTarget,
+                     originalModule: String,
+                     moduleMap: Map[String, DefModule],
+                     name2name: Map[String, String],
+                     renameMap: RenameMap): DefModule = {
+    val module = moduleMap(originalModule)
 
     // If black box, return it (it has no instances)
     if (module.isInstanceOf[ExtModule]) return module
@@ -187,7 +219,14 @@ object DedupModules {
       moduleMap(name2name(old))
     }
     // Define rename functions
-    def renameModule(name: String): String = getNewModule(name).name
+    def renameOfModule(instance: String, ofModule: String): String = {
+      val newOfModule = name2name(ofModule)
+      renameMap.record(
+        top.module(originalModule).instOf(instance, ofModule),
+        top.module(originalModule).instOf(instance, newOfModule)
+      )
+      newOfModule
+    }
     val typeMap = mutable.HashMap[String, Type]()
     def retype(name: String)(tpe: Type): Type = {
       if (typeMap.contains(name)) typeMap(name) else {
@@ -198,96 +237,155 @@ object DedupModules {
             case (old, nuu) => renameMap.rename(old.serialize, nuu.serialize)
           }
           newType
-        } else tpe
+        } else {
+          tpe
+        }
       }
     }
 
     renameMap.setModule(module.name)
     // Change module internals
-    changeInternals({n => n}, retype, {i => i}, renameModule)(module)
+    changeInternals({n => n}, retype, {i => i}, renameOfModule)(module)
   }
 
-  /**
-    * Deduplicate
-    * @param circuit Circuit
-    * @param noDedups list of modules to not dedup
-    * @param renameMap rename map to populate when deduping
-    * @return Map of original Module name -> Deduped Module
+  //scalastyle:off
+  /** Returns
+    *  1) map of tag to all matching module names,
+    *  2) renameMap of module name to tag (agnostic name)
+    *  3) maps module name to agnostic renameMap
+    * @param top CircuitTarget
+    * @param moduleLinearization Sequence of modules from leaf to top
+    * @param noDedups Set of modules to not dedup
+    * @param annotations All annotations to check if annotations are identical
+    * @return
     */
-  def deduplicate(circuit: Circuit,
-                  noDedups: Set[String],
-                  renameMap: RenameMap): Map[String, DefModule] = {
-
-    // Order of modules, from leaf to top
-    val moduleLinearization = new InstanceGraph(circuit).moduleOrder.map(_.name).reverse
+  def buildRTLTags(top: CircuitTarget,
+                   moduleLinearization: Seq[DefModule],
+                   noDedups: Set[String],
+                   annotations: Seq[Annotation]
+                  ): (collection.Map[String, collection.Set[String]], RenameMap, collection.Map[String, RenameMap]) = {
 
-    // Maps module name to original module
-    val moduleMap = circuit.modules.map(m => m.name -> m).toMap
-
-    // Maps a module's tag to its deduplicated module
-    val tag2name = mutable.HashMap.empty[String, String]
 
-    // Maps a module's name to its tag
-    val name2tag = mutable.HashMap.empty[String, String]
+    // Maps a module name to its agnostic name
+    val tagMap = RenameMap()
 
     // Maps a tag to all matching module names
-    val tag2all =  mutable.HashMap.empty[String, mutable.Set[String]]
+    val tag2all = mutable.HashMap.empty[String, mutable.HashSet[String]]
 
-    // Build dedupMap
-    moduleLinearization.foreach { moduleName =>
-      // Get original module
-      val originalModule = moduleMap(moduleName)
+    val module2Annotations = mutable.HashMap.empty[String, mutable.HashSet[Annotation]]
+    annotations.foreach { a =>
+      a.getTargets.foreach { t =>
+        val annos = module2Annotations.getOrElseUpdate(t.moduleOpt.get, mutable.HashSet.empty[Annotation])
+        annos += a
+      }
+    }
 
+    val agnosticModuleMap = RenameMap()
+    val agnosticRenames = mutable.HashMap[String, RenameMap]()
+
+    moduleLinearization.foreach { originalModule =>
       // Replace instance references to new deduped modules
       val dontcare = RenameMap()
       dontcare.setCircuit("dontcare")
-      //val fixedModule = DedupModules.dedupInstances(originalModule, tag2module, name2tag, name2module, dontcare)
+
+      val agnosticRename = RenameMap.create(agnosticModuleMap.getUnderlying)
+      agnosticRenames(originalModule.name) = agnosticRename
 
       if (noDedups.contains(originalModule.name)) {
         // Don't dedup. Set dedup module to be the same as fixed module
-        name2tag(originalModule.name) = originalModule.name
-        tag2name(originalModule.name) = originalModule.name
-        //templateModules += originalModule.name
+        tag2all(originalModule.name) = mutable.HashSet(originalModule.name)
       } else { // Try to dedup
 
         // Build name-agnostic module
-        val agnosticModule = DedupModules.agnostify(originalModule, name2tag, tag2name)
+        val agnosticModule = DedupModules.agnostify(top, originalModule, agnosticRename)
+        agnosticRename.record(top.module(originalModule.name), top.module("thisModule"))
+        val agnosticAnnos = module2Annotations.getOrElse(
+          originalModule.name, mutable.HashSet.empty[Annotation]
+        ).map(_.update(agnosticRename))
+        agnosticRename.delete(top.module(originalModule.name))
 
         // Build tag
-        val tag = (agnosticModule match {
-          case Module(i, n, ps, b) =>
-            ps.map(_.serialize).mkString + b.serialize
+        val builder = new mutable.ArrayBuffer[Any]()
+        agnosticModule.ports.foreach { builder ++= _.serialize }
+        builder ++= agnosticAnnos
+
+        agnosticModule match {
+          case Module(i, n, ps, b) => builder ++= b.serialize
           case ExtModule(i, n, ps, dn, p) =>
-            ps.map(_.serialize).mkString + dn + p.map(_.serialize).mkString
-        }).hashCode().toString
+            builder ++= dn
+            p.foreach { builder ++= _.serialize }
+        }
+        val tag = builder.hashCode().toString
 
         // Match old module name to its tag
-        name2tag(originalModule.name) = tag
+        agnosticRename.record(top.module(originalModule.name), top.module(tag))
+        agnosticModuleMap.record(top.module(originalModule.name), top.module(tag))
+        tagMap.record(top.module(originalModule.name), top.module(tag))
 
         // Set tag's module to be the first matching module
-        if (!tag2name.contains(tag)) {
-          tag2name(tag) = originalModule.name
-          tag2all(tag) = mutable.Set(originalModule.name)
-        } else {
-          tag2all(tag) += originalModule.name
-        }
+        val all = tag2all.getOrElseUpdate(tag, mutable.HashSet.empty[String])
+        all += originalModule.name
       }
     }
+    (tag2all, tagMap, agnosticRenames)
+  }
+  //scalastyle:on
 
+  /** Deduplicate
+    * @param circuit Circuit
+    * @param noDedups list of modules to not dedup
+    * @param renameMap rename map to populate when deduping
+    * @return Map of original Module name -> Deduped Module
+    */
+  def deduplicate(circuit: Circuit,
+                  noDedups: Set[String],
+                  annotations: Seq[Annotation],
+                  renameMap: RenameMap): Map[String, DefModule] = {
+
+    val (moduleMap, moduleLinearization) = {
+      val iGraph = new InstanceGraph(circuit)
+      (iGraph.moduleMap, iGraph.moduleOrder.reverse)
+    }
+    val top = CircuitTarget(circuit.main)
+
+    val (tag2all, tagMap, agnosticRenames) = buildRTLTags(top, moduleLinearization, noDedups, annotations)
 
     // Set tag2name to be the best dedup module name
     val moduleIndex = circuit.modules.zipWithIndex.map{case (m, i) => m.name -> i}.toMap
     def order(l: String, r: String): String = if (moduleIndex(l) < moduleIndex(r)) l else r
+
+    // Maps a module's tag to its deduplicated module
+    val tag2name = mutable.HashMap.empty[String, String]
     tag2all.foreach { case (tag, all) => tag2name(tag) = all.reduce(order)}
 
     // Create map from original to dedup name
-    val name2name = name2tag.map({ case (name, tag) => name -> tag2name(tag) })
+    val name2name = moduleMap.keysIterator.map{ originalModule =>
+      tagMap.get(top.module(originalModule)) match {
+        case Some(Seq(Target(_, Some(tag), Nil))) => originalModule -> tag2name(tag)
+        case None => originalModule -> originalModule
+        case other => throwInternalError(other.toString)
+      }
+    }.toMap
 
     // Build Remap for modules with deduped module references
-    val tag2module = tag2name.map({ case (tag, name) => tag -> DedupModules.dedupInstances(name, moduleMap, name2name, renameMap) })
+    val dedupedName2module = tag2name.map({ case (tag, name) => name -> DedupModules.dedupInstances(top, name, moduleMap, name2name, renameMap) })
 
     // Build map from original name to corresponding deduped module
-    val name2module = name2tag.map({ case (name, tag) => name -> tag2module(tag) })
+    val name2module = tag2all.flatMap({ case (tag, names) => names.map(n => n -> dedupedName2module(tag2name(tag))) })
+
+    val reversedAgnosticRenames = mutable.HashMap[String, RenameMap]()
+    name2module.foreach { case (originalModuleName, dedupedModule) =>
+      if(!reversedAgnosticRenames.contains(dedupedModule.name)) {
+        reversedAgnosticRenames(dedupedModule.name) = agnosticRenames(dedupedModule.name).getReverseRenameMap
+      }
+      agnosticRenames(originalModuleName).keys.foreach { key =>
+        if(key.isInstanceOf[IsComponent]) {
+          val tag = agnosticRenames(originalModuleName)(key).head
+          val newKey = reversedAgnosticRenames(dedupedModule.name).apply(tag)
+          renameMap.record(key.asInstanceOf[IsMember], newKey.asInstanceOf[Seq[IsMember]])
+        }
+      }
+    }
 
     name2module.toMap
   }
diff --git a/src/main/scala/firrtl/transforms/OptimizationAnnotations.scala b/src/main/scala/firrtl/transforms/OptimizationAnnotations.scala
index fab540da..a66bd4ce 100644
--- a/src/main/scala/firrtl/transforms/OptimizationAnnotations.scala
+++ b/src/main/scala/firrtl/transforms/OptimizationAnnotations.scala
@@ -4,6 +4,7 @@ package transforms
 
 import firrtl.annotations._
 import firrtl.passes.PassException
+import firrtl.transforms
 
 /** Indicate that DCE should not be run */
 case object NoDCEAnnotation extends NoTargetAnnotation
@@ -12,13 +13,14 @@ case object NoDCEAnnotation extends NoTargetAnnotation
   *
   * DCE treats the component as a top-level sink of the circuit
   */
-case class DontTouchAnnotation(target: ComponentName) extends SingleTargetAnnotation[ComponentName] {
-  def duplicate(n: ComponentName) = this.copy(n)
+case class DontTouchAnnotation(target: ReferenceTarget) extends SingleTargetAnnotation[ReferenceTarget] {
+  def targets = Seq(target)
+  def duplicate(n: ReferenceTarget) = this.copy(n)
 }
 
 object DontTouchAnnotation {
   class DontTouchNotFoundException(module: String, component: String) extends PassException(
-    s"Component marked dontTouch ($module.$component) not found!\n" +
+    s"Target marked dontTouch ($module.$component) not found!\n" +
     "It was probably accidentally deleted. Please check that your custom transforms are not" +
     "responsible and then file an issue on Github."
   )