Fix annotations (#366)

getMyAnnotations now returns Seq[Annotation]
Changed test to check number of annotations is the same

9 files changed, 44 insertions, 64 deletions

diff --git a/src/main/scala/firrtl/Annotations.scala b/src/main/scala/firrtl/Annotations.scala
index a8aecfc5..b502e055 100644
--- a/src/main/scala/firrtl/Annotations.scala
+++ b/src/main/scala/firrtl/Annotations.scala
@@ -232,23 +232,8 @@ object Annotations {
    * Container of all annotations for a Firrtl compiler.
    */
   case class AnnotationMap(annotations: Seq[Annotation]) {
-    type NamedMap = Map[Named, Map[Class[_], Annotation]]
-    type IDMap = Map[Class[_], Map[Named, Annotation]]
-
-    val (namedMap: NamedMap, idMap:IDMap) =
-      //annotations.foldLeft(Tuple2[NamedMap, IDMap](Map.empty, Map.empty)){
-      annotations.foldLeft((Map.empty: NamedMap, Map.empty: IDMap)){
-        (partialMaps: (NamedMap, IDMap), annotation: Annotation) => {
-          val transformToAnn = partialMaps._1.getOrElse(annotation.target, Map.empty)
-          val pNMap = partialMaps._1 + (annotation.target -> (transformToAnn + (annotation.transform -> annotation)))
-
-          val nToAnn = partialMaps._2.getOrElse(annotation.transform, Map.empty)
-          val ptransformMap = partialMaps._2 + (annotation.transform -> (nToAnn + (annotation.target -> annotation)))
-          Tuple2(pNMap, ptransformMap)
-        }
-      }
-    def get(id: Class[_]): Option[Map[Named, Annotation]] = idMap.get(id)
-    def get(named: Named): Option[Map[Class[_], Annotation]] = namedMap.get(named)
+    def get(id: Class[_]): Seq[Annotation] = annotations.filter(a => a.transform == id)
+    def get(named: Named): Seq[Annotation] = annotations.filter(n => n == named)
   }
 }
 
diff --git a/src/main/scala/firrtl/Compiler.scala b/src/main/scala/firrtl/Compiler.scala
index 316c563b..2e155885 100644
--- a/src/main/scala/firrtl/Compiler.scala
+++ b/src/main/scala/firrtl/Compiler.scala
@@ -95,11 +95,10 @@ abstract class Transform {
     * @param state The [[CircuitState]] form which to extract annotations
     * @return A collection of annotations
     */
-  final def getMyAnnotations(state: CircuitState): Option[Map[Named, Annotation]] =
-    for {
-      annotations <- state.annotations
-      myAnnotations <- annotations.get(this.getClass)
-    } yield myAnnotations
+  final def getMyAnnotations(state: CircuitState): Seq[Annotation] = state.annotations match {
+    case Some(annotations) => annotations.get(this.getClass)
+    case None => Nil
+  }
 }
 
 trait SimpleRun extends LazyLogging {
@@ -254,8 +253,7 @@ trait Compiler {
           // annotations with the names in rmap's value.
           for {
             (oldName, newNames) <- rmap.toSeq
-            transform2OldAnnos <- inAnnotationMap.get(oldName).toSeq
-            oldAnno <- transform2OldAnnos.values
+            oldAnno <- inAnnotationMap.get(oldName)
             newAnno <- oldAnno.update(newNames)
           } yield newAnno
         case _ => inAnnotationMap.annotations
diff --git a/src/main/scala/firrtl/passes/Inline.scala b/src/main/scala/firrtl/passes/Inline.scala
index a12e77cb..40bc7d7d 100644
--- a/src/main/scala/firrtl/passes/Inline.scala
+++ b/src/main/scala/firrtl/passes/Inline.scala
@@ -37,11 +37,12 @@ class InlineInstances extends Transform {
    def execute(state: CircuitState): CircuitState = {
      // TODO Add error check for more than one annotation for inlining
      // TODO Propagate other annotations
-     val result = for {
-       myAnnotations <- getMyAnnotations(state)
-       (modNames, instNames) = collectAnns(myAnnotations.values)
-     } yield run(state.circuit, modNames, instNames)
-     result getOrElse state // Return state if nothing to do
+     getMyAnnotations(state) match {
+       case Nil => CircuitState(state.circuit, state.form)
+       case myAnnotations =>
+         val (modNames, instNames) = collectAnns(myAnnotations)
+         run(state.circuit, modNames, instNames)
+     }
    }
 
    // Checks the following properties:
diff --git a/src/main/scala/firrtl/passes/memlib/InferReadWrite.scala b/src/main/scala/firrtl/passes/memlib/InferReadWrite.scala
index 16175e2c..6b56c5e8 100644
--- a/src/main/scala/firrtl/passes/memlib/InferReadWrite.scala
+++ b/src/main/scala/firrtl/passes/memlib/InferReadWrite.scala
@@ -153,12 +153,8 @@ class InferReadWrite extends Transform with PassBased {
     ResolveKinds,
     ResolveGenders
   )
-  def execute(state: CircuitState): CircuitState = {
-    val result = for {
-      myAnnotations <- getMyAnnotations(state)
-      InferReadWriteAnnotation(_) <- myAnnotations get CircuitName(state.circuit.main)
-      resCircuit = runPasses(state.circuit)
-    } yield state.copy(circuit = resCircuit)
-    result getOrElse state // Return state if nothing to do
+  def execute(state: CircuitState): CircuitState = getMyAnnotations(state) match {
+    case Nil => CircuitState(state.circuit, state.form)
+    case Seq(InferReadWriteAnnotation(_)) => CircuitState(runPasses(state.circuit), state.form)
   }
 }
diff --git a/src/main/scala/firrtl/passes/memlib/ReplaceMemMacros.scala b/src/main/scala/firrtl/passes/memlib/ReplaceMemMacros.scala
index ae872639..8d2c7200 100644
--- a/src/main/scala/firrtl/passes/memlib/ReplaceMemMacros.scala
+++ b/src/main/scala/firrtl/passes/memlib/ReplaceMemMacros.scala
@@ -219,12 +219,12 @@ class ReplaceMemMacros(writer: ConfWriter) extends Transform {
     // print conf
     writer.serialize()
     val pin = getMyAnnotations(state) match {
-      case Some(p) =>
-        p.values.head match {
-          case PinAnnotation(c, pin) => pin
-          case _ => error(s"Bad Annotations: ${p.values}")
-        }
-      case None => "pin"
+      case Nil => "pin"
+      case Seq(p) => p match {
+        case PinAnnotation(c, pin) => pin
+        case _ => error(s"Bad Annotation: ${p}")
+      }
+      case _ => throwInternalError
     }
     val annos = memMods.collect { case m: ExtModule => SinkAnnotation(ModuleName(m.name, CircuitName(c.main)), pin) }
     CircuitState(c.copy(modules = modules ++ memMods), inputForm, Some(AnnotationMap(annos)))
diff --git a/src/main/scala/firrtl/passes/memlib/ReplaceMemTransform.scala b/src/main/scala/firrtl/passes/memlib/ReplaceMemTransform.scala
index 818bd9cc..d9e8a50d 100644
--- a/src/main/scala/firrtl/passes/memlib/ReplaceMemTransform.scala
+++ b/src/main/scala/firrtl/passes/memlib/ReplaceMemTransform.scala
@@ -96,8 +96,9 @@ Optional Arguments:
 class SimpleTransform(p: Pass, form: CircuitForm) extends Transform {
   def inputForm = form
   def outputForm = form
-  def execute(state: CircuitState): CircuitState = state.copy(circuit = p.run(state.circuit))
+  def execute(state: CircuitState): CircuitState = CircuitState(p.run(state.circuit), state.form)
 }
+
 class SimpleMidTransform(p: Pass) extends SimpleTransform(p, MidForm)
 
 // SimpleRun instead of PassBased because of the arguments to passSeq
@@ -120,19 +121,21 @@ class ReplSeqMem extends Transform with SimpleRun {
         new SimpleMidTransform(ResolveKinds),
         new SimpleMidTransform(ResolveGenders))
   def run(state: CircuitState, xForms: Seq[Transform]): CircuitState = {
-    xForms.foldLeft(state) { case (curState: CircuitState, xForm: Transform) =>
-      val res = xForm.execute(state)
-      res.annotations match {
-        case None => CircuitState(res.circuit, res.form, state.annotations)
-        case Some(ann) => CircuitState(res.circuit, res.form, Some(
-          AnnotationMap(ann.annotations ++ curState.annotations.get.annotations)))
+    (xForms.foldLeft(state) { case (curState: CircuitState, xForm: Transform) =>
+      val res = xForm.execute(curState)
+      val newAnnotations = res.annotations match {
+        case None => curState.annotations
+        case Some(ann) => 
+          Some(AnnotationMap(ann.annotations ++ curState.annotations.get.annotations))
       }
-    }
+      CircuitState(res.circuit, res.form, newAnnotations)
+    }).copy(annotations = None)
   }
 
   def execute(state: CircuitState): CircuitState =
     getMyAnnotations(state) match {
-      case Some(p) => p get CircuitName(state.circuit.main) match {
+      case Nil => state.copy(annotations = None) // Do nothing if there are no annotations
+      case p => (p.collectFirst { case a if (a.target == CircuitName(state.circuit.main)) => a }) match {
         case Some(ReplSeqMemAnnotation(t)) =>
           val inputFileName = PassConfigUtil.getPassOptions(t).getOrElse(InputConfigFileName, "")
           val inConfigFile = {
@@ -144,6 +147,5 @@ class ReplSeqMem extends Transform with SimpleRun {
           run(state, passSeq(inConfigFile, outConfigFile))
         case _ => error("Unexpected transform annotation")
       }
-      case None => state // Do nothing if there are no annotations
     }
 }
diff --git a/src/main/scala/firrtl/passes/wiring/WiringTransform.scala b/src/main/scala/firrtl/passes/wiring/WiringTransform.scala
index 943c71c7..7a6fcc65 100644
--- a/src/main/scala/firrtl/passes/wiring/WiringTransform.scala
+++ b/src/main/scala/firrtl/passes/wiring/WiringTransform.scala
@@ -63,19 +63,18 @@ class WiringTransform extends Transform with SimpleRun {
         ResolveKinds,
         ResolveGenders)
   def execute(state: CircuitState): CircuitState = getMyAnnotations(state) match {
-    case Some(p) => 
+    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.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
-        }
+      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
       }
       (sources.size, tops.size, sinks.size, comp.size) match {
         case (0, 0, p, 0) => state
@@ -84,6 +83,5 @@ class WiringTransform extends Transform with SimpleRun {
           state.copy(circuit = runPasses(state.circuit, passSeq(winfo)))
         case _ => error("Wrong number of sources, tops, or sinks!")
       }
-      case None => state
   }
 }
diff --git a/src/main/scala/firrtl/transforms/Dedup.scala b/src/main/scala/firrtl/transforms/Dedup.scala
index 5d953e73..a9d3b4c9 100644
--- a/src/main/scala/firrtl/transforms/Dedup.scala
+++ b/src/main/scala/firrtl/transforms/Dedup.scala
@@ -23,7 +23,7 @@ case class DedupAnnotation(target: Named) extends Annotation with Loose with Uns
 class DedupModules extends Transform {
   def inputForm = HighForm
   def outputForm = HighForm
-  def execute(state: CircuitState): CircuitState = state.copy(circuit = run(state.circuit))
+  def execute(state: CircuitState): CircuitState = CircuitState(run(state.circuit), state.form)
   def run(c: Circuit): Circuit = {
     val moduleOrder = mutable.ArrayBuffer.empty[String]
     val moduleMap = c.modules.map(m => m.name -> m).toMap
diff --git a/src/test/scala/firrtlTests/AnnotationTests.scala b/src/test/scala/firrtlTests/AnnotationTests.scala
index 6c1adac3..a0a935a9 100644
--- a/src/test/scala/firrtlTests/AnnotationTests.scala
+++ b/src/test/scala/firrtlTests/AnnotationTests.scala
@@ -53,7 +53,7 @@ trait AnnotationSpec extends LowTransformSpec {
   }
   def execute(writer: Writer, annotations: AnnotationMap, input: String, check: Annotation) = {
     val cr = compile(CircuitState(parse(input), ChirrtlForm, Some(annotations)), writer)
-    (cr.annotations.get.annotations.head) should be (check)
+    (cr.annotations.get.annotations) should be (Seq(check))
   }
 }