1 files changed, 7 insertions, 1 deletions
diff --git a/src/main/scala/firrtl/analyses/InstanceGraph.scala b/src/main/scala/firrtl/analyses/InstanceGraph.scala
index 6eb67938..00689a51 100644
--- a/src/main/scala/firrtl/analyses/InstanceGraph.scala
+++ b/src/main/scala/firrtl/analyses/InstanceGraph.scala
@@ -3,12 +3,12 @@
 package firrtl.analyses
 
 import scala.collection.mutable
-
 import firrtl._
 import firrtl.ir._
 import firrtl.graph._
 import firrtl.Utils._
 import firrtl.Mappers._
+import firrtl.annotations.TargetToken.{Instance, OfModule}
 
 
 /** A class representing the instance hierarchy of a working IR Circuit
@@ -99,6 +99,12 @@ class InstanceGraph(c: Circuit) {
      */
   def getChildrenInstances: mutable.LinkedHashMap[String, mutable.LinkedHashSet[WDefInstance]] = childInstances
 
+  /** Given a circuit, returns a map from module name to children
+    * instance/module [[firrtl.annotations.TargetToken]]s
+    */
+  def getChildrenInstanceOfModule: mutable.LinkedHashMap[String, mutable.LinkedHashSet[(Instance, OfModule)]] =
+    childInstances.map(kv => kv._1 -> kv._2.map(i => (Instance(i.name), OfModule(i.module))))
+
 
 }