fixing context bug (#1874)

2 files changed, 36 insertions, 6 deletions

diff --git a/src/main/scala/chisel3/aop/injecting/InjectingAspect.scala b/src/main/scala/chisel3/aop/injecting/InjectingAspect.scala
index 170bfbad..768680ed 100644
--- a/src/main/scala/chisel3/aop/injecting/InjectingAspect.scala
+++ b/src/main/scala/chisel3/aop/injecting/InjectingAspect.scala
@@ -54,13 +54,14 @@ abstract class InjectorAspect[T <: RawModule, M <: RawModule](
     * @return
     */
   final def toAnnotation(modules: Iterable[M], circuit: String, moduleNames: Seq[String]): AnnotationSeq = {
-    val dynamicContext = new DynamicContext(annotationsInAspect)
-    // Add existing module names into the namespace. If injection logic instantiates new modules
-    //  which would share the same name, they will get uniquified accordingly
-    moduleNames.foreach { n =>
-      dynamicContext.globalNamespace.name(n)
-    }
     RunFirrtlTransformAnnotation(new InjectingTransform) +: modules.map { module =>
+      val dynamicContext = new DynamicContext(annotationsInAspect)
+      // Add existing module names into the namespace. If injection logic instantiates new modules
+      //  which would share the same name, they will get uniquified accordingly
+      moduleNames.foreach { n =>
+        dynamicContext.globalNamespace.name(n)
+      }
+
       val (chiselIR, _) = Builder.build(Module(new ModuleAspect(module) {
         module match {
           case x: Module => withClockAndReset(x.clock, x.reset) { injection(module) }
@@ -75,11 +76,15 @@ abstract class InjectorAspect[T <: RawModule, M <: RawModule](
 
       val annotations = chiselIR.annotations.map(_.toFirrtl).filterNot{ a => a.isInstanceOf[DesignAnnotation[_]] }
 
+      /** Statements to be injected via aspect. */
       val stmts = mutable.ArrayBuffer[ir.Statement]()
+      /** Modules to be injected via aspect. */
       val modules = Aspect.getFirrtl(chiselIR.copy(components = comps)).modules.flatMap {
+        // for "container" modules, inject their statements
         case m: firrtl.ir.Module if m.name == module.name =>
           stmts += m.body
           Nil
+        // for modules to be injected
         case other: firrtl.ir.DefModule =>
           Seq(other)
       }
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
+    )
+  }
 }