From 1d12f838833435716fff9cdea892a25330073abc Mon Sep 17 00:00:00 2001
From: Clo91eaf <Clo91eaf@qq.com>
Date: Mon, 8 Jul 2024 20:06:48 +0800
Subject: [PATCH] [difftest] use filter and for_each to iterate over the
 vrf/memory write

---
 difftest/offline/src/json_events.rs | 32 ++++++++++++++---------------
 1 file changed, 15 insertions(+), 17 deletions(-)

diff --git a/difftest/offline/src/json_events.rs b/difftest/offline/src/json_events.rs
index 026f472ad..69bf4381d 100644
--- a/difftest/offline/src/json_events.rs
+++ b/difftest/offline/src/json_events.rs
@@ -133,22 +133,23 @@ pub fn add_rtl_write(
   record_idx_base: usize,
   vrf_shadow: &mut Vec<u8>,
 ) {
-  vrf_write.mask.iter().enumerate().for_each(|(j, &mask)| {
-    if mask {
-      let written_byte = vrf_write.data[j];
-      let vrf_idx = record_idx_base + j;
+  vrf_write.mask.iter().enumerate()
+    .filter(|(_, &mask)| mask) 
+    .for_each(|(offset, _)| {
+      let written_byte = vrf_write.data[offset];
+      let vrf_idx = record_idx_base + offset;
       let original_byte = vrf_shadow[vrf_idx];
 
-      if let Some(record) = se.vrf_access_record.all_writes.get_mut(&(record_idx_base + j)) {
+      if let Some(record) = se.vrf_access_record.all_writes.get_mut(&(record_idx_base + offset)) {
         assert_eq!(
           record.byte,
           written_byte,
-          "[{}] {j}th byte incorrect ({:02X} != {written_byte:02X}) \
+          "[{}] {offset}th byte incorrect ({:02X} != {written_byte:02X}) \
           for vrf write (lane={}, vd={}, offset={}, mask={:?}) \
           [vrf_idx={}] (lsu_idx={}, disasm: {}, pc: {:#x}, bits: {:#x})",
           vrf_write.cycle,
           record.byte,
-          vrf_write.idx, vrf_write.vd, vrf_write.offset, vrf_write.mask, record_idx_base + j,
+          vrf_write.idx, vrf_write.vd, vrf_write.offset, vrf_write.mask, record_idx_base + offset,
           se.lsu_idx, se.disasm, se.pc, se.inst_bits);
         record.executed = true;
       } else if written_byte != original_byte {
@@ -162,7 +163,7 @@ pub fn add_rtl_write(
           vrf_write.mask,
           written_byte,
           original_byte,
-          record_idx_base + j,
+          record_idx_base + offset,
           se.issue_idx,
           se.describe_insn(),
         );
@@ -170,7 +171,6 @@ pub fn add_rtl_write(
 
       info!("vrf_shadow[{vrf_idx}] {original_byte} -> {written_byte}");
       vrf_shadow[vrf_idx] = written_byte;
-    }
   })
 }
 
@@ -269,10 +269,8 @@ impl JsonEventRunner for SpikeRunner {
   }
 
   fn peek_memory_write(&mut self, memory_write: &MemoryWriteEvent) -> anyhow::Result<()> {
-    let mut mask = memory_write.mask.to_owned();
-    mask.resize(self.dlen as usize / 8, false);
-    let mut data = memory_write.data.to_owned();
-    data.resize(self.dlen as usize / 8, 0u8);
+    let data = memory_write.data.to_owned();
+    let mask = memory_write.mask.to_owned();
     let cycle = memory_write.cycle;
     let base_addr = memory_write.address;
     let source = memory_write.source;
@@ -280,8 +278,9 @@ impl JsonEventRunner for SpikeRunner {
     if let Some(se) = self.to_rtl_queue.iter_mut().find(|se| se.lsu_idx == source) {
       info!("[{cycle}] MemoryWrite: address={base_addr:08x}, size={}, data={data:?}, mask={mask:?}, pc = {:#x}, disasm = {}", data.len(), se.pc, se.disasm);
       // compare with spike event record
-      for offset in 0..data.len() {
-        if mask[offset] {
+      mask.iter().enumerate()
+        .filter(|(_, &mask)| mask)
+        .for_each(|(offset, _)| {
           let byte_addr = base_addr + offset as u32;
           let data_byte = data[offset];
           let mem_write =
@@ -291,8 +290,7 @@ impl JsonEventRunner for SpikeRunner {
           let single_mem_write_val = mem_write.writes[mem_write.num_completed_writes].val;
           mem_write.num_completed_writes += 1;
           assert_eq!(single_mem_write_val, data_byte, "[{cycle}] expect mem write of byte {single_mem_write_val:02X}, actual byte {data_byte:02X} (byte_addr={byte_addr:08X}, pc = {:#x}, disasm = {})", se.pc, se.disasm);
-        }
-      }
+        });
       return Ok(());
     }