1 files changed, 12 insertions, 9 deletions
diff --git a/src/main/scala/firrtl/passes/memlib/InferReadWrite.scala b/src/main/scala/firrtl/passes/memlib/InferReadWrite.scala
index 28291135..2d6f4e96 100644
--- a/src/main/scala/firrtl/passes/memlib/InferReadWrite.scala
+++ b/src/main/scala/firrtl/passes/memlib/InferReadWrite.scala
@@ -38,10 +38,10 @@ import firrtl.passes.memlib.AnalysisUtils.{Connects, getConnects, getOrigin}
 import WrappedExpression.weq
 import Annotations._
 
-case class InferReadWriteAnnotation(t: String, tID: TransID)
-    extends Annotation with Loose with Unstable {
+case class InferReadWriteAnnotation(t: String) extends Annotation with Loose with Unstable {
   val target = CircuitName(t)
   def duplicate(n: Named) = this.copy(t=n.name)
+  def transform = classOf[InferReadWrite]
 }
 
 // This pass examine the enable signals of the read & write ports of memories
@@ -168,7 +168,9 @@ object InferReadWritePass extends Pass {
 
 // Transform input: Middle Firrtl. Called after "HighFirrtlToMidleFirrtl"
 // To use this transform, circuit name should be annotated with its TransId.
-class InferReadWrite(transID: TransID) extends Transform with SimpleRun {
+class InferReadWrite extends Transform with PassBased {
+  def inputForm = MidForm
+  def outputForm = MidForm
   def passSeq = Seq(
     InferReadWritePass,
     CheckInitialization,
@@ -176,11 +178,12 @@ class InferReadWrite(transID: TransID) extends Transform with SimpleRun {
     ResolveKinds,
     ResolveGenders
   )
-  def execute(c: Circuit, map: AnnotationMap) = map get transID match {
-    case Some(p) => p get CircuitName(c.main) match {
-      case Some(InferReadWriteAnnotation(_, _)) => run(c, passSeq)
-      case _ => sys.error("Unexpected annotation for InferReadWrite")
-    }
-    case _ => TransformResult(c)
+  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
   }
 }