Add Select APIs for Hierarchy package (#2210)

* Add Hierarchy trait

* Add Hierarchy trait

* Add Hierarchy scaladoc

* Add license

* Add isA and tests

* Add back isA

* Add new Select APIs for hierarchy package

* Update scaladoc

* Write outlines for tests

* Add tests and fixes to new Select functions

* Make calculate via lazy val

* Apply suggestions from code review

Co-authored-by: Megan Wachs <megan@sifive.com>

* Apply suggestions from code review

Co-authored-by: Megan Wachs <megan@sifive.com>

* Clean up scaladoc

* Add shouldNot compile

* Apply suggestions from code review

Co-authored-by: Megan Wachs <megan@sifive.com>

* Bugfix all funcs should analyze root too

* Add mdoc, bugfix toDefinition

* Make func private, add scaladoc

* Update src/test/scala/chiselTests/experimental/hierarchy/InstanceSpec.scala

Co-authored-by: Jack Koenig <koenig@sifive.com>

* Made protected vals private

* Apply suggestions from code review

Co-authored-by: Jack Koenig <koenig@sifive.com>

* Address code review comments

* Added additional null check

Co-authored-by: Megan Wachs <megan@sifive.com>
Co-authored-by: Jack Koenig <koenig@sifive.com>

4 files changed, 264 insertions, 24 deletions

diff --git a/src/main/scala/chisel3/aop/Select.scala b/src/main/scala/chisel3/aop/Select.scala
index 9c7320ce..8bdf4344 100644
--- a/src/main/scala/chisel3/aop/Select.scala
+++ b/src/main/scala/chisel3/aop/Select.scala
@@ -6,10 +6,12 @@ import chisel3._
 import chisel3.internal.{HasId}
 import chisel3.experimental.BaseModule
 import chisel3.experimental.FixedPoint
-import chisel3.internal.firrtl._
+import chisel3.internal.firrtl.{Definition => DefinitionIR, _}
+import chisel3.experimental.hierarchy._
 import chisel3.internal.PseudoModule
 import chisel3.internal.BaseModule.ModuleClone
 import firrtl.annotations.ReferenceTarget
+import scala.reflect.runtime.universe.TypeTag
 
 import scala.collection.mutable
 import chisel3.internal.naming.chiselName
@@ -22,7 +24,6 @@ object Select {
   /** Return just leaf components of expanded node
     *
     * @param d Component to find leafs if aggregate typed. Intermediate fields/indicies are not included
-    * @return
     */
   def getLeafs(d: Data): Seq[Data] = d match {
     case r: Record => r.getElements.flatMap(getLeafs)
@@ -33,7 +34,6 @@ object Select {
   /** Return all expanded components, including intermediate aggregate nodes
     *
     * @param d Component to find leafs if aggregate typed. Intermediate fields/indicies ARE included
-    * @return
     */
   def getIntermediateAndLeafs(d: Data): Seq[Data] = d match {
     case r: Record => r +: r.getElements.flatMap(getIntermediateAndLeafs)
@@ -41,15 +41,156 @@ object Select {
     case other => Seq(other)
   }
 
+  /** Selects all instances/modules directly instantiated within given definition
+    *
+    * @param parent
+    */
+  def instancesIn(parent: Hierarchy[BaseModule]): Seq[Instance[BaseModule]] = {
+    check(parent)
+    implicit val mg = new chisel3.internal.MacroGenerated{}
+    parent.proto._component.get match {
+      case d: DefModule => d.commands.collect {
+        case d: DefInstance =>
+          d.id match {
+            case p: chisel3.internal.BaseModule.IsClone[_] =>
+              parent._lookup { x => new Instance(Right(p)).asInstanceOf[Instance[BaseModule]] }
+            case other: BaseModule =>
+              parent._lookup { x => other }
+          }
+      }
+      case other => Nil
+    }
+  }
+
+  /** Selects all Instances of instances/modules directly instantiated within given module, of provided type
+    *
+    * @note IMPORTANT: this function requires summoning a TypeTag[T], which will fail if T is an inner class.
+    * @param parent hierarchy which instantiates the returned Definitions
+    */
+  def instancesOf[T <: BaseModule : TypeTag](parent: Hierarchy[BaseModule]): Seq[Instance[T]] = {
+    check(parent)
+    implicit val mg = new chisel3.internal.MacroGenerated{}
+    parent.proto._component.get match {
+      case d: DefModule => d.commands.flatMap {
+        case d: DefInstance =>
+          d.id match {
+            case p: chisel3.internal.BaseModule.IsClone[_] =>
+              val i = parent._lookup { x => new Instance(Right(p)).asInstanceOf[Instance[BaseModule]] }
+              if(i.isA[T]) Some(i.asInstanceOf[Instance[T]]) else None
+            case other: BaseModule =>
+              val i = parent._lookup { x => other }
+              if(i.isA[T]) Some(i.asInstanceOf[Instance[T]]) else None
+          }
+        case other => None
+      }
+      case other => Nil
+    }
+  }
+
+  /** Selects all Instances directly and indirectly instantiated within given root hierarchy, of provided type
+    *
+    * @note IMPORTANT: this function requires summoning a TypeTag[T], which will fail if T is an inner class.
+    * @param root top of the hierarchy to search for instances/modules of given type
+    */
+  def allInstancesOf[T <: BaseModule : TypeTag](root: Hierarchy[BaseModule]): Seq[Instance[T]] = {
+    val soFar = if(root.isA[T]) Seq(root.toInstance.asInstanceOf[Instance[T]]) else Nil
+    val allLocalInstances = instancesIn(root)
+    soFar ++ (allLocalInstances.flatMap(allInstancesOf[T]))
+  }
+
+  /** Selects the Definitions of all instances/modules directly instantiated within given module
+    * 
+    * @param parent
+    */
+  def definitionsIn(parent: Hierarchy[BaseModule]): Seq[Definition[BaseModule]] = {
+    type DefType = Definition[BaseModule]
+    implicit val mg = new chisel3.internal.MacroGenerated{}
+    check(parent)
+    val defs = parent.proto._component.get match {
+      case d: DefModule => d.commands.collect {
+        case i: DefInstance =>
+          i.id match {
+            case p: chisel3.internal.BaseModule.IsClone[_] =>
+              parent._lookup { x => new Definition(Left(p.getProto)).asInstanceOf[Definition[BaseModule]] }
+            case other: BaseModule =>
+              parent._lookup { x => other.toDefinition }
+          }
+      }
+      case other => Nil
+    }
+    val (_, defList) = defs.foldLeft((Set.empty[DefType], List.empty[DefType])) { case ((set, list), definition: Definition[BaseModule]) =>
+      if(set.contains(definition)) (set, list) else (set + definition, definition +: list)
+    }
+    defList.reverse
+  }
+
+
+  /** Selects all Definitions of instances/modules directly instantiated within given module, of provided type
+    *
+    * @note IMPORTANT: this function requires summoning a TypeTag[T], which will fail if T is an inner class.
+    * @param parent hierarchy which instantiates the returned Definitions
+    */
+  def definitionsOf[T <: BaseModule : TypeTag](parent: Hierarchy[BaseModule]): Seq[Definition[T]] = {
+    check(parent)
+    implicit val mg = new chisel3.internal.MacroGenerated{}
+    type DefType = Definition[T]
+    val defs = parent.proto._component.get match {
+      case d: DefModule => d.commands.flatMap {
+        case d: DefInstance =>
+          d.id match {
+            case p: chisel3.internal.BaseModule.IsClone[_] =>
+              val d = parent._lookup { x => new Definition(Right(p)).asInstanceOf[Definition[BaseModule]] }
+              if(d.isA[T]) Some(d.asInstanceOf[Definition[T]]) else None
+            case other: BaseModule =>
+              val d = parent._lookup { x => other.toDefinition }
+              if(d.isA[T]) Some(d.asInstanceOf[Definition[T]]) else None
+          }
+        case other => None
+      }
+    }
+    val (_, defList) = defs.foldLeft((Set.empty[DefType], List.empty[DefType])) { case ((set, list), definition: Definition[T]) =>
+      if(set.contains(definition)) (set, list) else (set + definition, definition +: list)
+    }
+    defList.reverse
+  }
+
+  /** Selects all Definition's directly and indirectly instantiated within given root hierarchy, of provided type
+    *
+    * @note IMPORTANT: this function requires summoning a TypeTag[T], which will fail if T is an inner class, i.e.
+    *   a class defined within another class.
+    * @param root top of the hierarchy to search for definitions of given type
+    */
+  def allDefinitionsOf[T <: BaseModule : TypeTag](root: Hierarchy[BaseModule]): Seq[Definition[T]] = {
+    type DefType = Definition[T]
+    val allDefSet = mutable.HashSet[Definition[BaseModule]]()
+    val defSet = mutable.HashSet[DefType]()
+    val defList = mutable.ArrayBuffer[DefType]()
+    def rec(hier: Definition[BaseModule]): Unit = {
+      if(hier.isA[T] && !defSet.contains(hier.asInstanceOf[DefType])) {
+        defSet  += hier.asInstanceOf[DefType]
+        defList += hier.asInstanceOf[DefType]
+      }
+      allDefSet += hier
+      val allDefs = definitionsIn(hier)
+      allDefs.collect {
+        case d if !allDefSet.contains(d) => rec(d)
+      }
+    }
+    rec(root.toDefinition)
+    defList.toList
+  }
+
 
   /** Collects all components selected by collector within module and all children modules it instantiates
     *   directly or indirectly
     * Accepts a collector function, rather than a collector partial function (see [[collectDeep]])
+    *
+    * @note This API will not work with the new experimental hierarchy package. Instead, use allInstancesOf or allDefinitionsOf.
+    *
     * @param module Module to collect components, as well as all children module it directly and indirectly instantiates
     * @param collector Collector function to pick, given a module, which components to collect
     * @param tag Required for generics to work, should ignore this
     * @tparam T Type of the component that will be collected
-    * @return
     */
   def getDeep[T](module: BaseModule)(collector: BaseModule => Seq[T]): Seq[T] = {
     check(module)
@@ -63,11 +204,13 @@ object Select {
   /** Collects all components selected by collector within module and all children modules it instantiates
     *   directly or indirectly
     * Accepts a collector partial function, rather than a collector function (see [[getDeep]])
+    *
+    * @note This API will not work with the new experimental hierarchy package. Instead, use allInstancesOf or allDefinitionsOf.
+    *
     * @param module Module to collect components, as well as all children module it directly and indirectly instantiates
     * @param collector Collector partial function to pick, given a module, which components to collect
     * @param tag Required for generics to work, should ignore this
     * @tparam T Type of the component that will be collected
-    * @return
     */
   def collectDeep[T](module: BaseModule)(collector: PartialFunction[BaseModule, T]): Iterable[T] = {
     check(module)
@@ -78,9 +221,11 @@ object Select {
     myItems ++ deepChildrenItems
   }
 
-  /** Selects all instances directly instantiated within given module
+  /** Selects all modules directly instantiated within given module
+    *
+    * @note This API will not work with the new experimental hierarchy package. Instead, use instancesIn or definitionsIn.
+    *
     * @param module
-    * @return
     */
   def instances(module: BaseModule): Seq[BaseModule] = {
     check(module)
@@ -88,7 +233,7 @@ object Select {
       case d: DefModule => d.commands.flatMap {
         case i: DefInstance => i.id match {
           case m: ModuleClone[_] if !m._madeFromDefinition => None
-          case _: PseudoModule => throw new Exception("Aspect APIs are currently incompatible with Definition/Instance")
+          case _: PseudoModule => throw new Exception("instances, collectDeep, and getDeep are currently incompatible with Definition/Instance!")
           case other          => Some(other)
         }
         case _ => None
@@ -99,7 +244,6 @@ object Select {
 
   /** Selects all registers directly instantiated within given module
     * @param module
-    * @return
     */
   def registers(module: BaseModule): Seq[Data] = {
     check(module)
@@ -111,7 +255,6 @@ object Select {
 
   /** Selects all ios directly contained within given module
     * @param module
-    * @return
     */
   def ios(module: BaseModule): Seq[Data] = {
     check(module)
@@ -120,7 +263,6 @@ object Select {
 
   /** Selects all SyncReadMems directly contained within given module
     * @param module
-    * @return
     */
   def syncReadMems(module: BaseModule): Seq[SyncReadMem[_]] = {
     check(module)
@@ -131,7 +273,6 @@ object Select {
 
   /** Selects all Mems directly contained within given module
     * @param module
-    * @return
     */
   def mems(module: BaseModule): Seq[Mem[_]] = {
     check(module)
@@ -142,7 +283,6 @@ object Select {
 
   /** Selects all arithmetic or logical operators directly instantiated within given module
     * @param module
-    * @return
     */
   def ops(module: BaseModule): Seq[(String, Data)] = {
     check(module)
@@ -155,7 +295,6 @@ object Select {
     * The kind of operators are contained in [[chisel3.internal.firrtl.PrimOp]]
     * @param opKind the kind of operator, e.g. "mux", "add", or "bits"
     * @param module
-    * @return
     */
   def ops(opKind: String)(module: BaseModule): Seq[Data] = {
     check(module)
@@ -166,7 +305,6 @@ object Select {
 
   /** Selects all wires in a module
     * @param module
-    * @return
     */
   def wires(module: BaseModule): Seq[Data] = {
     check(module)
@@ -177,7 +315,6 @@ object Select {
 
   /** Selects all memory ports, including their direction and memory
     * @param module
-    * @return
     */
   def memPorts(module: BaseModule): Seq[(Data, MemPortDirection, MemBase[_])] = {
     check(module)
@@ -189,7 +326,6 @@ object Select {
   /** Selects all memory ports of a given direction, including their memory
     * @param dir The direction of memory ports to select
     * @param module
-    * @return
     */
   def memPorts(dir: MemPortDirection)(module: BaseModule): Seq[(Data, MemBase[_])] = {
     check(module)
@@ -200,7 +336,6 @@ object Select {
 
   /** Selects all components who have been set to be invalid, even if they are later connected to
     * @param module
-    * @return
     */
   def invalids(module: BaseModule): Seq[Data] = {
     check(module)
@@ -211,7 +346,6 @@ object Select {
 
   /** Selects all components who are attached to a given signal, within a module
     * @param module
-    * @return
     */
   def attachedTo(module: BaseModule)(signal: Data): Set[Data] = {
     check(module)
@@ -226,7 +360,6 @@ object Select {
     * E.g. if signal = io.foo.bar, connectionsTo will return all connections to io, io.foo, and io.bar
     * @param module
     * @param signal
-    * @return
     */
   def connectionsTo(module: BaseModule)(signal: Data): Seq[PredicatedConnect] = {
     check(module)
@@ -237,7 +370,7 @@ object Select {
     var seenDef = isPort
     searchWhens(module, (cmd: Command, preds) => {
       cmd match {
-        case cmd: Definition if cmd.id.isInstanceOf[Data] =>
+        case cmd: DefinitionIR if cmd.id.isInstanceOf[Data] =>
           val x = getIntermediateAndLeafs(cmd.id.asInstanceOf[Data])
           if(x.contains(signal)) prePredicates = preds
         case Connect(_, loc@Node(d: Data), exp) =>
@@ -263,7 +396,6 @@ object Select {
   /** Selects all stop statements, and includes the predicates surrounding the stop statement
     *
     * @param module
-    * @return
     */
   def stops(module: BaseModule): Seq[Stop]  = {
     val stops = mutable.ArrayBuffer[Stop]()
@@ -279,7 +411,6 @@ object Select {
   /** Selects all printf statements, and includes the predicates surrounding the printf statement
     *
     * @param module
-    * @return
     */
   def printfs(module: BaseModule): Seq[Printf] = {
     val printfs = mutable.ArrayBuffer[Printf]()
@@ -297,6 +428,7 @@ object Select {
     require(module.isClosed, "Can't use Selector on modules that have not finished construction!")
     require(module._component.isDefined, "Can't use Selector on modules that don't have components!")
   }
+  private def check(hierarchy: Hierarchy[BaseModule]): Unit = check(hierarchy.proto)
 
   // Given a loc, return all subcomponents of id that could be assigned to in connect
   private def getEffected(a: Arg): Seq[Data] = a match {
diff --git a/src/main/scala/chisel3/aop/injecting/InjectingAspect.scala b/src/main/scala/chisel3/aop/injecting/InjectingAspect.scala
index 1a476f61..dc7e6487 100644
--- a/src/main/scala/chisel3/aop/injecting/InjectingAspect.scala
+++ b/src/main/scala/chisel3/aop/injecting/InjectingAspect.scala
@@ -43,7 +43,7 @@ abstract class InjectorAspect[T <: RawModule, M <: RawModule](
     injection: M => Unit
 ) extends Aspect[T] {
   final def toAnnotation(top: T): AnnotationSeq = {
-    val moduleNames = Select.collectDeep(top) { case i => i.name }.toSeq
+    val moduleNames = Select.allDefinitionsOf[chisel3.experimental.BaseModule](top.toDefinition).map{i => i.toTarget.module }.toSeq
     toAnnotation(selectRoots(top), top.name, moduleNames)
   }
 
diff --git a/src/test/scala/chiselTests/experimental/hierarchy/Examples.scala b/src/test/scala/chiselTests/experimental/hierarchy/Examples.scala
index 7c9396cf..d8ae7322 100644
--- a/src/test/scala/chiselTests/experimental/hierarchy/Examples.scala
+++ b/src/test/scala/chiselTests/experimental/hierarchy/Examples.scala
@@ -57,6 +57,18 @@ object Examples {
     i1.in := i0.out
     out := i1.out
   }
+
+  @instantiable
+  class AddFour extends Module {
+    @public val in  = IO(Input(UInt(32.W)))
+    @public val out = IO(Output(UInt(32.W)))
+    @public val definition = Definition(new AddTwoMixedModules)
+    @public val i0 = Instance(definition)
+    @public val i1 = Instance(definition)
+    i0.in := in
+    i1.in := i0.out
+    out := i1.out
+  }
   @instantiable
   class AggregatePortModule extends Module {
     @public val io = IO(new Bundle {
diff --git a/src/test/scala/chiselTests/experimental/hierarchy/InstanceSpec.scala b/src/test/scala/chiselTests/experimental/hierarchy/InstanceSpec.scala
index 0795e76c..83084468 100644
--- a/src/test/scala/chiselTests/experimental/hierarchy/InstanceSpec.scala
+++ b/src/test/scala/chiselTests/experimental/hierarchy/InstanceSpec.scala
@@ -755,5 +755,101 @@ class InstanceSpec extends ChiselFunSpec with Utils {
       getFirrtlAndAnnos(new Top)
     }
   }
+  describe("10: Select APIs") {
+    it("10.0: instancesOf") {
+      val aspect = aop.inspecting.InspectingAspect({ m: AddTwoMixedModules =>
+        val targets = aop.Select.instancesOf[AddOne](m.toDefinition).map { i: Instance[AddOne] => i.toTarget }
+        targets should be (Seq(
+          "~AddTwoMixedModules|AddTwoMixedModules/i0:AddOne".it,
+          "~AddTwoMixedModules|AddTwoMixedModules/i1:AddOne_2".it,
+        ))
+      })
+      getFirrtlAndAnnos(new AddTwoMixedModules, Seq(aspect))
+    }
+    it("10.1: instancesIn") {
+      val aspect = aop.inspecting.InspectingAspect({ m: AddTwoMixedModules =>
+        val insts = aop.Select.instancesIn(m.toDefinition)
+        val abs = insts.map { i: Instance[BaseModule] => i.toAbsoluteTarget }
+        val rel = insts.map { i: Instance[BaseModule] => i.toTarget }
+        abs should be (Seq(
+          "~AddTwoMixedModules|AddTwoMixedModules/i0:AddOne".it,
+          "~AddTwoMixedModules|AddTwoMixedModules/i1:AddOne_2".it,
+        ))
+        rel should be (Seq(
+          "~AddTwoMixedModules|AddTwoMixedModules/i0:AddOne".it,
+          "~AddTwoMixedModules|AddTwoMixedModules/i1:AddOne_2".it,
+        ))
+      })
+      getFirrtlAndAnnos(new AddTwoMixedModules, Seq(aspect))
+    }
+    it("10.2: allInstancesOf") {
+      val aspect = aop.inspecting.InspectingAspect({ m: AddFour =>
+        val insts = aop.Select.allInstancesOf[AddOne](m.toDefinition)
+        val abs = insts.map { i: Instance[AddOne] => i.in.toAbsoluteTarget }
+        val rel = insts.map { i: Instance[AddOne] => i.in.toTarget }
+        rel should be (Seq(
+          "~AddFour|AddFour/i0:AddTwoMixedModules/i0:AddOne>in".rt,
+          "~AddFour|AddFour/i0:AddTwoMixedModules/i1:AddOne_2>in".rt,
+          "~AddFour|AddFour/i1:AddTwoMixedModules/i0:AddOne>in".rt,
+          "~AddFour|AddFour/i1:AddTwoMixedModules/i1:AddOne_2>in".rt,
+        ))
+        abs should be (Seq(
+          "~AddFour|AddFour/i0:AddTwoMixedModules/i0:AddOne>in".rt,
+          "~AddFour|AddFour/i0:AddTwoMixedModules/i1:AddOne_2>in".rt,
+          "~AddFour|AddFour/i1:AddTwoMixedModules/i0:AddOne>in".rt,
+          "~AddFour|AddFour/i1:AddTwoMixedModules/i1:AddOne_2>in".rt,
+        ))
+      })
+      getFirrtlAndAnnos(new AddFour, Seq(aspect))
+    }
+    it("10.3: definitionsOf") {
+      val aspect = aop.inspecting.InspectingAspect({ m: AddTwoMixedModules =>
+        val targets = aop.Select.definitionsOf[AddOne](m.toDefinition).map { i: Definition[AddOne] => i.in.toTarget }
+        targets should be (Seq(
+          "~AddTwoMixedModules|AddOne>in".rt,
+          "~AddTwoMixedModules|AddOne_2>in".rt,
+        ))
+      })
+      getFirrtlAndAnnos(new AddTwoMixedModules, Seq(aspect))
+    }
+    it("10.4: definitionsIn") {
+      val aspect = aop.inspecting.InspectingAspect({ m: AddTwoMixedModules =>
+        val targets = aop.Select.definitionsIn(m.toDefinition).map { i: Definition[BaseModule] => i.toTarget }
+        targets should be (Seq(
+          "~AddTwoMixedModules|AddOne".mt,
+          "~AddTwoMixedModules|AddOne_2".mt,
+        ))
+      })
+      getFirrtlAndAnnos(new AddTwoMixedModules, Seq(aspect))
+    }
+    it("10.5: allDefinitionsOf") {
+      val aspect = aop.inspecting.InspectingAspect({ m: AddFour =>
+        val targets = aop.Select.allDefinitionsOf[AddOne](m.toDefinition).map { i: Definition[AddOne] => i.in.toTarget }
+        targets should be (Seq(
+          "~AddFour|AddOne>in".rt,
+          "~AddFour|AddOne_2>in".rt,
+        ))
+      })
+      getFirrtlAndAnnos(new AddFour, Seq(aspect))
+    }
+    it("10.6: Select.collectDeep should fail when combined with hierarchy package") {
+      val aspect = aop.inspecting.InspectingAspect({ m: AddFour =>
+        aop.Select.collectDeep(m) { case m: AddOne => m.toTarget }
+      })
+      intercept[Exception] { getFirrtlAndAnnos(new AddFour, Seq(aspect)) }
+    }
+    it("10.7: Select.getDeep should fail when combined with hierarchy package") {
+      val aspect = aop.inspecting.InspectingAspect({ m: AddFour =>
+        aop.Select.getDeep(m) { m: BaseModule => Nil }
+      })
+      intercept[Exception] { getFirrtlAndAnnos(new AddFour, Seq(aspect)) }
+    }
+    it("10.8: Select.instances should fail when combined with hierarchy package") {
+      val aspect = aop.inspecting.InspectingAspect({ m: AddFour =>
+        aop.Select.instances(m)
+      })
+      intercept[Exception] { getFirrtlAndAnnos(new AddFour, Seq(aspect)) }
+    }
+  }
 }