From b058bdcf8c142641aacd92f9938fb350a90e0762 Mon Sep 17 00:00:00 2001
From: Aditya Naik
Date: Wed, 24 Jul 2024 02:05:26 -0700
Subject: Working on Module.scala

---
 core/src/main/scala/chisel3/Module.scala           | 30 +++++++++++++++++++++-
 core/src/main/scala/chisel3/internal/Builder.scala | 21 +++++++++++++++
 .../main/scala/chisel3/internal/firrtl/IR.scala    |  2 ++
 3 files changed, 52 insertions(+), 1 deletion(-)

(limited to 'core/src')

diff --git a/core/src/main/scala/chisel3/Module.scala b/core/src/main/scala/chisel3/Module.scala
index 727a1320..18749afa 100644
--- a/core/src/main/scala/chisel3/Module.scala
+++ b/core/src/main/scala/chisel3/Module.scala
@@ -3,7 +3,7 @@
 package chisel3
 
 import scala.collection.immutable.ListMap
-import scala.collection.mutable.{ArrayBuffer, HashMap}
+import scala.collection.mutable.{ArrayBuffer, HashMap, LinkedHashSet}
 
 import chisel3.internal._
 import chisel3.internal.Builder._
@@ -22,6 +22,34 @@ object Module {
     * @return the input module `m` with Chisel metadata properly set
     */
   def apply[T <: BaseModule](bc: => T): T = {
+    // Instantiate the module definition.
+    val module = evaluate[T](bc)
+
+    // Handle connections at enclosing scope
+    // We use _component because Modules that don't generate them may still have one
+    if (Builder.currentModule.isDefined && module._component.isDefined) {
+      // Class only uses the Definition API, and is not allowed here.
+      module match {
+        case _: Class[_] => throwException("Module() cannot be called on a Class. Please use Definition().")
+        case _ => ()
+      }
+
+      val component = module._component.get
+      component match {
+        case DefClass(_, name, _, _) =>
+          Builder.referenceUserContainer match {
+            case rm:  RawModule => rm.addCommand(DefObject(module, name))
+          }
+        case _ => pushCommand(DefInstance(module, component.ports))
+      }
+      module.initializeInParent()
+    }
+
+    module
+  }
+
+  /** Build a module definition */
+  private[chisel3] def evaluate[T <: BaseModule](bc: => T): T = {
     if (Builder.readyForModuleConstr) {
       throwException(
         "Error: Called Module() twice without instantiating a Module."
diff --git a/core/src/main/scala/chisel3/internal/Builder.scala b/core/src/main/scala/chisel3/internal/Builder.scala
index 5782a2f5..1c78f322 100644
--- a/core/src/main/scala/chisel3/internal/Builder.scala
+++ b/core/src/main/scala/chisel3/internal/Builder.scala
@@ -623,6 +623,27 @@ private[chisel3] object Builder extends LazyLogging {
       )
   }
 
+  def referenceUserModule: RawModule = {
+    currentModule match {
+      case Some(module: RawModule) => module
+      case _ =>
+        throwException(
+          "Error: Not in a RawModule. Likely cause: Missed Module() wrap, bare chisel API call, or attempting to construct hardware inside a BlackBox."
+          // A bare api call is, e.g. calling Wire() from the scala console).
+        )
+    }
+  }
+  def referenceUserContainer: BaseModule = {
+    currentModule match {
+      case Some(module: RawModule) => module
+      case _ =>
+        throwException(
+          "Error: Not in a RawModule. Likely cause: Missed Module() or Definition() wrap, bare chisel API call, or attempting to construct hardware inside a BlackBox."
+          // A bare api call is, e.g. calling Wire() from the scala console).
+        )
+    }
+  }
+
   def forcedUserModule: RawModule = currentModule match {
     case Some(module: RawModule) => module
     case _ =>
diff --git a/core/src/main/scala/chisel3/internal/firrtl/IR.scala b/core/src/main/scala/chisel3/internal/firrtl/IR.scala
index 6789b041..de3d8716 100644
--- a/core/src/main/scala/chisel3/internal/firrtl/IR.scala
+++ b/core/src/main/scala/chisel3/internal/firrtl/IR.scala
@@ -325,6 +325,7 @@ case class DefMemPort[T <: Data](
   clock:      Arg)
     extends Definition
 @nowarn("msg=class Port") // delete when Port becomes private
+case class DefObject(id: HasId, className: String) extends Definition
 case class DefInstance(id: BaseModule, ports: Seq[Port]) extends Definition
 case class WhenBegin(pred: Arg) extends Command
 case class WhenEnd(firrtlDepth: Int, hasAlt: Boolean = false) extends Command
@@ -379,6 +380,7 @@ case class Circuit(
   ) = Circuit(name, components, annotations, renames, newAnnotations)
 
 }
+case class DefClass(id: Class[_], name: String, ports: Seq[Port], commands: Seq[Command]) extends Component
 object Circuit
     extends scala.runtime.AbstractFunction4[String, Seq[Component], Seq[ChiselAnnotation], RenameMap, Circuit] {
   def unapply(c: Circuit): Option[(String, Seq[Component], Seq[ChiselAnnotation], RenameMap)] = {
-- 
cgit v1.2.3