Replace AtomicU64 with u64

AtomicU64 causes compilation issues when used as a subcrate of another
crate when the toolchain configurations conflict. The problem was observed
with ID_COUNTER when Polkadot SDK was compiled for riscv32emac, which could
not be found from core.

The documentation also states that:

"This type is only available on platforms that support atomic loads and
stores of u64."

Further, register state itself is an atomic state. Using atomic for
different registers does not guarantee this any better than u64.

RwLock from the crate Spin is used, which is pretty standard choice in
firmware, shim and similar artifacts (ad-hoc Lock has bee denoted but I
think this is just a better idea).

Signed-off-by: Jarkko Sakkinen <[email protected]>

Cargo.toml

@@ -60,6 +60,7 @@ ruzstd = { version = "0.4.0", default-features = false }
 schnellru = { version = "0.2.3" }
 serde = { version = "1.0.203", features = ["derive"] }
 serde_json = { version = "1.0.117" }
+spin = { version = "0.9.8", default-features = false, features = ["lock_api", "spin_mutex", "rwlock", "lazy"] }
 syn = "2.0.25"
 yansi = "0.5.1"
 

crates/polkavm-common/Cargo.toml

@@ -12,6 +12,7 @@ description = "The common crate for PolkaVM"
 log = { workspace = true, optional = true }
 polkavm-assembler = { workspace = true, optional = true }
 blake3 = { workspace = true, optional = true }
+spin = { workspace = true }
 
 [features]
 default = []

crates/polkavm-common/src/program.rs

