From b4e5cf575232d16c2becaafa18f22c0cd2285d5a Mon Sep 17 00:00:00 2001
From: Jack Koenig <koenig@sifive.com>
Date: Wed, 22 May 2024 13:38:49 -0700
Subject: [PATCH] Materialize wires for .ref of Aggregate views

Fixes muxing of Aggregate views, and fixes ProbeValues of Aggregate
views.
---
 core/src/main/scala/chisel3/Data.scala        | 11 ++--
 .../chiselTests/experimental/DataView.scala   | 57 +++++++++++++++++++
 2 files changed, 63 insertions(+), 5 deletions(-)

diff --git a/core/src/main/scala/chisel3/Data.scala b/core/src/main/scala/chisel3/Data.scala
index ce96400bb00..8d8816c6949 100644
--- a/core/src/main/scala/chisel3/Data.scala
+++ b/core/src/main/scala/chisel3/Data.scala
@@ -658,11 +658,12 @@ abstract class Data extends HasId with NamedComponent with SourceInfoDoc {
     topBindingOpt match {
       // DataView
       case Some(ViewBinding(target)) => reify(target).ref
-      case Some(AggregateViewBinding(viewMap)) =>
-        viewMap.get(this) match {
-          case None => materializeWire() // FIXME FIRRTL doesn't have Aggregate Init expressions
-          // This should not be possible because Element does the lookup in .topBindingOpt
-          case x: Some[_] => throwException(s"Internal Error: In .ref for $this got '$topBindingOpt' and '$x'")
+      case Some(_: AggregateViewBinding) =>
+        reifySingleData(this) match {
+          // If this is an identity view (a view of something of the same type), return ref of target
+          case Some(target) if this.typeEquivalent(target) => target.ref
+          // Otherwise, we need to materialize hardware of the correct type
+          case _ => materializeWire()
         }
       // Literals
       case Some(ElementLitBinding(litArg)) => litArg
diff --git a/src/test/scala/chiselTests/experimental/DataView.scala b/src/test/scala/chiselTests/experimental/DataView.scala
index 1b503f8c54d..fb95555bcd8 100644
--- a/src/test/scala/chiselTests/experimental/DataView.scala
+++ b/src/test/scala/chiselTests/experimental/DataView.scala
@@ -554,6 +554,63 @@ class DataViewSpec extends ChiselFlatSpec {
     verilog should include("assign z = sel ? b : d;")
   }
 
+  it should "support muxing between views of Bundles" in {
+    import SimpleBundleDataView._
+    class MyModule extends Module {
+      val cond = IO(Input(Bool()))
+      val in1 = IO(Input(new BundleA(8)))
+      val in2 = IO(Input(new BundleA(8)))
+      val out = IO(Output(new BundleB(8)))
+
+      out := Mux(cond, in1.viewAs[BundleB], in2.viewAs[BundleB])
+    }
+    val firrtl = ChiselStage.emitCHIRRTL(new MyModule)
+    val lines = Seq(
+      "wire _out_WIRE : { bar : UInt<8>}",
+      "connect _out_WIRE.bar, in1.foo",
+      "wire _out_WIRE_1 : { bar : UInt<8>}",
+      "connect _out_WIRE_1.bar, in2.foo",
+      "node _out_T = mux(cond, _out_WIRE, _out_WIRE_1)"
+    )
+    for (line <- lines) {
+      firrtl should include(line)
+    }
+  }
+
+  it should "not generate extra wires when muxing between identity views of Bundles" in {
+    import SimpleBundleDataView._
+    class MyModule extends Module {
+      val cond = IO(Input(Bool()))
+      val in1 = IO(Input(new BundleA(8)))
+      val in2 = IO(Input(new BundleA(8)))
+      val out = IO(Output(new BundleA(8)))
+
+      out := Mux(cond, in1.viewAs[BundleA], in2.viewAs[BundleA])
+    }
+    val firrtl = ChiselStage.emitCHIRRTL(new MyModule)
+    firrtl should include("node _out_T = mux(cond, in1, in2)")
+    firrtl shouldNot include("wire")
+  }
+
+  it should "handle Probe of a view of a Bundle" in {
+    import SimpleBundleDataView._
+    class MyModule extends Module {
+      val in = IO(Input(new BundleA(8)))
+      val out_probe = IO(Output(Probe(new BundleB(8))))
+      val view = in.viewAs[BundleB]
+      define(out_probe, ProbeValue(view))
+    }
+    val firrtl = ChiselStage.emitCHIRRTL(new MyModule)
+    val lines = Seq(
+      "wire _WIRE : { bar : UInt<8>}",
+      "connect _WIRE.bar, in.foo",
+      "define out_probe = probe(_WIRE)"
+    )
+    for (line <- lines) {
+      firrtl should include(line)
+    }
+  }
+
   it should "support primitive types in the view target" in {
     class MyBundle(val name: String) extends Bundle {
       val x = UInt(8.W)