The Rust Memory Model

use std::thread;

fn main() {
    let mut totals = vec![0u64; 4];
    let input: Vec<u64> = (1..=1000).collect();

    thread::scope(|s| {
        for (i, slot) in totals.iter_mut().enumerate() {
            let chunk = &input[i * 250..(i + 1) * 250]; // written before spawn: visible
            s.spawn(move || *slot = chunk.iter().sum());
        }
    }); // all threads joined here

    // Everything the threads wrote is visible after the join.
    println!("{totals:?} = {}", totals.iter().sum::<u64>());
}

use std::sync::Mutex;
use std::thread;

#[derive(Default, Debug)]
struct Inventory {
    apples: u32,
    pears: u32,
}

fn main() {
    let inv = Mutex::new(Inventory::default());
    thread::scope(|s| {
        for _ in 0..4 {
            s.spawn(|| {
                for _ in 0..1000 {
                    let mut g = inv.lock().unwrap();
                    g.apples += 1; // two plain writes, published together at unlock
                    g.pears += 2;
                } // unlock: Release
            });
        }
    });
    let g = inv.lock().unwrap(); // lock: Acquire
    println!("{g:?}, consistent: {}", g.pears == 2 * g.apples);
}

use std::collections::HashMap;
use std::sync::mpsc;
use std::thread;

fn main() {
    let (tx, rx) = mpsc::channel::<HashMap<String, u32>>();

    let producer = thread::spawn(move || {
        let mut index = HashMap::new(); // built with plain writes
        for word in "the cat and the hat and the bat".split(' ') {
            *index.entry(word.to_string()).or_insert(0) += 1;
        }
        tx.send(index).unwrap(); // publish: the whole map moves
    });

    let index = rx.recv().unwrap(); // everything written before `send` is visible
    let mut words: Vec<_> = index.into_iter().collect();
    words.sort();
    println!("{words:?}");
    producer.join().unwrap();
}

use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
use std::thread;

static PROGRESS: AtomicU64 = AtomicU64::new(0);

fn main() {
    thread::scope(|s| {
        s.spawn(|| {
            for i in 1..=100_000 {
                PROGRESS.store(i, Relaxed);
            }
        });
        s.spawn(|| {
            let mut last = 0;
            let mut went_backwards = false;
            while last < 100_000 {
                let now = PROGRESS.load(Relaxed);
                went_backwards |= now < last; // never true: coherence of one variable
                last = now;
            }
            println!("reader reached {last}, went backwards: {went_backwards}");
        });
    });
}

use std::sync::atomic::{AtomicBool, AtomicU64, Ordering::Relaxed};
use std::thread;

static DATA: AtomicU64 = AtomicU64::new(0);
static READY: AtomicBool = AtomicBool::new(false);

fn main() {
    thread::scope(|s| {
        s.spawn(|| {
            DATA.store(42, Relaxed);
            READY.store(true, Relaxed); // BUG: should be Release
        });
        s.spawn(|| {
            while !READY.load(Relaxed) {} // BUG: should be Acquire
            assert_eq!(DATA.load(Relaxed), 42); // may fail: the reader can see READY before DATA
        });
    });
}

use std::cell::UnsafeCell;
use std::ops::{Deref, DerefMut};
use std::sync::atomic::{AtomicBool, Ordering::{Acquire, Relaxed, Release}};
use std::thread;

pub struct SpinLock<T> {
    locked: AtomicBool,
    value: UnsafeCell<T>,
}

// Safety: the lock ensures only one thread accesses `value` at a time.
unsafe impl<T: Send> Sync for SpinLock<T> {}

pub struct Guard<'a, T> {
    lock: &'a SpinLock<T>,
}

impl<T> SpinLock<T> {
    pub const fn new(value: T) -> Self {
        Self { locked: AtomicBool::new(false), value: UnsafeCell::new(value) }
    }

