Fix using Definitions as arguments to Definitions (backport #3726) (#…

…3731)

* Fix using Definitions as arguments to Definitions (#3726)

(cherry picked from commit a050b8c)

# Conflicts:
#	core/src/main/scala/chisel3/Module.scala
#	core/src/main/scala/chisel3/experimental/hierarchy/core/Definition.scala
#	core/src/main/scala/chisel3/experimental/hierarchy/core/Instance.scala
#	core/src/main/scala/chisel3/internal/Builder.scala

* Resolve backport conflicts

---------

Co-authored-by: Jack Koenig <[email protected]>

core/src/main/scala/chisel3/Module.scala

@@ -280,8 +280,15 @@ package experimental {
   object BaseModule {
     implicit class BaseModuleExtensions[T <: BaseModule](b: T) {
       import chisel3.experimental.hierarchy.core.{Definition, Instance}
-      def toInstance:   Instance[T] = new Instance(Proto(b))
-      def toDefinition: Definition[T] = new Definition(Proto(b))
+      def toInstance: Instance[T] = new Instance(Proto(b))
+      def toDefinition: Definition[T] = {
+        val result = new Definition(Proto(b))
+        // .toDefinition is sometimes called in Select APIs outside of Chisel elaboration
+        if (Builder.inContext) {
+          Builder.definitions += result
+        }
+        result
+      }
     }
   }
 

core/src/main/scala/chisel3/experimental/hierarchy/core/Definition.scala

@@ -101,7 +101,13 @@ object Definition extends SourceInfoDoc {
   ): Definition[T] = {
     val dynamicContext = {
       val context = Builder.captureContext()
-      new DynamicContext(Nil, context.throwOnFirstError, context.warningFilters, context.sourceRoots)
+      new DynamicContext(
+        Nil,
+        context.throwOnFirstError,
+        context.warningFilters,
+        context.sourceRoots,
+        Builder.allDefinitions
+      )
     }
     Builder.globalNamespace.copyTo(dynamicContext.globalNamespace)
     dynamicContext.inDefinition = true
@@ -110,7 +116,7 @@ object Definition extends SourceInfoDoc {
     Builder.annotations ++= ir.annotations: @nowarn // this will go away when firrtl is merged
     module._circuit = Builder.currentModule
     dynamicContext.globalNamespace.copyTo(Builder.globalNamespace)
-    new Definition(Proto(module))
+    module.toDefinition
   }
 
 }

core/src/main/scala/chisel3/experimental/hierarchy/core/Instance.scala

@@ -5,13 +5,14 @@ package chisel3.experimental.hierarchy.core
 import scala.language.experimental.macros
 import chisel3._
 import chisel3.experimental.hierarchy.{InstantiableClone, ModuleClone}
-import chisel3.internal.{throwException, Builder}
+import chisel3.internal.{throwException, BaseBlackBox, Builder}
 import chisel3.experimental.{BaseModule, ExtModule, SourceInfo, UnlocatableSourceInfo}
 import chisel3.internal.sourceinfo.InstanceTransform
 import chisel3.internal.firrtl.{Component, DefBlackBox, DefIntrinsicModule, DefModule, Port}
 import firrtl.annotations.IsModule
 
 import scala.annotation.nowarn
+import chisel3.experimental.BaseIntrinsicModule
 
 /** User-facing Instance type.
   * Represents a unique instance of type `A` which are marked as @instantiable
@@ -120,20 +121,20 @@ object Instance extends SourceInfoDoc {
     implicit sourceInfo: SourceInfo
   ): Instance[T] = {
     // Check to see if the module is already defined internally or externally
-    val existingMod = Builder.components.map {
-      case c: DefModule if c.id == definition.proto                 => Some(c)
-      case c: DefBlackBox if c.name == definition.proto.name        => Some(c)
-      case c: DefIntrinsicModule if c.name == definition.proto.name => Some(c)
-      case _ => None
-    }.flatten
-
-    if (existingMod.isEmpty) {
+    val existingMod = Builder.allDefinitions.view.flatten.map(_.proto).exists {
+      case c: RawModule           => c == definition.proto
+      case c: BaseBlackBox        => c.name == definition.proto.name
+      case c: BaseIntrinsicModule => c.name == definition.proto.name
+      case _ => false
+    }
+
+    if (!existingMod) {
       // Add a Definition that will get emitted as an ExtModule so that FIRRTL
       // does not complain about a missing element
       val extModName = Builder.importedDefinitionMap.getOrElse(
         definition.proto.name,
         throwException(
-          "Imported Definition information not found - possibly forgot to add ImportDefinition annotation?"
+          s"Imported Definition information not found for ${definition.proto.name} - possibly forgot to add ImportDefinition annotation?"
         )
       )
       class EmptyExtModule extends ExtModule {

core/src/main/scala/chisel3/internal/Builder.scala

@@ -6,7 +6,7 @@ import scala.util.DynamicVariable
 import scala.collection.mutable.ArrayBuffer
 import chisel3._
 import chisel3.experimental._
-import chisel3.experimental.hierarchy.core.{Clone, ImportDefinitionAnnotation, Instance}
+import chisel3.experimental.hierarchy.core.{Clone, Definition, ImportDefinitionAnnotation, Instance}
 import chisel3.internal.firrtl._
 import chisel3.internal.naming._
 import _root_.firrtl.annotations.{CircuitName, ComponentName, IsMember, ModuleName, Named, ReferenceTarget}
@@ -418,7 +418,9 @@ private[chisel3] class DynamicContext(
   val annotationSeq:     AnnotationSeq,
   val throwOnFirstError: Boolean,
   val warningFilters:    Seq[WarningFilter],
-  val sourceRoots:       Seq[File]) {
+  val sourceRoots:       Seq[File],
+  // Definitions from other scopes in the same elaboration, use allDefinitions below
+  val outerScopeDefinitions: List[Iterable[Definition[_]]]) {
   val importedDefinitionAnnos = annotationSeq.collect { case a: ImportDefinitionAnnotation[_] => a }
 
   // Map from proto module name to ext-module name
@@ -462,6 +464,7 @@ private[chisel3] class DynamicContext(
   }
 
   val components = ArrayBuffer[Component]()
+  val definitions = ArrayBuffer[Definition[_]]()
   val annotations = ArrayBuffer[ChiselAnnotation]()
   val newAnnotations = ArrayBuffer[ChiselMultiAnnotation]()
   var currentModule: Option[BaseModule] = None
@@ -501,6 +504,9 @@ private[chisel3] object Builder extends LazyLogging {
     dynamicContextVar.value.get
   }
 
+  /** Check if we are in a Builder context */
+  def inContext: Boolean = dynamicContextVar.value.isDefined
+
   // Used to suppress warnings when casting from a UInt to an Enum
   var suppressEnumCastWarning: Boolean = false
 
@@ -536,6 +542,10 @@ private[chisel3] object Builder extends LazyLogging {
   def globalIdentifierNamespace: Namespace = dynamicContext.globalIdentifierNamespace
   def components:                ArrayBuffer[Component] = dynamicContext.components
   def annotations:               ArrayBuffer[ChiselAnnotation] = dynamicContext.annotations
+  def definitions:               ArrayBuffer[Definition[_]] = dynamicContext.definitions
+
+  /** All definitions from current elaboration, including Definitions passed as an argument to this one */
+  def allDefinitions: List[Iterable[Definition[_]]] = definitions :: dynamicContext.outerScopeDefinitions
 
   def contextCache: BuilderContextCache = dynamicContext.contextCache
 

src/main/scala/chisel3/aop/injecting/InjectingAspect.scala

@@ -62,7 +62,8 @@ abstract class InjectorAspect[T <: RawModule, M <: RawModule](
           annotationsInAspect,
           chiselOptions.throwOnFirstError,
           chiselOptions.warningFilters,
-          chiselOptions.sourceRoots
+          chiselOptions.sourceRoots,
+          Nil // FIXME this maybe should somehow grab definitions from earlier elaboration
         )
       // Add existing module names into the namespace. If injection logic instantiates new modules
       //  which would share the same name, they will get uniquified accordingly

src/main/scala/chisel3/stage/phases/Elaborate.scala

@@ -34,7 +34,8 @@ class Elaborate extends Phase {
             annotations,
             chiselOptions.throwOnFirstError,
             chiselOptions.warningFilters,
-            chiselOptions.sourceRoots
+            chiselOptions.sourceRoots,
+            Nil
           )
         val (circuit, dut) =
           Builder.build(Module(gen()), context)

src/test/scala/chiselTests/experimental/hierarchy/DefinitionSpec.scala

@@ -90,6 +90,34 @@ class DefinitionSpec extends ChiselFunSpec with Utils {
       chirrtl.serialize should include("module AddOneWithNested_2 :")
       chirrtl.serialize should include("module AddOneWithNested_3 :")
     }
+
+    it("(0.g): definitions should work as arguments to definitions") {
+      class Top extends Module {
+        val addOne = Definition(new AddOne)
+        val addTwo = Definition(new AddTwoDefinitionArgument(addOne))
+        val inst = Instance(addTwo)
+        inst.in := 12.U
+      }
+      val chirrtl = getFirrtlAndAnnos(new Top)._1.serialize
+      chirrtl should include("module AddOne :")
+      chirrtl shouldNot include("module AddOne_")
+      chirrtl should include("module AddTwoDefinitionArgument :")
+      chirrtl should include("module Top :")
+    }
+    it("(0.h): definitions created from Modules should work as arguments to definitions") {
+      class Top extends Module {
+        val addOne = Module(new AddOne)
+        val addTwo = Definition(new AddTwoDefinitionArgument(addOne.toDefinition))
+        val inst = Instance(addTwo)
+        inst.in := 12.U
+      }
+      val chirrtl = getFirrtlAndAnnos(new Top)._1.serialize
+      chirrtl should include("module AddOne :")
+      chirrtl shouldNot include("module AddOne_")
+      chirrtl should include("module AddTwoDefinitionArgument :")
+      chirrtl should include("module Top :")
+      chirrtl should include("inst addOne of AddOne")
+    }
   }
   describe("(1): Annotations on definitions in same chisel compilation") {
     it("(1.a): should work on a single definition, annotating the definition") {

src/test/scala/chiselTests/experimental/hierarchy/Examples.scala

@@ -93,6 +93,16 @@ object Examples {
     val out = IO(Output(UInt(width.W)))
     val addOnes = makeParameterizedOnes(width)
   }
+  @instantiable
+  class AddTwoDefinitionArgument(definition: Definition[AddOne]) extends Module {
+    @public val in = IO(Input(UInt(32.W)))
+    @public val out = IO(Output(UInt(32.W)))
+    @public val i0: Instance[AddOne] = Instance(definition)
+    @public val i1: Instance[AddOne] = Instance(definition)
+    i0.in := in
+    i1.in := i0.out
+    out := i1.out
+  }
 
   @instantiable
   class AddFour extends Module {