From 09ad1b2d4d4e26220cc645d8ca389bcd68d485b2 Mon Sep 17 00:00:00 2001
From: Yinan Xu <xuyinan@ict.ac.cn>
Date: Sat, 20 Jan 2024 18:57:02 +0800
Subject: [PATCH] Support needUpdate and defaultUpdate for DifftestBaseBundle
 (#262)

The valid bits are actually defined in the base bundle.
Thus, we move them from DifftestBundle to DifftestBaseBundle.
---
 src/main/scala/Bundles.scala  | 15 +++++++++++++++
 src/main/scala/Difftest.scala |  3 ---
 src/main/scala/Gateway.scala  |  2 +-
 src/main/scala/Squash.scala   |  2 +-
 4 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/src/main/scala/Bundles.scala b/src/main/scala/Bundles.scala
index 9d6e3bb7f..576a0da3c 100644
--- a/src/main/scala/Bundles.scala
+++ b/src/main/scala/Bundles.scala
@@ -38,6 +38,14 @@ sealed trait DifftestBaseBundle extends Bundle {
     }
   }
 
+  def needUpdate: Option[Bool] = if (hasValid) Some(getValid) else None
+  def defaultUpdate(): DifftestBaseBundle = {
+    if (hasValid) {
+      getValid := false.B
+    }
+    this
+  }
+
   def hasAddress: Boolean = this.isInstanceOf[HasAddress]
   def getNumElements: Int = {
     this match {
@@ -89,6 +97,13 @@ class TrapEvent extends DifftestBaseBundle {
 
   val code = UInt(3.W)
   val pc = UInt(64.W)
+
+  override def needUpdate: Option[Bool] = Some(hasTrap || hasWFI)
+  override def defaultUpdate(): TrapEvent = {
+    hasTrap := false.B
+    hasWFI := false.B
+    this
+  }
 }
 
 class CSRState extends DifftestBaseBundle {
diff --git a/src/main/scala/Difftest.scala b/src/main/scala/Difftest.scala
index 7210d296b..455e23c4b 100644
--- a/src/main/scala/Difftest.scala
+++ b/src/main/scala/Difftest.scala
@@ -56,8 +56,6 @@ sealed trait DifftestBundle extends Bundle with DifftestWithCoreid { this: Difft
     }
   }
 
-  def needUpdate: Option[Bool] = if (hasValid) Some(getValid) else None
-
   protected val needFlatten: Boolean = false
   def isFlatten: Boolean = hasAddress && this.needFlatten
 
@@ -142,7 +140,6 @@ class DiffInstrCommit(nPhyRegs: Int = 32) extends InstrCommit(nPhyRegs)
 
 class DiffTrapEvent extends TrapEvent with DifftestBundle {
   override val desiredCppName: String = "trap"
-  override def needUpdate: Option[Bool] = Some(hasTrap || hasWFI)
   override def supportsSquashBase: Bool = !hasTrap && !hasWFI
 }
 
diff --git a/src/main/scala/Gateway.scala b/src/main/scala/Gateway.scala
index c747cd4c5..35e950673 100644
--- a/src/main/scala/Gateway.scala
+++ b/src/main/scala/Gateway.scala
@@ -183,7 +183,7 @@ class GatewayEndpoint(signals: Seq[DifftestBundle], config: GatewayConfig) exten
 
   val global_enable = WireInit(true.B)
   if(config.hasGlobalEnable) {
-    global_enable := VecInit(out.filter(_.needUpdate.isDefined).map(_.needUpdate.get).toSeq).asUInt.orR
+    global_enable := VecInit(out.flatMap(_.bits.needUpdate).toSeq).asUInt.orR
   }
 
   val batch_data = Option.when(config.isBatch)(Mem(config.batchSize, out_pack.cloneType))
diff --git a/src/main/scala/Squash.scala b/src/main/scala/Squash.scala
index d5dd7274d..3b9f6b17d 100644
--- a/src/main/scala/Squash.scala
+++ b/src/main/scala/Squash.scala
@@ -96,7 +96,7 @@ class SquashEndpoint(bundles: Seq[DifftestBundle], config: GatewayConfig) extend
     // ignore useless unsquashed data when hasGlobalEnable
     val needStore = WireInit(true.B)
     if (config.hasGlobalEnable) {
-      needStore := VecInit(in.filter(_.needUpdate.isDefined).map(_.needUpdate.get).toSeq).asUInt.orR
+      needStore := VecInit(in.flatMap(_.bits.needUpdate).toSeq).asUInt.orR
     }
     when ((should_tick || do_squash.asUInt.orR) && needStore && !control.replay.get) {
       replay_data(replay_ptr) := in