    pub fn lock(&self) -> Guard<'_, T> {
        // Acquire: if we win, we see everything the previous holder wrote.
        while self.locked.compare_exchange_weak(false, true, Acquire, Relaxed).is_err() {
            std::hint::spin_loop();
        }
        Guard { lock: self }
    }
}

impl<T> Deref for Guard<'_, T> {
    type Target = T;
    fn deref(&self) -> &T {
        unsafe { &*self.lock.value.get() } // Safety: we hold the lock
    }
}

impl<T> DerefMut for Guard<'_, T> {
    fn deref_mut(&mut self) -> &mut T {
        unsafe { &mut *self.lock.value.get() } // Safety: we hold the lock
    }
}

impl<T> Drop for Guard<'_, T> {
    fn drop(&mut self) {
        self.lock.locked.store(false, Release); // Release: publish our writes
    }
}

fn main() {
    let log = SpinLock::new(Vec::new());
    thread::scope(|s| {
        for t in 0..4 {
            let log = &log;
            s.spawn(move || {
                for i in 0..1000 {
                    log.lock().push(t * 1000 + i);
                }
            });
        }
    });
    let entries = log.lock();
    println!("{} entries, none lost: {}", entries.len(), entries.iter().copied().collect::<std::collections::HashSet<_>>().len() == 4000);
}

use std::sync::atomic::{AtomicBool, Ordering::SeqCst};
use std::thread;

fn main() {
    let mut both_missed = 0;
    for _ in 0..20_000 {
        let (x, y) = (AtomicBool::new(false), AtomicBool::new(false));
        let (saw_y, saw_x) = thread::scope(|s| {
            let a = s.spawn(|| {
                x.store(true, SeqCst);
                y.load(SeqCst)
            });
            let b = s.spawn(|| {
                y.store(true, SeqCst);
                x.load(SeqCst)
            });
            (a.join().unwrap(), b.join().unwrap())
        });
        if !saw_y && !saw_x {
            both_missed += 1;
        }
    }
    println!("runs where both threads missed the other's flag: {both_missed}");
}

use std::sync::atomic::{AtomicBool, AtomicU64, Ordering::{Acquire, Relaxed, Release}, fence};
use std::thread;

static RESULTS: [AtomicU64; 4] = [const { AtomicU64::new(0) }; 4];
static DONE: [AtomicBool; 4] = [const { AtomicBool::new(false) }; 4];

fn main() {
    thread::scope(|s| {
        for w in 0..4 {
            s.spawn(move || {
                let sum: u64 = (w * 1000..(w + 1) * 1000).sum();
                RESULTS[w as usize].store(sum, Relaxed);
                DONE[w as usize].store(true, Release); // publish this worker's result
            });
        }

        // Poll all flags with Relaxed loads...
        while !DONE.iter().all(|d| d.load(Relaxed)) {
            std::hint::spin_loop();
        }
        // ...then one Acquire fence pairs with every Release store we observed.
        fence(Acquire);
        let total: u64 = RESULTS.iter().map(|r| r.load(Relaxed)).sum();
        println!("total: {total}");
    });
}

use std::collections::HashMap;
use std::sync::LazyLock;
use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
use std::thread;

static INIT_RUNS: AtomicUsize = AtomicUsize::new(0);

static COUNTRY_NAMES: LazyLock<HashMap<&'static str, &'static str>> = LazyLock::new(|| {
    INIT_RUNS.fetch_add(1, Relaxed);
    HashMap::from([("DE", "Germany"), ("FR", "France"), ("IN", "India")])
});

fn main() {
    let names: Vec<&str> = thread::scope(|s| {
        let handles: Vec<_> = ["DE", "FR", "IN", "DE"].into_iter().map(|code| s.spawn(move || COUNTRY_NAMES[code])).collect();
        handles.into_iter().map(|h| h.join().unwrap()).collect()
    });
    println!("{names:?}");
    println!("initializer ran {} time(s)", INIT_RUNS.load(Relaxed));
}

use std::mem::{align_of, size_of};
use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
use std::thread;

// Four counters in 32 bytes: usually on one cache line, contended by four threads.
struct Packed {
    counters: [AtomicU64; 4],
}

