[core] Enable returning values from layer blocks

Add support for a limited form of returning values from layer blocks.
With this patch, when a `Data` is returned from a layer block, this will
now cause the creation of a wire of layer-colored probe type that is
driven by a probe of the return value.  There is additional handling to
directly return a probe type as well (which enables returning from an
inner layer block up through an outer layer block).  If a non-`Data` is
returned, then this will be converted to a `Unit`.  This latter behavior
is the current behavior.

Signed-off-by: Schuyler Eldridge <[email protected]>

Author: seldridge

core/src/main/scala/chisel3/Layer.scala

@@ -311,6 +311,15 @@ object layer {
     * be the current layer which causes no layer block to be created.  (This is
     * not a _proper_ ancestor requirement.)
     *
+    * By default, the return of the layer block will be either a subtype of
+    * [[Data]] or [[Unit]], depending on the `thunk` provided.  If the `thunk`
+    * (the hardware that should be constructed inside the layer block) returns a
+    * [[Data]], then this will either return a [[Wire]] of layer-colored
+    * [[Probe]] type if a layer block was created or the underlying [[Data]] if
+    * no layer block was created.  If the `thunk` returns anything else, this
+    * will return [[Unit]].  This is controlled by the implicit argument `tc`
+    * and may be customized by advanced users to do other things.
+    *
     * @param layer the layer this block is associated with
     * @param skipIfAlreadyInBlock if true, then this will not create a layer if
     * this `block` is already inside another layerblock
@@ -321,21 +330,23 @@ object layer {
     * @param sourceInfo a source locator
     * @throws java.lang.IllegalArgumentException if the layer of the currnet
     * layerblock is not an ancestor of the desired layer
+    * @return either a subtype of [[Data]] or [[Unit]] depending on the `thunk`
+    * return type
     */
   def block[A](
     layer:                Layer,
     skipIfAlreadyInBlock: Boolean = false,
     skipIfLayersEnabled:  Boolean = false
   )(thunk: => A)(
-    implicit sourceInfo: SourceInfo
-  ): Unit = {
+    implicit tc: BlockReturnHandler[A] = BlockReturnHandler.unit[A],
+    sourceInfo:  SourceInfo
+  ): tc.R = {
     // Do nothing if we are already in a layer block and are not supposed to
     // create new layer blocks.
     if (
       skipIfAlreadyInBlock && Builder.layerStack.size > 1 || skipIfLayersEnabled && Builder.enabledLayers.nonEmpty || Builder.elideLayerBlocks
     ) {
-      thunk
-      return
+      return tc.identity(thunk)
     }
 
     val _layer = Builder.layerMap.getOrElse(layer, layer)
@@ -350,20 +361,37 @@ object layer {
       s"a layerblock associated with layer '${_layer.fullName}' cannot be created under a layerblock of non-ancestor layer '${Builder.layerStack.head.fullName}'"
     )
 
+    if (layersToCreate.isEmpty)
+      return tc.identity(thunk)
+
     addLayer(_layer)
 
+    // Save the append point _before_ the layer block so that we can insert a
+    // layer-colored wire once the `thunk` executes.
+    val beforeLayerBlock = new Placeholder
+
+    // Track the current layer block.  When this is used, this will be the
+    // innermost layer block that will be created.  This is guaranteed to be
+    // non-null as long as `layersToCreate` is not empty.
+    var layerBlock: LayerBlock = null
+
+    // Recursively create any necessary layers.  There are two cases:
+    //
+    // 1. There are no layers left to create.  Run the thunk, create the
+    //    layer-colored wire, and define it.
+    // 2. There are layers left to create.  Create the next layer and recurse.
     def createLayers(layers: List[Layer])(thunk: => A): A = layers match {
       case Nil => thunk
       case head :: tail =>
-        val layerBlock = new LayerBlock(sourceInfo, head)
-        Builder.pushCommand(layerBlock)
+        layerBlock = Builder.pushCommand(new LayerBlock(sourceInfo, head))
         Builder.layerStack = head :: Builder.layerStack
         val result = Builder.forcedUserModule.withRegion(layerBlock.region)(createLayers(tail)(thunk))
         Builder.layerStack = Builder.layerStack.tail
         result
     }
 
-    createLayers(layersToCreate)(thunk)
+    val result = createLayers(layersToCreate)(thunk)
+    return tc.apply(beforeLayerBlock, layerBlock, result)
   }
 
   /** API that will cause any calls to `block` in the `thunk` to not create new

src/test/scala/chiselTests/LayerSpec.scala

@@ -10,6 +10,7 @@ import chisel3.probe.{define, Probe, ProbeValue}
 import chiselTests.{ChiselFlatSpec, FileCheck, Utils}
 import java.nio.file.{FileSystems, Paths}
 import _root_.circt.stage.ChiselStage
+import chisel3.reflect.DataMirror.internal.chiselTypeClone
 
 class LayerSpec extends ChiselFlatSpec with Utils with FileCheck {
 
@@ -157,10 +158,10 @@ class LayerSpec extends ChiselFlatSpec with Utils with FileCheck {
       """|CHECK:      module Foo :
          |CHECK:        define x = probe(a)
          |CHECK-NEXT:   define y = probe(b)
-         |CHECK-NEXT:   layerblock A :
+         |CHECK:        layerblock A :
          |CHECK-NEXT:     define x = probe(c)
          |CHECK-NEXT:     define y = probe(d)
-         |CHECK-NEXT:     layerblock B :
+         |CHECK:          layerblock B :
          |CHECK-NEXT:       define y = probe(e)
          |""".stripMargin
     }
@@ -306,6 +307,133 @@ class LayerSpec extends ChiselFlatSpec with Utils with FileCheck {
     ChiselStage.emitCHIRRTL(new Foo) should include("""layer ExpensiveAsserts, bind, "verification/ExpensiveAsserts"""")
   }
 
+  "Values returned by layer blocks" should "be layer-colored wires" in {
+    class Foo extends RawModule {
+      val out = IO {
+        Output {
+          new Bundle {
+            val a = Probe(UInt(1.W), A)
+            val b = Probe(UInt(2.W), A.B)
+            val c = Probe(UInt(3.W), A.B)
+            val d = Probe(UInt(4.W), A)
+          }
+        }
+      }
+
+      // A single layer block
+      val a = layer.block(A) {
+        Wire(UInt(1.W))
+      }
+      define(out.a, a)
+
+      // Auto-creation of parent layer blocks
+      val b = layer.block(A.B) {
+        Wire(UInt(2.W))
+      }
+      define(out.b, b)
+
+      // Nested layer blocks
+      val c = layer.block(A) {
+        layer.block(A.B) {
+          Wire(UInt(3.W))
+        }
+      }
+      define(out.c, c)
+
+      // Return of ProbeValue
+      val d = layer.block(A) {
+        ProbeValue(Wire(UInt(4.W)))
+      }
+      define(out.d, d)
+
+      // No layers created.  Check all generator code paths.
+      layer.block(A) {
+        // Path 1: `skipIfAlreadyInBlock` set
+        val e = layer.block(C, skipIfAlreadyInBlock = true) {
+          Wire(UInt(5.W))
+        }
+        // Path 2: requested layer is the current layer
+        val f = layer.block(A) {
+          Wire(UInt(6.W))
+        }
+      }
+      // Path 3: `skipIfLayersEnabled` is set and a layer is enabled
+      layer.enable(C)
+      val g = layer.block(C, skipIfLayersEnabled = true) {
+        Wire(UInt(7.W))
+      }
+      // Path 4: the `elideBlocks` API is used
+      val h = layer.elideBlocks {
+        layer.block(A) {
+          Wire(UInt(8.W))
+        }
+      }
+
+      // Return non-`Data` as `Unit`.
+      val i = layer.block(A) {
+        42
+      }
+      assert(i.getClass() == classOf[Unit])
+
+      // Return `Unit` as `Unit`.
+      val j = layer.block(A) { () }
+      assert(j.getClass() == classOf[Unit])
+    }
+
+    generateFirrtlAndFileCheck(new Foo) {
+      s"""|CHECK:      module Foo
+          |CHECK:        wire a : Probe<UInt<1>, A>
+          |CHECK-NEXT:   layerblock A :
+          |CHECK-NEXT:     wire [[a:.*]] : UInt<1>
+          |CHECK-NEXT:     define a = probe([[a]])
+          |CHECK:        define out.a = a
+          |
+          |CHECK:        wire b : Probe<UInt<2>, A.B>
+          |CHECK-NEXT:   layerblock A :
+          |CHECK-NEXT:     layerblock B :
+          |CHECK-NEXT:       wire [[b:.*]] : UInt<2>
+          |CHECK-NEXT:       define b = probe([[b]])
+          |CHECK:        define out.b = b
+          |
+          |CHECK:        wire c : Probe<UInt<3>, A.B>
+          |CHECK-NEXT:   layerblock A :
+          |CHECK-NEXT:     wire [[c0:.*]] : Probe<UInt<3>, A.B>
+          |CHECK-NEXT:     layerblock B :
+          |CHECK-NEXT:       wire [[c:.*]] : UInt<3>
+          |CHECK-NEXT:       define [[c0]] = probe([[c]])
+          |CHECK:          define c = [[c0]]
+          |CHECK:        define out.c = c
+          |
+          |CHECK:        wire d : Probe<UInt<4>, A>
+          |CHECK-NEXT:   layerblock A :
+          |CHECK-NEXT:     wire [[d:.*]] : UInt<4>
+          |CHECK-NEXT:     define d = probe([[d]])
+          |CHECK:        define out.d = d
+          |
+          |CHECK:        layerblock A :
+          |CHECK-NEXT:     wire e : UInt<5>
+          |CHECK-NEXT:     wire f : UInt<6>
+          |CHECK-NOT:    layerblock
+          |CHECK:        wire g : UInt<7>
+          |CHECK-NOT:    layerblock
+          |CHECK:        wire h : UInt<8>
+          |""".stripMargin
+    }
+  }
+
+  they should "require compatible color if returning a layer-colored wire" in {
+    class Foo extends RawModule {
+
+      val a = layer.block(A) {
+        Wire(Probe(Bool(), C))
+      }
+
+    }
+
+    intercept[ChiselException] { ChiselStage.convert(new Foo, Array("--throw-on-first-error")) }
+      .getMessage() should include("cannot return probe of color 'C' from a layer block associated with layer 'A'")
+  }
+
   "addLayer API" should "add a layer to the output CHIRRTL even if no layer block references that layer" in {
     class Foo extends RawModule {
       layer.addLayer(A)