[Module] Add afterModuleBuilt hook (#4479)

Add the `afterModuleBuilt` hook to `RawModule`, similar to the existing
`atModuleBodyEnd`. This hook is called after a module has been fully
constructed and its resulting component definition has been added to the
builder. Thus code inside the body of `afterModuleBuilt` can get the
definition of the surrounding module and instantiate it, which is
interesting for generating additional collateral like unit tests
alongside a module.

core/src/main/scala/chisel3/ModuleImpl.scala

@@ -114,6 +114,7 @@ private[chisel3] trait ObjectModuleImpl {
     Builder.setPrefix(savePrefix)
     Builder.enabledLayers = saveEnabledLayers
 
+    module.moduleBuilt()
     module
   }
 
@@ -644,6 +645,11 @@ package experimental {
       }
     }
 
+    /** Called once the module's definition has been fully built. At this point
+      * the module can be instantiated through its definition.
+      */
+    protected[chisel3] def moduleBuilt(): Unit = {}
+
     //
     // Chisel Internals
     //

core/src/main/scala/chisel3/RawModule.scala

@@ -82,6 +82,46 @@ abstract class RawModule extends BaseModule {
   }
   private val _atModuleBodyEnd = new ArrayBuffer[() => Unit]
 
+  /** Hook to invoke hardware generators after a Module has been constructed and closed.
+    *
+    * This is useful for running hardware generators after a Module's constructor has run and its Definition is available, while still having access to arguments and definitions in the constructor. The Module itself can no longer be modified at this point.
+    *
+    * An interesting application of this is the generation of unit tests whenever a module is instantiated. For example:
+    *
+    * {{{
+    *   class Example(N: int) extends RawModule {
+    *     private val someSecret: Int = ...
+    *
+    *     afterModuleBuilt {
+    *       // Executes once the surrounding module is closed.
+    *       // We can get its definition at this point and pass it to another module.
+    *       Definition(ExampleTest(this.definition, someSecret))
+    *     }
+    *   }
+    *
+    *   class ExampleTest(unitDef: Definition[Example], someSecret: Int) extends RawModule {
+    *     // Instantiate the generated module and test it.
+    *     val unit = Instance(unitDef)
+    *     ...
+    *   }
+    *
+    *   class Parent extends RawModule {
+    *     Instantiate(Example(42))
+    *   }
+    *
+    *   // Resulting modules:
+    *   // - Parent (top-level)
+    *   //   - instantiates Example
+    *   // - ExampleTest (top-level)
+    *   //   - instantiates Example
+    *   // - Example
+    * }}}
+    */
+  protected def afterModuleBuilt(gen: => Unit): Unit = {
+    _afterModuleBuilt += { () => gen }
+  }
+  private val _afterModuleBuilt = new ArrayBuffer[() => Unit]
+
   //
   // RTL construction internals
   //
@@ -253,6 +293,9 @@ abstract class RawModule extends BaseModule {
     Builder.currentBlock.get.addSecretCommand(rhs)
   }
 
+  // Evaluate any afterModuleBuilt generators.
+  protected[chisel3] override def moduleBuilt(): Unit = _afterModuleBuilt.foreach(_())
+
   private[chisel3] def initializeInParent(): Unit = {}
 }
 

src/test/scala/chiselTests/RawModuleSpec.scala

@@ -4,7 +4,7 @@ package chiselTests
 
 import chisel3._
 import chisel3.aop.Select
-import chisel3.experimental.hierarchy.Definition
+import chisel3.experimental.hierarchy.{instantiable, public, Definition, Instance}
 import chisel3.reflect.DataMirror
 import chisel3.testers.BasicTester
 import circt.stage.ChiselStage
@@ -170,4 +170,57 @@ class RawModuleSpec extends ChiselFlatSpec with Utils with MatchesAndOmits {
       }
     }
   }
+
+  "RawModule with afterModuleBuilt" should "be able to create other modules" in {
+    val chirrtl = ChiselStage.emitCHIRRTL(new RawModule {
+      override def desiredName = "Foo"
+      val port0 = IO(Input(Bool()))
+
+      afterModuleBuilt {
+        Module(new RawModule {
+          override def desiredName = "Bar"
+          val port1 = IO(Input(Bool()))
+        })
+      }
+    })
+
+    matchesAndOmits(chirrtl)(
+      "module Foo",
+      "input port0 : UInt<1>",
+      "module Bar",
+      "input port1 : UInt<1>"
+    )()
+  }
+
+  "RawModule with afterModuleBuilt" should "be able to instantiate the surrounding module" in {
+    @instantiable
+    class Foo extends RawModule {
+      override def desiredName = "Foo"
+      @public val port0 = IO(Input(Bool()))
+
+      afterModuleBuilt {
+        val fooDef = this.toDefinition
+        Module(new RawModule {
+          override def desiredName = "Bar"
+          val port1 = IO(Input(Bool()))
+          val foo1 = Instance(fooDef)
+          val foo2 = Instance(fooDef)
+          foo1.port0 := port1
+          foo2.port0 := port1
+        })
+      }
+    }
+    val chirrtl = ChiselStage.emitCHIRRTL(new Foo)
+
+    matchesAndOmits(chirrtl)(
+      "module Foo :",
+      "input port0 : UInt<1>",
+      "module Bar",
+      "input port1 : UInt<1>",
+      "inst foo1 of Foo",
+      "inst foo2 of Foo",
+      "connect foo1.port0, port1",
+      "connect foo2.port0, port1"
+    )()
+  }
 }