// Each counter on its own cache line.
#[repr(align(64))]
struct CachePadded(AtomicU64);

struct Padded {
    counters: [CachePadded; 4],
}

fn main() {
    println!("Packed: {} bytes; CachePadded: {} bytes, align {}", size_of::<Packed>(), size_of::<CachePadded>(), align_of::<CachePadded>());

    let packed = Packed { counters: [const { AtomicU64::new(0) }; 4] };
    let padded = Padded { counters: [const { CachePadded(AtomicU64::new(0)) }; 4] };

    thread::scope(|s| {
        for t in 0..4 {
            let (packed, padded) = (&packed, &padded);
            s.spawn(move || {
                for _ in 0..100_000 {
                    packed.counters[t].fetch_add(1, Relaxed); // fights over one line
                    padded.counters[t].0.fetch_add(1, Relaxed); // private line per thread
                }
            });
        }
    });
    let total = |c: &[u64]| c.iter().sum::<u64>();
    println!(
        "totals: {} {}",
        total(&packed.counters.iter().map(|c| c.load(Relaxed)).collect::<Vec<_>>()),
        total(&padded.counters.iter().map(|c| c.0.load(Relaxed)).collect::<Vec<_>>())
    );
}

#[cfg(loom)]
#[test]
fn publication_is_visible() {
    use loom::sync::Arc;
    use loom::sync::atomic::{AtomicBool, AtomicU64, Ordering::{Acquire, Relaxed, Release}};
    use loom::thread;

    loom::model(|| {
        let data = Arc::new(AtomicU64::new(0));
        let ready = Arc::new(AtomicBool::new(false));
        let (d, r) = (Arc::clone(&data), Arc::clone(&ready));
        let writer = thread::spawn(move || {
            d.store(42, Relaxed);
            r.store(true, Release);
        });
        if ready.load(Acquire) {
            assert_eq!(data.load(Relaxed), 42);
        }
        writer.join().unwrap();
    });
}

use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering::{Acquire, Relaxed, Release}};
use std::thread;

const CAP: usize = 8;

struct Spsc {
    slots: [AtomicU64; CAP],
    head: AtomicUsize, // next slot to read (written by the consumer only)
    tail: AtomicUsize, // next slot to write (written by the producer only)
}

impl Spsc {
    const fn new() -> Self {
        Self { slots: [const { AtomicU64::new(0) }; CAP], head: AtomicUsize::new(0), tail: AtomicUsize::new(0) }
    }

    fn try_push(&self, value: u64) -> bool {
        let tail = self.tail.load(Relaxed); // only we write tail
        if tail - self.head.load(Acquire) == CAP {
            return false; // full: wait until the consumer frees a slot
        }
        self.slots[tail % CAP].store(value, Relaxed);
        self.tail.store(tail + 1, Release); // publish the slot
        true
    }

    fn try_pop(&self) -> Option<u64> {
        let head = self.head.load(Relaxed); // only we write head
        if head == self.tail.load(Acquire) {
            return None; // empty
        }
        let value = self.slots[head % CAP].load(Relaxed); // visible thanks to Acquire on tail
        self.head.store(head + 1, Release); // hand the slot back to the producer
        Some(value)
    }
}

fn main() {
    let queue = Spsc::new();
    let (sum, count, in_order) = thread::scope(|s| {
        s.spawn(|| {
            for v in 1..=100_000u64 {
                while !queue.try_push(v) {
                    std::hint::spin_loop();
                }
            }
        });
        let consumer = s.spawn(|| {
            let (mut sum, mut count, mut last, mut in_order) = (0u64, 0u64, 0u64, true);
            while count < 100_000 {
                if let Some(v) = queue.try_pop() {
                    in_order &= v == last + 1;
                    last = v;
                    sum += v;
                    count += 1;
                }
            }
            (sum, count, in_order)
        });
        consumer.join().unwrap()
    });
    println!("received {count} values, sum {sum}, in order: {in_order}");
}