From b064eb3994bc1b8fe03fac87c9307f6552e6d367 Mon Sep 17 00:00:00 2001
From: David CARLIER <devnexen@gmail.com>
Date: Mon, 24 Jul 2023 12:14:07 +0100
Subject: [PATCH] read_time_counter port for the RISCV family. (#1378)

---
 libafl/src/bolts/cpu.rs | 55 +++++++++++++++++++++++++++++++++++++++--
 1 file changed, 53 insertions(+), 2 deletions(-)

diff --git a/libafl/src/bolts/cpu.rs b/libafl/src/bolts/cpu.rs
index 2a9c46dc55..d5a409e220 100644
--- a/libafl/src/bolts/cpu.rs
+++ b/libafl/src/bolts/cpu.rs
@@ -3,7 +3,13 @@
 #[cfg(target_arch = "aarch64")]
 use core::arch::asm;
 
-#[cfg(not(any(target_arch = "x86_64", target_arch = "x86", target_arch = "aarch64")))]
+#[cfg(not(any(
+    target_arch = "x86_64",
+    target_arch = "x86",
+    target_arch = "aarch64",
+    target_arch = "riscv64",
+    target_arsch = "riscv32"
+)))]
 use crate::bolts::current_nanos;
 
 // TODO: Add more architectures, using C code, see
@@ -44,13 +50,58 @@ pub fn read_time_counter() -> u64 {
     v
 }
 
+/// Read a timestamp for measurements
+///
+/// Fetches the full 64 bits of the cycle counter in one instruction.
+#[cfg(target_arch = "riscv64")]
+#[must_use]
+pub fn read_time_counter() -> u64 {
+    let mut v: u64;
+    unsafe {
+        asm!("rdcycle {v}", v = out(reg) v);
+    }
+    v
+}
+
+/// Read a timestamp for measurements
+///
+/// Fetches the high 32 bits of the cycle counter, its low end
+/// If the high part has changed, we branch again.
+/// FIXME: see if the latter is overkill.
+#[cfg(target_arch = "riscv32")]
+#[must_use]
+pub fn read_time_counter() -> u64 {
+    let mut v: u64;
+    let mut hg: u32;
+    let mut lw: u32;
+    let mut cmp: u32;
+    unsafe {
+        asm!("jmp%=:",
+             "rdcycleh {hg}",
+             "rdcycle {lw}",
+             "rdcycleh {cmp}",
+             "bne {hg} {cmp}, jmp%=",
+             hg = out(reg) hg,
+             lw = out(reg) lw,
+             cmp = out(reg) cmp);
+        v = ((hg as u64) << 32) | lw as u64;
+    }
+    v
+}
+
 /// Read a timestamp for measurements.
 ///
 /// This function is a wrapper around different ways to get a timestamp, fast.
 /// In this way, an experiment only has to
 /// change this implementation rather than every instead of [`read_time_counter`]
 /// On unsupported architectures, it's falling back to normal system time, in millis.
-#[cfg(not(any(target_arch = "x86_64", target_arch = "x86", target_arch = "aarch64")))]
+#[cfg(not(any(
+    target_arch = "x86_64",
+    target_arch = "x86",
+    target_arch = "aarch64",
+    target_arch = "riscv64",
+    target_arch = "riscv32"
+)))]
 #[must_use]
 pub fn read_time_counter() -> u64 {
     current_nanos()