Skip to main content

salmon_core/
atomic.rs

1//! Lock-free `f64` accumulation, used for concurrent mass/count updates.
2
3use crate::math::log_add;
4use std::sync::atomic::{AtomicU64, Ordering};
5
6/// An `f64` stored as its bit pattern in an `AtomicU64`, supporting lock-free
7/// updates via a CAS loop. Salmon updates per-transcript mass and the fragment
8/// length distribution concurrently from worker threads; this is the building
9/// block for that.
10#[derive(Debug)]
11pub struct AtomicF64(AtomicU64);
12
13impl AtomicF64 {
14    pub fn new(v: f64) -> Self {
15        Self(AtomicU64::new(v.to_bits()))
16    }
17
18    pub fn load(&self) -> f64 {
19        f64::from_bits(self.0.load(Ordering::Relaxed))
20    }
21
22    pub fn store(&self, v: f64) {
23        self.0.store(v.to_bits(), Ordering::Relaxed);
24    }
25
26    /// Atomically replace the stored value with `f(current)`, returning the new value.
27    pub fn fetch_update<F: Fn(f64) -> f64>(&self, f: F) -> f64 {
28        let mut cur = self.0.load(Ordering::Relaxed);
29        loop {
30            let next = f(f64::from_bits(cur)).to_bits();
31            match self
32                .0
33                .compare_exchange_weak(cur, next, Ordering::Relaxed, Ordering::Relaxed)
34            {
35                Ok(_) => return f64::from_bits(next),
36                Err(actual) => cur = actual,
37            }
38        }
39    }
40
41    /// Atomically accumulate `log_inc` into the stored log-space value:
42    /// `self = log(exp(self) + exp(log_inc))`. Mirrors salmon's `incLoopLog`.
43    pub fn log_add_assign(&self, log_inc: f64) {
44        self.fetch_update(|cur| log_add(cur, log_inc));
45    }
46
47    /// Atomically add `inc` in linear space. Mirrors salmon's `incLoop`.
48    pub fn add_assign(&self, inc: f64) {
49        self.fetch_update(|cur| cur + inc);
50    }
51}
52
53impl Default for AtomicF64 {
54    fn default() -> Self {
55        Self::new(0.0)
56    }
57}