Added Circuit mappers

2 files changed, 18 insertions, 1 deletions

diff --git a/src/main/scala/firrtl/Mappers.scala b/src/main/scala/firrtl/Mappers.scala
index d1e74064..cf43d9dc 100644
--- a/src/main/scala/firrtl/Mappers.scala
+++ b/src/main/scala/firrtl/Mappers.scala
@@ -98,6 +98,21 @@ object Mappers {
   }
   implicit class ModuleMap(val module: DefModule) extends AnyVal {
     def map[T](f: T => T)(implicit magnet: (T => T) => ModuleMagnet): DefModule = magnet(f).map(module)
-  }
+  } 
 
+  // ********** Circuit Mappers **********
+  private trait CircuitMagnet {
+    def map(module: Circuit): Circuit
+  }
+  private object CircuitMagnet {
+    implicit def forModules(f: DefModule => DefModule): CircuitMagnet = new CircuitMagnet {
+      override def map(circuit: Circuit): Circuit = circuit mapModule f
+    }
+    implicit def forString(f: String => String): CircuitMagnet = new CircuitMagnet {
+      override def map(circuit: Circuit): Circuit = circuit mapString f
+    }
+  }
+  implicit class CircuitMap(val circuit: Circuit) extends AnyVal {
+    def map[T](f: T => T)(implicit magnet: (T => T) => CircuitMagnet): Circuit = magnet(f).map(circuit)
+  } 
 }
diff --git a/src/main/scala/firrtl/ir/IR.scala b/src/main/scala/firrtl/ir/IR.scala
index 938b5848..aafa17ca 100644
--- a/src/main/scala/firrtl/ir/IR.scala
+++ b/src/main/scala/firrtl/ir/IR.scala
@@ -496,4 +496,6 @@ case class Circuit(info: Info, modules: Seq[DefModule], main: String) extends Fi
   def serialize: String =
     s"circuit $main :" + info.serialize +
     (modules map ("\n" + _.serialize) map indent mkString "\n") + "\n"
+  def mapModule(f: DefModule => DefModule): Circuit = this.copy(modules = modules map f)
+  def mapString(f: String => String): Circuit = this.copy(main = f(main))
 }