@@ -4,6 +4,13 @@ use crate::varint::{read_simple_varint, read_varint, write_simple_varint, write_
 use core::fmt::Write;
 use core::ops::Range;
 
+#[cfg(feature = "unique-id")]
+use spin::RwLock;
+#[cfg(feature = "unique-id")]
+struct UniqueId(u64);
+#[cfg(feature = "unique-id")]
+static ID_COUNTER: RwLock<UniqueId> = RwLock::new(UniqueId(0));
+
 #[derive(Copy, Clone)]
 #[repr(transparent)]
 pub struct RawReg(u32);
@@ -3980,8 +3987,10 @@ impl ProgramBlob {
 
         #[cfg(feature = "unique-id")]
         {
-            static ID_COUNTER: core::sync::atomic::AtomicU64 = core::sync::atomic::AtomicU64::new(0);
-            blob.unique_id = ID_COUNTER.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
+            let mut counter = ID_COUNTER.write();
+            *counter = UniqueId(counter.0 + 1);
+            blob.unique_id = counter.0;
+            // The lock is dropper here.
         }
 
         Ok(blob)

crates/polkavm-common/src/zygote.rs

@@ -4,7 +4,7 @@
 //! is recompiled.
 
 use core::cell::UnsafeCell;
-use core::sync::atomic::{AtomicBool, AtomicI64, AtomicU32, AtomicU64};
+use core::sync::atomic::{AtomicBool, AtomicU32};
 
 // Due to the limitations of Rust's compile time constant evaluation machinery
 // we need to define this struct multiple times.
@@ -137,24 +137,24 @@ pub const VM_SANDBOX_MAXIMUM_NATIVE_CODE_SIZE: u32 = 2176 * 1024 * 1024 - 1;
 
 #[repr(C)]
 pub struct JmpBuf {
-    pub rip: AtomicU64,
-    pub rbx: AtomicU64,
-    pub rsp: AtomicU64,
-    pub rbp: AtomicU64,
-    pub r12: AtomicU64,
-    pub r13: AtomicU64,
-    pub r14: AtomicU64,
-    pub r15: AtomicU64,
+    pub rip: u64,
+    pub rbx: u64,
+    pub rsp: u64,
+    pub rbp: u64,
+    pub r12: u64,
+    pub r13: u64,
+    pub r14: u64,
+    pub r15: u64,
 }
 
 #[repr(C)]
 pub struct VmInit {
-    pub stack_address: AtomicU64,
-    pub stack_length: AtomicU64,
-    pub vdso_address: AtomicU64,
-    pub vdso_length: AtomicU64,
-    pub vvar_address: AtomicU64,
-    pub vvar_length: AtomicU64,
+    pub stack_address: u64,
+    pub stack_length: u64,
+    pub vdso_address: u64,
+    pub vdso_length: u64,
+    pub vvar_address: u64,
+    pub vvar_length: u64,
 
     /// Whether userfaultfd-based memory management is available.
     pub uffd_available: AtomicBool,
@@ -230,15 +230,15 @@ pub struct VmCtx {
     _align_1: CacheAligned<()>,
 
     /// The current gas counter.
-    pub gas: AtomicI64,
+    pub gas: i64,
 
     _align_2: CacheAligned<()>,
 
     /// The futex used to synchronize the sandbox with the host process.
     pub futex: AtomicU32,
 
     /// Address to which to jump to.
-    pub jump_into: AtomicU64,
+    pub jump_into: u64,
 
     /// The address of the instruction currently being executed.
     pub program_counter: AtomicU32,
@@ -253,31 +253,31 @@ pub struct VmCtx {
     pub arg: AtomicU32,
 
     /// A dump of all of the registers of the VM.
-    pub regs: [AtomicU64; REG_COUNT],
+    pub regs: [u64; REG_COUNT],
 
     /// The address of the native code to call inside of the VM process, if non-zero.
-    pub next_native_program_counter: AtomicU64,
+    pub next_native_program_counter: u64,
 
     /// The state of the program's heap.
     pub heap_info: VmCtxHeapInfo,
 
     pub arg2: AtomicU32,
 
     /// Offset in shared memory to this sandbox's memory map.
-    pub shm_memory_map_offset: AtomicU64,
+    pub shm_memory_map_offset: u64,
     /// Number of maps to map.
-    pub shm_memory_map_count: AtomicU64,
+    pub shm_memory_map_count: u64,
     /// Offset in shared memory to this sandbox's code.
-    pub shm_code_offset: AtomicU64,
+    pub shm_code_offset: u64,
     /// Length this sandbox's code.
-    pub shm_code_length: AtomicU64,
+    pub shm_code_length: u64,
     /// Offset in shared memory to this sandbox's jump table.
-    pub shm_jump_table_offset: AtomicU64,
+    pub shm_jump_table_offset: u64,
     /// Length of sandbox's jump table, in bytes.
-    pub shm_jump_table_length: AtomicU64,
+    pub shm_jump_table_length: u64,
 
     /// Address of the sysreturn routine.
-    pub sysreturn_address: AtomicU64,
+    pub sysreturn_address: u64,
 
     /// Whether userfaultfd-based memory management is enabled.
     pub uffd_enabled: AtomicBool,
@@ -328,7 +328,7 @@ pub const VMCTX_FUTEX_GUEST_SIGNAL: u32 = VMCTX_FUTEX_IDLE | (3 << 1);
 pub const VMCTX_FUTEX_GUEST_STEP: u32 = VMCTX_FUTEX_IDLE | (4 << 1);
 
 #[allow(clippy::declare_interior_mutable_const)]
-const ATOMIC_U64_ZERO: AtomicU64 = AtomicU64::new(0);
+const ATOMIC_U64_ZERO: u64 = 0;
 
 #[allow(clippy::new_without_default)]
 impl VmCtx {
@@ -338,25 +338,25 @@ impl VmCtx {
             _align_1: CacheAligned(()),
             _align_2: CacheAligned(()),
 
-            gas: AtomicI64::new(0),
+            gas: 0,
             program_counter: AtomicU32::new(0),
             next_program_counter: AtomicU32::new(0),
             arg: AtomicU32::new(0),
             arg2: AtomicU32::new(0),
             regs: [ATOMIC_U64_ZERO; REG_COUNT],
-            jump_into: AtomicU64::new(0),
-            next_native_program_counter: AtomicU64::new(0),
+            jump_into: 0,
+            next_native_program_counter: 0,
 
             futex: AtomicU32::new(VMCTX_FUTEX_BUSY),
 
-            shm_memory_map_offset: AtomicU64::new(0),
-            shm_memory_map_count: AtomicU64::new(0),
-            shm_code_offset: AtomicU64::new(0),
-            shm_code_length: AtomicU64::new(0),
-            shm_jump_table_offset: AtomicU64::new(0),
-            shm_jump_table_length: AtomicU64::new(0),
+            shm_memory_map_offset: 0,
+            shm_memory_map_count: 0,
+            shm_code_offset: 0,
+            shm_code_length: 0,
+            shm_jump_table_offset: 0,
+            shm_jump_table_length: 0,
             uffd_enabled: AtomicBool::new(false),
-            sysreturn_address: AtomicU64::new(0),
+            sysreturn_address: 0,
             heap_base: UnsafeCell::new(0),
             heap_initial_threshold: UnsafeCell::new(0),
             heap_max_size: UnsafeCell::new(0),
@@ -373,24 +373,24 @@ impl VmCtx {
             }),
 
             init: VmInit {
-                stack_address: AtomicU64::new(0),
-                stack_length: AtomicU64::new(0),
-                vdso_address: AtomicU64::new(0),
-                vdso_length: AtomicU64::new(0),
-                vvar_address: AtomicU64::new(0),
-                vvar_length: AtomicU64::new(0),
+                stack_address: 0,
+                stack_length: 0,
+                vdso_address: 0,
+                vdso_length: 0,
+                vvar_address: 0,
+                vvar_length: 0,
                 uffd_available: AtomicBool::new(false),
                 sandbox_disabled: AtomicBool::new(false),
                 logging_enabled: AtomicBool::new(false),
                 idle_regs: JmpBuf {
-                    rip: AtomicU64::new(0),
-                    rbx: AtomicU64::new(0),
-                    rsp: AtomicU64::new(0),
-                    rbp: AtomicU64::new(0),
-                    r12: AtomicU64::new(0),
-                    r13: AtomicU64::new(0),
-                    r14: AtomicU64::new(0),
-                    r15: AtomicU64::new(0),
+                    rip: 0,
+                    rbx: 0,
+                    rsp: 0,
+                    rbp: 0,
+                    r12: 0,
+                    r13: 0,
+                    r14: 0,
+                    r15: 0,
                 },
             },
 

crates/polkavm-zygote/src/main.rs

@@ -4,7 +4,7 @@
 
 use core::ptr::addr_of_mut;
 use core::sync::atomic::Ordering;
-use core::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize};
+use core::sync::atomic::{AtomicBool, AtomicUsize};
 
 #[rustfmt::skip]
 use polkavm_common::{
@@ -353,14 +353,14 @@ unsafe extern "C" fn signal_handler(signal: u32, _info: &linux_raw::siginfo_t, c
 }
 
 static mut RESUME_MAIN_LOOP_JMPBUF: JmpBuf = JmpBuf {
-    rip: AtomicU64::new(0),
-    rbx: AtomicU64::new(0),
-    rsp: AtomicU64::new(0),
-    rbp: AtomicU64::new(0),
-    r12: AtomicU64::new(0),
-    r13: AtomicU64::new(0),
-    r14: AtomicU64::new(0),
-    r15: AtomicU64::new(0),
+    rip: 0,
+    rbx: 0,
+    rsp: 0,
+    rbp: 0,
+    r12: 0,
+    r13: 0,
+    r14: 0,
+    r15: 0,
 };
 
 extern "C" {

crates/polkavm/Cargo.toml

@@ -15,6 +15,7 @@ log = { workspace = true }
 polkavm-assembler = { workspace = true, features = ["alloc"] }
 polkavm-common = { workspace = true, features = ["alloc", "logging", "regmap", "unique-id"] }
 schnellru = { workspace = true, optional = true }
+spin = { workspace = true }
 
 [target.'cfg(all(not(miri), target_arch = "x86_64", target_os = "linux"))'.dependencies]
 polkavm-linux-raw = { workspace = true, features = ["std"] }