1 files changed, 12 insertions, 10 deletions
diff --git a/src/main/scala/firrtl/passes/memlib/DecorateMems.scala b/src/main/scala/firrtl/passes/memlib/DecorateMems.scala
index 10cc8f88..c98dd4ca 100644
--- a/src/main/scala/firrtl/passes/memlib/DecorateMems.scala
+++ b/src/main/scala/firrtl/passes/memlib/DecorateMems.scala
@@ -5,20 +5,22 @@ import ir._
 import Annotations._
 import wiring._
 
-class CreateMemoryAnnotations(reader: Option[YamlFileReader], replaceID: TransID, wiringID: TransID) extends Transform {
-  def name = "Create Memory Annotations"
-  def execute(c: Circuit, map: AnnotationMap): TransformResult = reader match {
-    case None => TransformResult(c)
+class CreateMemoryAnnotations(reader: Option[YamlFileReader]) extends Transform {
+  def inputForm = MidForm
+  def outputForm = MidForm
+  override def name = "Create Memory Annotations"
+  def execute(state: CircuitState): CircuitState = reader match {
+    case None => state
     case Some(r) =>
       import CustomYAMLProtocol._
       r.parse[Config] match {
         case Seq(config) =>
-          val cN = CircuitName(c.main)
-          val top = TopAnnotation(ModuleName(config.top.name, cN), wiringID)
-          val source = SourceAnnotation(ComponentName(config.source.name, ModuleName(config.source.module, cN)), wiringID)
-          val pin = PinAnnotation(cN, replaceID, config.pin.name)
-          TransformResult(c, None, Some(AnnotationMap(Seq(top, source, pin))))
-        case Nil => TransformResult(c, None, None)
+          val cN = CircuitName(state.circuit.main)
+          val top = TopAnnotation(ModuleName(config.top.name, cN))
+          val source = SourceAnnotation(ComponentName(config.source.name, ModuleName(config.source.module, cN)))
+          val pin = PinAnnotation(cN, config.pin.name)
+          state.copy(annotations = Some(AnnotationMap(Seq(top, source, pin))))
+        case Nil => state
         case _ => error("Can only have one config in yaml file")
       }
   }