tilelink2: use LazyModule(new ...) just like Chisel Module(new ...)

uncore/src/main/scala/tilelink2/Buffer.scala

@@ -8,7 +8,7 @@ class TLBuffer(entries: Int = 2, pipe: Boolean = false) extends LazyModule
 {
   val node = TLIdentityNode()
 
-  lazy val module = Module(new LazyModuleImp(this) {
+  lazy val module = new LazyModuleImp(this) {
     val io = new Bundle {
       val in  = node.bundleIn
       val out = node.bundleOut
@@ -26,15 +26,14 @@ class TLBuffer(entries: Int = 2, pipe: Boolean = false) extends LazyModule
       out.c <> Queue(in .c, entries, pipe)
       out.e <> Queue(out.e, entries, pipe)
     }
-  })
+  }
 }
 
 object TLBuffer
 {
   // applied to the TL source node; connect (TLBuffer(x.node) -> y.node)
   def apply(x: TLBaseNode, entries: Int = 2, pipe: Boolean = false)(implicit lazyModule: LazyModule): TLBaseNode = {
-    val buffer = new TLBuffer(entries, pipe)
-    lazyModule.addChild(buffer)
+    val buffer = LazyModule(new TLBuffer(entries, pipe))
     lazyModule.connect(x -> buffer.node)
     buffer.node
   }

uncore/src/main/scala/tilelink2/HintHandler.scala

@@ -11,7 +11,7 @@ class TLHintHandler(supportManagers: Boolean = true, supportClients: Boolean = f
     clientFn  = { case Seq(c) => if (supportClients)  c.copy(clients  = c.clients .map(_.copy(supportsHint = true))) else c },
     managerFn = { case Seq(m) => if (supportManagers) m.copy(managers = m.managers.map(_.copy(supportsHint = true))) else m })
 
-  lazy val module = Module(new LazyModuleImp(this) {
+  lazy val module = new LazyModuleImp(this) {
     val io = new Bundle {
       val in  = node.bundleIn
       val out = node.bundleOut
@@ -76,15 +76,14 @@ class TLHintHandler(supportManagers: Boolean = true, supportClients: Boolean = f
       in.e.ready := out.e.ready
       out.e.bits := in.e.bits
     }
-  })
+  }
 }
 
 object TLHintHandler
 {
   // applied to the TL source node; connect (TLHintHandler(x.node) -> y.node)
   def apply(x: TLBaseNode, supportManagers: Boolean = true, supportClients: Boolean = false, passthrough: Boolean = true)(implicit lazyModule: LazyModule): TLBaseNode = {
-    val hints = new TLHintHandler(supportManagers, supportClients, passthrough)
-    lazyModule.addChild(hints)
+    val hints = LazyModule(new TLHintHandler(supportManagers, supportClients, passthrough))
     lazyModule.connect(x -> hints.node)
     hints.node
   }

uncore/src/main/scala/tilelink2/LazyModule.scala

@@ -3,45 +3,58 @@
 package uncore.tilelink2
 
 import Chisel._
-import scala.collection.mutable.ListBuffer
-import chisel3.internal.sourceinfo.SourceInfo
+import chisel3.internal.sourceinfo._
 
 abstract class LazyModule
 {
-  private val bindings = ListBuffer[() => Unit]()
-  private val extraChildren = ListBuffer[LazyModule]()
+  protected[tilelink2] var bindings = List[() => Unit]()
+  protected[tilelink2] var children = List[LazyModule]()
+  protected[tilelink2] var info: SourceInfo = UnlocatableSourceInfo
+  protected[tilelink2] val parent = LazyModule.stack.headOption
+
+  LazyModule.stack = this :: LazyModule.stack
+  parent.foreach(p => p.children = this :: p.children)
 
   // Use as: connect(source -> sink, source2 -> sink2, ...)
   def connect[PO, PI, EO, EI, B <: Bundle](edges: (BaseNode[PO, PI, EO, EI, B], BaseNode[PO, PI, EO, EI, B])*)(implicit sourceInfo: SourceInfo) = {
     edges.foreach { case (source, sink) =>
-      bindings += (source edge sink)
+      bindings = (source edge sink) :: bindings
     }
   }
 
   def module: LazyModuleImp
+  implicit val lazyModule = this
 
   protected[tilelink2] def instantiate() = {
-    // Find all LazyModule members of self
-    for (m <- getClass.getMethods) {
-      if (m.getParameterTypes.isEmpty && 
-          !java.lang.reflect.Modifier.isStatic(m.getModifiers) &&
-          !(m.getName contains '$') &&
-          !(m.getName == "lazyModule") &&
-          classOf[LazyModule].isAssignableFrom(m.getReturnType)) {
-        // ... and force their lazy module members to exist
-        m.invoke(this).asInstanceOf[LazyModule].module
+    children.reverse.foreach { c => 
+      // !!! fix chisel3 so we can pass the desired sourceInfo
+      // implicit val sourceInfo = c.module.outer.info
+      Module(c.module)
     }
+    bindings.reverse.foreach { f => f () }
   }
-    extraChildren.foreach { _.module }
-    bindings.foreach { f => f () }
 }
 
-  implicit val lazyModule = this
-  def addChild(x: LazyModule) = extraChildren += x
-}
-
-abstract class LazyModuleImp(outer: LazyModule) extends Module
+object LazyModule
 {
+  protected[tilelink2] var stack = List[LazyModule]()
+
+  def apply[T <: LazyModule](bc: T)(implicit sourceInfo: SourceInfo): T = {
+    // Make sure the user put LazyModule around modules in the correct order
+    // If this require fails, probably some grandchild was missing a LazyModule
+    // ... or you applied LazyModule twice
+    require (!stack.isEmpty && (stack.head eq bc))
+    stack = stack.tail
+    bc.info = sourceInfo
+    bc
+  }
+}
+
+abstract class LazyModuleImp(val outer: LazyModule) extends Module
+{
+  // .module had better not be accessed while LazyModules are still being built!
+  require (LazyModule.stack.isEmpty)
+
   override def desiredName = outer.getClass.getName.split('.').last
   outer.instantiate()
 }

uncore/src/main/scala/tilelink2/Legacy.scala

@@ -14,7 +14,7 @@ class TLLegacy(implicit val p: Parameters) extends LazyModule with HasTileLinkPa
   val node = TLClientNode(TLClientParameters(
     sourceId = IdRange(0, 1 << tlClientXactIdBits)))
 
-  lazy val module = Module(new LazyModuleImp(this) with HasTileLinkParameters {
+  lazy val module = new LazyModuleImp(this) with HasTileLinkParameters {
     val p = outer_p
     val io = new Bundle {
       val legacy = new ClientUncachedTileLinkIO()(p).flip
@@ -103,5 +103,5 @@ class TLLegacy(implicit val p: Parameters) extends LazyModule with HasTileLinkPa
     grant.manager_xact_id := out.d.bits.sink
     grant.data            := out.d.bits.data
     grant.addr_beat       := beatCounter
-  })
+  }
 }

uncore/src/main/scala/tilelink2/RegisterRouter.scala

@@ -83,5 +83,5 @@ class TLRegisterRouter[B <: Bundle, M <: LazyModuleImp]
   require (isPow2(size))
   require (size >= 4096)
 
-  lazy val module = Module(moduleBuilder(bundleBuilder(node.bundleIn), this))
+  lazy val module = moduleBuilder(bundleBuilder(node.bundleIn), this)
 }

uncore/src/main/scala/tilelink2/SRAM.scala

@@ -17,7 +17,7 @@ class TLRAM(address: AddressSet, beatBytes: Int = 4) extends LazyModule
   // We require the address range to include an entire beat (for the write mask)
   require ((address.mask & (beatBytes-1)) == beatBytes-1)
 
-  lazy val module = Module(new LazyModuleImp(this) {
+  lazy val module = new LazyModuleImp(this) {
     val io = new Bundle {
       val in = node.bundleIn
     }
@@ -62,5 +62,5 @@ class TLRAM(address: AddressSet, beatBytes: Int = 4) extends LazyModule
         mem.write(memAddress, wdata, in.a.bits.mask.toBools)
       }
     }
-  })
+  }
 }

uncore/src/main/scala/tilelink2/Xbar.scala

@@ -65,7 +65,7 @@ class TLXbar(policy: (Vec[Bool], Bool) => Seq[Bool] = TLXbar.lowestIndex) extend
       TLManagerPortParameters(managers, seq(0).beatBytes)
     })
 
-  lazy val module = Module(new LazyModuleImp(this) {
+  lazy val module = new LazyModuleImp(this) {
     val io = new Bundle {
       val in  = node.bundleIn
       val out = node.bundleOut
@@ -207,5 +207,5 @@ class TLXbar(policy: (Vec[Bool], Bool) => Seq[Bool] = TLXbar.lowestIndex) extend
 
       muxState
     }
-  })
+  }
 }