Add a Select.ios that works with Definition/Instance, fix isA behavior (#2315)

* New Feature: Add a Select.ios that works with Definition/Instance
* BugFix: isA now truly ignores type parameters

6 files changed, 207 insertions, 5 deletions

diff --git a/core/src/main/scala/chisel3/experimental/hierarchy/Hierarchy.scala b/core/src/main/scala/chisel3/experimental/hierarchy/Hierarchy.scala
index 503e437b..4b8d3ccc 100644
--- a/core/src/main/scala/chisel3/experimental/hierarchy/Hierarchy.scala
+++ b/core/src/main/scala/chisel3/experimental/hierarchy/Hierarchy.scala
@@ -36,7 +36,8 @@ sealed trait Hierarchy[+A] {
     */
   def isA[B : TypeTag]: Boolean = {
     val tptag = implicitly[TypeTag[B]]
-    val name = tptag.tpe.toString
+    // drop any type information for the comparison, because the proto will not have that information.
+    val name = tptag.tpe.toString.takeWhile(_ != '[') 
     inBaseClasses(name)
   }
 
diff --git a/docs/src/cookbooks/hierarchy.md b/docs/src/cookbooks/hierarchy.md
index 350d20eb..b8d2db63 100644
--- a/docs/src/cookbooks/hierarchy.md
+++ b/docs/src/cookbooks/hierarchy.md
@@ -209,13 +209,14 @@ chisel3.stage.ChiselStage.emitVerilog(new AddTwo(Definition(new AddOne(10))))
 
 Select functions can be applied after a module has been elaborated, either in a Chisel Aspect or in a parent module applied to a child module.
 
-There are six hierarchy-specific functions, which either return `Instance`'s or `Definition`'s:
+There are seven hierarchy-specific functions, which (with the exception of `ios`) either return `Instance`'s or `Definition`'s:
  - `instancesIn(parent)`: Return all instances directly instantiated locally within `parent`
  - `instancesOf[type](parent)`: Return all instances of provided `type` directly instantiated locally within `parent`
  - `allInstancesOf[type](root)`: Return all instances of provided `type` directly and indirectly instantiated, locally and deeply, starting from `root`
  - `definitionsIn`: Return definitions of all instances directly instantiated locally within `parent`
  - `definitionsOf[type]`: Return definitions of all instances of provided `type` directly instantiated locally within `parent`
  - `allDefinitionsOf[type]`: Return all definitions of instances of provided `type` directly and indirectly instantiated, locally and deeply, starting from `root`
+ - `ios`: Returns all the I/Os of the provided definition or instance.
 
 To demonstrate this, consider the following. We mock up an example where we are using the `Select.allInstancesOf` and `Select.allDefinitionsOf` to annotate instances and the definition of `EmptyModule`. When converting the `ChiselAnnotation` to firrtl's `Annotation`, we print out the resulting `Target`. As shown, despite `EmptyModule` actually only being elaborated once, we still provide different targets depending on how the instance or definition is selected.
 
@@ -259,3 +260,51 @@ println("```")
 val x = chisel3.stage.ChiselStage.emitFirrtl(new Top)
 println("```")
 ```
+
+You can also use `Select.ios` on either a `Definition` or an `Instance` to annotate the I/Os appropriately:
+
+```scala mdoc
+case class MyIOAnnotation(m: Data, tag: String) extends experimental.ChiselAnnotation {
+  def toFirrtl = {
+    println(tag + ": " + m.toTarget)
+    EmptyAnnotation
+  }
+}
+
+@instantiable
+class InOutModule extends Module {
+  @public val in = IO(Input(Bool()))
+  @public val out = IO(Output(Bool()))
+  out := in
+}
+
+@instantiable
+class TwoInOutModules extends Module {
+  val in = IO(Input(Bool()))
+  val out = IO(Output(Bool()))
+  val definition = Definition(new InOutModule)
+  val i0         = Instance(definition)
+  val i1         = Instance(definition)
+  i0.in := in
+  i1.in := i0.out
+  out := i1.out
+}
+
+class InOutTop extends Module {
+  val definition = Definition(new TwoInOutModules)
+  val instance   = Instance(definition)
+  aop.Select.allInstancesOf[InOutModule](instance).foreach { i =>
+    aop.Select.ios(i).foreach { io =>
+      experimental.annotate(MyIOAnnotation(io, "instance io"))
+  }}
+  aop.Select.allDefinitionsOf[InOutModule](instance).foreach { d =>
+    aop.Select.ios(d).foreach {io =>
+      experimental.annotate(MyIOAnnotation(io, "definition io"))
+  }}
+}
+```
+```scala mdoc:passthrough
+println("```")
+val y = chisel3.stage.ChiselStage.emitFirrtl(new InOutTop)
+println("```")
+```
\ No newline at end of file
diff --git a/src/main/scala/chisel3/aop/Select.scala b/src/main/scala/chisel3/aop/Select.scala
index 8f5a2577..6bd13445 100644
--- a/src/main/scala/chisel3/aop/Select.scala
+++ b/src/main/scala/chisel3/aop/Select.scala
@@ -65,6 +65,8 @@ object Select {
   /** 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.
+    * @note IMPORTANT: this function ignores type parameters. E.g. instancesOf[List[Int]] would return List[String].
+    * 
     * @param parent hierarchy which instantiates the returned Definitions
     */
   def instancesOf[T <: BaseModule : TypeTag](parent: Hierarchy[BaseModule]): Seq[Instance[T]] = {
@@ -90,6 +92,8 @@ object Select {
   /** 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.
+    * @note IMPORTANT: this function ignores type parameters. E.g. allInstancesOf[List[Int]] would return List[String].
+    *
     * @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]] = {
@@ -128,6 +132,8 @@ object Select {
   /** 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.
+    * @note IMPORTANT: this function ignores type parameters. E.g. definitionsOf[List[Int]] would return List[String].
+    *
     * @param parent hierarchy which instantiates the returned Definitions
     */
   def definitionsOf[T <: BaseModule : TypeTag](parent: Hierarchy[BaseModule]): Seq[Definition[T]] = {
@@ -158,6 +164,8 @@ object Select {
     *
     * @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.
+    * @note IMPORTANT: this function ignores type parameters. E.g. allDefinitionsOf[List[Int]] would return List[String].
+    *
     * @param root top of the hierarchy to search for definitions of given type
     */
   def allDefinitionsOf[T <: BaseModule : TypeTag](root: Hierarchy[BaseModule]): Seq[Definition[T]] = {
@@ -253,7 +261,7 @@ object Select {
     }
   }
 
-  /** Selects all ios directly contained within given module
+  /** Selects all ios on a given module
     * @param module
     */
   def ios(module: BaseModule): Seq[Data] = {
@@ -261,6 +269,15 @@ object Select {
     module._component.get.asInstanceOf[DefModule].ports.map(_.id)
   }
 
+  /** Selects all ios directly on a given Instance or Definition of a module
+    * @param parent the Definition or Instance to get the IOs of
+    */
+  def ios[T <: BaseModule](parent: Hierarchy[T]): Seq[Data] = {
+    check(parent)
+    implicit val mg = new chisel3.internal.MacroGenerated{}
+    parent._lookup { x => ios(parent.proto) }
+  }
+
   /** Selects all SyncReadMems directly contained within given module
     * @param module
     */
diff --git a/src/main/scala/chisel3/aop/injecting/InjectingAspect.scala b/src/main/scala/chisel3/aop/injecting/InjectingAspect.scala
index dc7e6487..ed59d4fb 100644
--- a/src/main/scala/chisel3/aop/injecting/InjectingAspect.scala
+++ b/src/main/scala/chisel3/aop/injecting/InjectingAspect.scala
@@ -4,7 +4,6 @@ package chisel3.aop.injecting
 
 import chisel3.{Module, ModuleAspect, RawModule, withClockAndReset}
 import chisel3.aop._
-import chisel3.experimental.hierarchy.IsInstantiable
 import chisel3.internal.{Builder, DynamicContext}
 import chisel3.internal.firrtl.DefModule
 import chisel3.stage.DesignAnnotation
diff --git a/src/test/scala/chiselTests/experimental/hierarchy/Examples.scala b/src/test/scala/chiselTests/experimental/hierarchy/Examples.scala
index c0f504ff..996b36ee 100644
--- a/src/test/scala/chiselTests/experimental/hierarchy/Examples.scala
+++ b/src/test/scala/chiselTests/experimental/hierarchy/Examples.scala
@@ -236,4 +236,14 @@ object Examples {
   class HasTypeParams[D <: Data](d: D) extends Module {
     @public val blah = Wire(d)
   }
+
+  @instantiable
+  class HasMultipleTypeParamsInside extends Module {
+    val tpDef0 = Definition(new HasTypeParams(Bool()))
+    val tpDef1 = Definition(new HasTypeParams(UInt(4.W)))
+    val i00 = Instance(tpDef0)
+    val i01 = Instance(tpDef0)
+    val i10 = Instance(tpDef1)
+    val i11 = Instance(tpDef1)
+  }
 }
diff --git a/src/test/scala/chiselTests/experimental/hierarchy/InstanceSpec.scala b/src/test/scala/chiselTests/experimental/hierarchy/InstanceSpec.scala
index 9ceb9b40..94af9a8b 100644
--- a/src/test/scala/chiselTests/experimental/hierarchy/InstanceSpec.scala
+++ b/src/test/scala/chiselTests/experimental/hierarchy/InstanceSpec.scala
@@ -754,6 +754,16 @@ class InstanceSpec extends ChiselFunSpec with Utils {
       }
       getFirrtlAndAnnos(new Top)
     }
+    it("9.3 it should ignore type parameters (even though it would be nice if it didn't)") {
+      class Top extends Module {
+        val d0: Definition[Module] = Definition(new HasTypeParams(Bool()))
+        require(d0.isA[HasTypeParams[Bool]])
+        require(d0.isA[HasTypeParams[_]])
+        require(d0.isA[HasTypeParams[UInt]])
+        require(!d0.isA[HasBlah])
+      }
+      getFirrtlAndAnnos(new Top)
+    }
   }
   describe("10: Select APIs") {
     it("10.0: instancesOf") {
@@ -850,6 +860,122 @@ class InstanceSpec extends ChiselFunSpec with Utils {
       })
       intercept[Exception] { getFirrtlAndAnnos(new AddFour, Seq(aspect)) }
     }
+    it("10.9: allInstancesOf.ios") {
+      val aspect = aop.inspecting.InspectingAspect({ m: AddFour =>
+        val abs = aop.Select.allInstancesOf[AddOne](m.toDefinition).flatMap { i: Instance[AddOne] => aop.Select.ios(i).map(_.toAbsoluteTarget) }
+        val rel = aop.Select.allInstancesOf[AddOne](m.toDefinition).flatMap { i: Instance[AddOne] => aop.Select.ios(i).map(_.toTarget) }
+        abs should be (Seq(
+          "~AddFour|AddFour/i0:AddTwoMixedModules/i0:AddOne>clock".rt,
+          "~AddFour|AddFour/i0:AddTwoMixedModules/i0:AddOne>reset".rt,
+          "~AddFour|AddFour/i0:AddTwoMixedModules/i0:AddOne>in".rt,
+          "~AddFour|AddFour/i0:AddTwoMixedModules/i0:AddOne>out".rt,
+
+          "~AddFour|AddFour/i0:AddTwoMixedModules/i1:AddOne_1>clock".rt,
+          "~AddFour|AddFour/i0:AddTwoMixedModules/i1:AddOne_1>reset".rt,
+          "~AddFour|AddFour/i0:AddTwoMixedModules/i1:AddOne_1>in".rt,
+          "~AddFour|AddFour/i0:AddTwoMixedModules/i1:AddOne_1>out".rt,
+
+          "~AddFour|AddFour/i1:AddTwoMixedModules/i0:AddOne>clock".rt,
+          "~AddFour|AddFour/i1:AddTwoMixedModules/i0:AddOne>reset".rt,
+          "~AddFour|AddFour/i1:AddTwoMixedModules/i0:AddOne>in".rt,
+          "~AddFour|AddFour/i1:AddTwoMixedModules/i0:AddOne>out".rt,
+
+          "~AddFour|AddFour/i1:AddTwoMixedModules/i1:AddOne_1>clock".rt,
+          "~AddFour|AddFour/i1:AddTwoMixedModules/i1:AddOne_1>reset".rt,
+          "~AddFour|AddFour/i1:AddTwoMixedModules/i1:AddOne_1>in".rt,
+          "~AddFour|AddFour/i1:AddTwoMixedModules/i1:AddOne_1>out".rt,
+        ))
+
+        rel should be (Seq(
+          "~AddFour|AddFour/i0:AddTwoMixedModules/i0:AddOne>clock".rt,
+          "~AddFour|AddFour/i0:AddTwoMixedModules/i0:AddOne>reset".rt,
+          "~AddFour|AddFour/i0:AddTwoMixedModules/i0:AddOne>in".rt,
+          "~AddFour|AddFour/i0:AddTwoMixedModules/i0:AddOne>out".rt,
+
+          "~AddFour|AddFour/i0:AddTwoMixedModules/i1:AddOne_1>clock".rt,
+          "~AddFour|AddFour/i0:AddTwoMixedModules/i1:AddOne_1>reset".rt,
+          "~AddFour|AddFour/i0:AddTwoMixedModules/i1:AddOne_1>in".rt,
+          "~AddFour|AddFour/i0:AddTwoMixedModules/i1:AddOne_1>out".rt,
+
+          "~AddFour|AddFour/i1:AddTwoMixedModules/i0:AddOne>clock".rt,
+          "~AddFour|AddFour/i1:AddTwoMixedModules/i0:AddOne>reset".rt,
+          "~AddFour|AddFour/i1:AddTwoMixedModules/i0:AddOne>in".rt,
+          "~AddFour|AddFour/i1:AddTwoMixedModules/i0:AddOne>out".rt,
+             
+          "~AddFour|AddFour/i1:AddTwoMixedModules/i1:AddOne_1>clock".rt,
+          "~AddFour|AddFour/i1:AddTwoMixedModules/i1:AddOne_1>reset".rt,
+          "~AddFour|AddFour/i1:AddTwoMixedModules/i1:AddOne_1>in".rt,
+          "~AddFour|AddFour/i1:AddTwoMixedModules/i1:AddOne_1>out".rt,
+        ))
+      })
+      getFirrtlAndAnnos(new AddFour, Seq(aspect))
+    }
+    it("10.10: allDefinitionsOf.ios") {
+      val aspect = aop.inspecting.InspectingAspect({ m: AddFour =>
+        val abs = aop.Select.allDefinitionsOf[AddOne](m.toDefinition).flatMap { i: Definition[AddOne] => aop.Select.ios(i).map(_.toAbsoluteTarget) }
+        val rel = aop.Select.allDefinitionsOf[AddOne](m.toDefinition).flatMap { i: Definition[AddOne] => aop.Select.ios(i).map(_.toTarget) }
+        abs should be (Seq(
+          "~AddFour|AddOne>clock".rt,
+          "~AddFour|AddOne>reset".rt,
+          "~AddFour|AddOne>in".rt,
+          "~AddFour|AddOne>out".rt,
+
+          "~AddFour|AddOne_1>clock".rt,
+          "~AddFour|AddOne_1>reset".rt,
+          "~AddFour|AddOne_1>in".rt,
+          "~AddFour|AddOne_1>out".rt,
+        ))
+
+        rel should be (Seq(
+          "~AddFour|AddOne>clock".rt,
+          "~AddFour|AddOne>reset".rt,
+          "~AddFour|AddOne>in".rt,
+          "~AddFour|AddOne>out".rt,
+
+          "~AddFour|AddOne_1>clock".rt,
+          "~AddFour|AddOne_1>reset".rt,
+          "~AddFour|AddOne_1>in".rt,
+          "~AddFour|AddOne_1>out".rt,
+        ))
+        
+      })
+      getFirrtlAndAnnos(new AddFour, Seq(aspect))
+    }
+    it("10.11 Select.instancesIn for typed BaseModules") {
+      val aspect = aop.inspecting.InspectingAspect({ m: HasMultipleTypeParamsInside =>
+        val targets = aop.Select.instancesIn(m.toDefinition).map { i: Instance[BaseModule] => i.toTarget }
+        targets should be (Seq(
+          "~HasMultipleTypeParamsInside|HasMultipleTypeParamsInside/i00:HasTypeParams".it,
+          "~HasMultipleTypeParamsInside|HasMultipleTypeParamsInside/i01:HasTypeParams".it,
+          "~HasMultipleTypeParamsInside|HasMultipleTypeParamsInside/i10:HasTypeParams_1".it,
+          "~HasMultipleTypeParamsInside|HasMultipleTypeParamsInside/i11:HasTypeParams_1".it,
+        ))     
+      })
+      getFirrtlAndAnnos(new HasMultipleTypeParamsInside, Seq(aspect))
+    }
+    it("10.12 Select.instancesOf for typed BaseModules if type is ignored") {
+      val aspect = aop.inspecting.InspectingAspect({ m: HasMultipleTypeParamsInside =>
+        val targets = aop.Select.instancesOf[HasTypeParams[_]](m.toDefinition).map { i: Instance[HasTypeParams[_]] => i.toTarget }
+        targets should be (Seq(
+          "~HasMultipleTypeParamsInside|HasMultipleTypeParamsInside/i00:HasTypeParams".it,
+          "~HasMultipleTypeParamsInside|HasMultipleTypeParamsInside/i01:HasTypeParams".it,
+          "~HasMultipleTypeParamsInside|HasMultipleTypeParamsInside/i10:HasTypeParams_1".it,
+          "~HasMultipleTypeParamsInside|HasMultipleTypeParamsInside/i11:HasTypeParams_1".it,
+        ))       
+      })
+      getFirrtlAndAnnos(new HasMultipleTypeParamsInside, Seq(aspect))
+    }
+    it("10.13 Select.instancesOf for typed BaseModules even type is specified wrongly (should be ignored, even though we wish it weren't)") {
+      val aspect = aop.inspecting.InspectingAspect({ m: HasMultipleTypeParamsInside =>
+        val targets = aop.Select.instancesOf[HasTypeParams[SInt]](m.toDefinition).map { i: Instance[HasTypeParams[_]] => i.toTarget }
+        targets should be (Seq(
+          "~HasMultipleTypeParamsInside|HasMultipleTypeParamsInside/i00:HasTypeParams".it,
+          "~HasMultipleTypeParamsInside|HasMultipleTypeParamsInside/i01:HasTypeParams".it,
+          "~HasMultipleTypeParamsInside|HasMultipleTypeParamsInside/i10:HasTypeParams_1".it,
+          "~HasMultipleTypeParamsInside|HasMultipleTypeParamsInside/i11:HasTypeParams_1".it,
+        ))       
+      })
+      getFirrtlAndAnnos(new HasMultipleTypeParamsInside, Seq(aspect)) 
+    }
   }
 }
-