From 98cc79ff7461064559690e17763708592a1bb8c5 Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Fri, 14 Jun 2024 12:45:44 -0700
Subject: [PATCH] Error when calling define targeting a child of a probe
 (backport #4175) (#4176)

* Error when calling define targeting a child of a probe (#4175)

(cherry picked from commit 4de35818cb73ff7b0c4c2140fe766e7e4230d6be)

# Conflicts:
#	core/src/main/scala/chisel3/internal/package.scala

* Resolve backport conflicts

---------

Co-authored-by: Jack Koenig <koenig@sifive.com>
---
 .../main/scala/chisel3/internal/package.scala   | 17 +++++++++++++++++
 core/src/main/scala/chisel3/probe/package.scala |  1 +
 src/test/scala/chiselTests/ProbeSpec.scala      | 16 ++++++++++++++++
 3 files changed, 34 insertions(+)

diff --git a/core/src/main/scala/chisel3/internal/package.scala b/core/src/main/scala/chisel3/internal/package.scala
index 308b5628420..a0daed9e18a 100644
--- a/core/src/main/scala/chisel3/internal/package.scala
+++ b/core/src/main/scala/chisel3/internal/package.scala
@@ -132,6 +132,23 @@ package object internal {
     case leaf => leaf.probeInfo.nonEmpty
   }
 
+  private[chisel3] def requireNotChildOfProbe(
+    probe:        Data,
+    errorMessage: => String = ""
+  )(
+    implicit sourceInfo: SourceInfo
+  ): Unit = {
+    probe.binding match {
+      case Some(ChildBinding(parent)) =>
+        if (parent.probeInfo.nonEmpty) {
+          val providedMsg = errorMessage // only evaluate by-name argument once
+          val msg = if (providedMsg.isEmpty) "Expected a root of a probe." else providedMsg
+          Builder.error(msg)
+        }
+      case _ => ()
+    }
+  }
+
   // TODO this exists in cats.Traverse, should we just use that?
   private[chisel3] implicit class ListSyntax[A](xs: List[A]) {
     def mapAccumulate[B, C](z: B)(f: (B, A) => (B, C)): (B, List[C]) = {
diff --git a/core/src/main/scala/chisel3/probe/package.scala b/core/src/main/scala/chisel3/probe/package.scala
index 7091f22434a..6057154013c 100644
--- a/core/src/main/scala/chisel3/probe/package.scala
+++ b/core/src/main/scala/chisel3/probe/package.scala
@@ -37,6 +37,7 @@ package object probe extends SourceInfoDoc {
       Builder.error("Cannot define a probe on a non-equivalent type.")
     }
     requireHasProbeTypeModifier(sink, "Expected sink to be a probe.")
+    requireNotChildOfProbe(sink, "Expected sink to be the root of a probe.")
     requireHasProbeTypeModifier(probeExpr, "Expected source to be a probe expression.")
     if (sink.probeInfo.get.writable) {
       requireHasWritableProbeTypeModifier(
diff --git a/src/test/scala/chiselTests/ProbeSpec.scala b/src/test/scala/chiselTests/ProbeSpec.scala
index 9f83c8af6fa..9194147852f 100644
--- a/src/test/scala/chiselTests/ProbeSpec.scala
+++ b/src/test/scala/chiselTests/ProbeSpec.scala
@@ -584,6 +584,22 @@ class ProbeSpec extends ChiselFlatSpec with Utils {
     exc.getMessage should include("Probe width unknown.")
   }
 
+  it should "error trying to define a child of an Aggregate probe" in {
+    val exc = intercept[chisel3.ChiselException] {
+      ChiselStage.emitSystemVerilog(
+        new RawModule {
+          val in = IO(Input(UInt(16.W)))
+          val p = IO(Output(Probe(new Bundle {
+            val a = UInt(16.W)
+          })))
+          define(p.a, ProbeValue(in))
+        },
+        Array("--throw-on-first-error")
+      )
+    }
+    exc.getMessage should include("Expected sink to be the root of a probe.")
+  }
+
   "Probe force/release reg example" should "work in simulator" in {
     // Simple example forcing a register and checking basic behavior.