From ad6ee736926d5a658c3ab8c8f2d7711e191531c4 Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Mon, 25 Nov 2024 10:43:29 -0800
Subject: [PATCH] Handle define on views of Probes (#4308) (#4311)

(cherry picked from commit 73e96ea5db53ad176e6a9011fc42c888c901d7e2)

Co-authored-by: Jack Koenig <koenig@sifive.com>
---
 .../main/scala/chisel3/probe/package.scala    | 16 +++++++++++-----
 .../chiselTests/experimental/DataView.scala   | 19 ++++++++++++++++++-
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/core/src/main/scala/chisel3/probe/package.scala b/core/src/main/scala/chisel3/probe/package.scala
index a9194f23dfd..5f8c98f44cb 100644
--- a/core/src/main/scala/chisel3/probe/package.scala
+++ b/core/src/main/scala/chisel3/probe/package.scala
@@ -8,6 +8,7 @@ import chisel3.internal.Builder.pushCommand
 import chisel3.internal.firrtl.ir._
 import chisel3.Data.ProbeInfo
 import chisel3.experimental.{requireIsHardware, SourceInfo}
+import chisel3.experimental.dataview.reifyIdentityView
 import chisel3.reflect.DataMirror.{checkTypeEquivalence, collectAllMembers, hasProbeTypeModifier}
 
 import scala.language.experimental.macros
@@ -33,7 +34,12 @@ package object probe extends SourceInfoDoc {
     * @param probeExpr value to initialize the sink to
     */
   def define[T <: Data](sink: T, probeExpr: T)(implicit sourceInfo: SourceInfo): Unit = {
-    val typeCheckResult = sink.findFirstTypeMismatch(
+    val (realSink, writable) = reifyIdentityView(sink).getOrElse {
+      Builder.error(s"Define only supports identity views for the sink, $sink has multiple targets.")
+      return // This error is recoverable and this function returns Unit, just continue elaboration.
+    }
+    writable.reportIfReadOnlyUnit(())
+    val typeCheckResult = realSink.findFirstTypeMismatch(
       probeExpr,
       strictTypes = true,
       strictWidths = true,
@@ -42,16 +48,16 @@ package object probe extends SourceInfoDoc {
     typeCheckResult.foreach { msg =>
       Builder.error(s"Cannot define a probe on a non-equivalent type.\n$msg")
     }
-    requireHasProbeTypeModifier(sink, "Expected sink to be a probe.")
-    requireNotChildOfProbe(sink, "Expected sink to be the root of a probe.")
+    requireHasProbeTypeModifier(realSink, "Expected sink to be a probe.")
+    requireNotChildOfProbe(realSink, "Expected sink to be the root of a probe.")
     requireHasProbeTypeModifier(probeExpr, "Expected source to be a probe expression.")
-    if (sink.probeInfo.get.writable) {
+    if (realSink.probeInfo.get.writable) {
       requireHasWritableProbeTypeModifier(
         probeExpr,
         "Cannot use a non-writable probe expression to define a writable probe."
       )
     }
-    pushCommand(ProbeDefine(sourceInfo, sink.lref, probeExpr.ref))
+    pushCommand(ProbeDefine(sourceInfo, realSink.lref, probeExpr.ref))
   }
 
   /** Access the value of a probe.
diff --git a/src/test/scala/chiselTests/experimental/DataView.scala b/src/test/scala/chiselTests/experimental/DataView.scala
index 6a2bceab6fa..a0bcb13093c 100644
--- a/src/test/scala/chiselTests/experimental/DataView.scala
+++ b/src/test/scala/chiselTests/experimental/DataView.scala
@@ -939,7 +939,24 @@ class DataViewSpec extends ChiselFlatSpec {
     chirrtl should include("define a = rwprobe(w)")
   }
 
-  it should "error if attempting to define a viewed a Probe as a RWProbe" in {
+  it should "support defining identity views of Probes" in {
+    class InnerBundle extends Bundle {
+      val a = Bool()
+    }
+    class OuterBundle extends Bundle {
+      val a = Probe(new InnerBundle)
+    }
+    class MyModule extends Module {
+      val p = Wire(new OuterBundle)
+      val view = p.viewAs[OuterBundle]
+      val w = Wire(new InnerBundle)
+      define(view.a, ProbeValue(w))
+    }
+    val chirrtl = ChiselStage.emitCHIRRTL(new MyModule)
+    chirrtl should include("define p.a = probe(w)")
+  }
+
+  it should "error if attempting to define a Probe viewed as a RWProbe" in {
     class MyModule extends Module {
       val a = IO(Output(RWProbe(Bool())))
       val w = WireInit(Bool(), false.B)