fixing context bug (#1874)

1 files changed, 25 insertions, 0 deletions

diff --git a/src/test/scala/chiselTests/aop/InjectionSpec.scala b/src/test/scala/chiselTests/aop/InjectionSpec.scala
index c9fa2e5e..a28501a5 100644
--- a/src/test/scala/chiselTests/aop/InjectionSpec.scala
+++ b/src/test/scala/chiselTests/aop/InjectionSpec.scala
@@ -5,6 +5,7 @@ package chiselTests.aop
 import chisel3.testers.{BasicTester, TesterDriver}
 import chiselTests.{ChiselFlatSpec, Utils}
 import chisel3._
+import chisel3.aop.Select
 import chisel3.aop.injecting.InjectingAspect
 import logger.{LogLevel, LogLevelAnnotation}
 
@@ -14,6 +15,11 @@ object InjectionHierarchy {
     val moduleSubmoduleA = Module(new SubmoduleA)
   }
 
+  class MultiModuleInjectionTester extends BasicTester {
+    val subA0 = Module(new SubmoduleA)
+    val subA1 = Module(new SubmoduleA)
+  }
+
   class SubmoduleA extends Module {
     val io = IO(new Bundle {
       val out = Output(Bool())
@@ -104,6 +110,17 @@ class InjectionSpec extends ChiselFlatSpec with Utils {
     }
   )
 
+  val multiModuleInjectionAspect = InjectingAspect(
+    { top: MultiModuleInjectionTester =>
+      Select.collectDeep(top) { case m: SubmoduleA => m }
+    },
+    { m: Module =>
+      val wire = Wire(Bool())
+      wire := m.reset.asBool()
+      dontTouch(wire)
+      stop()
+    }
+  )
 
   "Test" should "pass if inserted the correct values" in {
     assertTesterPasses{ new AspectTester(Seq(0, 1, 2)) }
@@ -142,4 +159,12 @@ class InjectionSpec extends ChiselFlatSpec with Utils {
       Seq(addingExternalModules) ++ TesterDriver.verilatorOnly
     )
   }
+
+  "Injection into multiple submodules of the same class" should "work" in {
+    assertTesterPasses(
+      {new MultiModuleInjectionTester},
+      Nil,
+      Seq(multiModuleInjectionAspect) ++ TesterDriver.verilatorOnly
+    )
+  }
 }