Last active 1784621034

Revision 01a71c789aa94eb4a15bddb7a4bbf94db96bc76b

Gemma4-C++-to-rust-test.md Raw

This is a lock-free Triple Buffer implementation. Its primary purpose is to allow one thread (the Producer) to write data and another thread (the Consumer) to read data without blocking each other, while ensuring the consumer always gets the most recent "complete" frame.

Here is the translation into Rust.

The Rust Implementation

To implement this in Rust, we need UnsafeCell because we are mutating elements of an array across different threads based on a manual state machine rather than using a Mutex. We also use repr(align(64)) to match the C++ alignas(64), which prevents "false sharing" (where the CPU caches multiple buffers in one line, causing performance degradation).

use std::cell::UnsafeCell;
use std::sync::atomic::{AtomicU8, Ordering};

#[repr(align(64))]
pub struct TripleBuffer<T> {
    // UnsafeCell allows interior mutability for lock-free access
    buffers: [UnsafeCell<T>; 3],
    state: AtomicU8,
}

impl<T: Copy> TripleBuffer<T> {
    const FREE: u8 = 0;
    const WRITING: u8 = 1;
    const READABLE: u8 = 2;
    const READING: u8 = 3;

    const MASKS: [u8; 3] = [0x03, 0x0C, 0x30];
    const SHIFTS: [u8; 3] = [0, 2, 4];

    pub fn new(initial_value: T) -> Self {
        Self {
            buffers: [
                UnsafeCell::new(initial_value),
                UnsafeCell::new(initial_value),
                UnsafeCell::new(initial_value),
            ],
            state: AtomicU8::new(0),
        }
    }

    pub fn produce(&self, data: T) {
        let mut old_state: u8;
        let mut idx: usize;

        // Step 1: Find a FREE slot and mark it as WRITING
        loop {
            old_state = self.state.load(Ordering::Relaxed);
            idx = None;
            for i in 0..3 {
                if (old_state & Self::MASKS[i]) == Self::FREE {
                    idx = Some(i);
                    break;
                }
            }

            let idx = match idx {
                Some(i) => i,
                None => continue, // Spin until a slot is free
            };

            let new_state = old_state | (Self::WRITING << Self::SHIFTS[idx]);
            if self.state.compare_exchange_weak(
                old_state, new_state, Ordering::AcqRel, Ordering::Relaxed
            ).is_ok() {
                break;
            }
        }

        // Step 2: Write data to the buffer (Safe because we hold the WRITING state)
        unsafe {
            *self.buffers[idx].get() = data;
        }

        // Step 3: Mark the slot as READABLE
        loop {
            old_state = self.state.load(Ordering::Relaxed);
            if (old_state & Self::MASKS[idx]) != Self::WRITING {
                break;
            }
            let new_state = (old_state & !Self::MASKS[idx]) | (Self::READABLE << Self::SHIFTS[idx]);
            if self.state.compare_exchange_weak(
                old_state, new_state, Ordering::AcqRel, Ordering::Relaxed
            ).is_ok() {
                break;
            }
        }
    }

    pub fn consume(&self, out: &mut T) -> bool {
        let mut old_state: u8;
        let mut idx: usize;

        // Step 1: Find a READABLE slot and mark it as READING
        loop {
            old_state = self.state.load(Ordering::Relaxed);
            idx = None;
            for i in 0..3 {
                if (old_state & Self::MASKS[i]) == Self::READABLE {
                    idx = Some(i);
                    break;
                }
            }

            let idx = match idx {
                Some(i) => i,
                None => return false, // No complete frame available
            };

            let new_state = old_state | (Self::READING << Self::SHIFTS[idx]);
            if self.state.compare_exchange_weak(
                old_state, new_state, Ordering::AcqRel, Ordering::Relaxed
            ).is_ok() {
                break;
            }
        }

        // Step 2: Read data from the buffer (Safe because we hold the READING state)
        unsafe {
            *out = *self.buffers[idx].get();
        }

        // Step 3: Mark the slot as FREE
        loop {
            old_state = self.state.load(Ordering::Relaxed);
            if (old_state & Self::MASKS[idx]) != Self::READING {
                break;
            }
            let new_state = old_state & !Self::MASKS[idx]; // Set to FREE (0)
            if self.state.compare_exchange_weak(
                old_state, new_state, Ordering::AcqRel, Ordering::Relaxed
            ).is_ok() {
                break;
            }
        }

        true
    }
}

// Required to allow the buffer to be shared between threads
unsafe impl<T: Sync> Sync for TripleBuffer<T> {}

Explanation of the Logic

1. The State Machine (The "Bit-Packing" Trick)

Instead of using three separate atomic variables, this code uses a single u8 as a bitfield. Each buffer is allocated 2 bits of space:

  • 00: Free
  • 01: Writing (Producer owns it)
  • 10: Readable (Ready for Consumer)
  • 11: Reading (Consumer owns it)

The MASKS and SHIFTS constants are used to isolate these 2-bit chunks. For example, Buffer 1 is checked by shifting the state 2 bits to the right and masking with 0x03.

2. The Workflow

  • Producer (produce):
    1. Scans the bitfield for any buffer marked as FREE.
    2. Uses CAS (Compare-And-Swap) via compare_exchange_weak to atomically change FREE $\rightarrow$ WRITING.
    3. Writes the data into that specific buffer index.
    4. Uses CAS to change WRITING $\rightarrow$ READABLE.
  • Consumer (consume):
    1. Scans for any buffer marked as READABLE.
    2. Uses CAS to atomically change READABLE $\rightarrow$ READING.
    3. Copies the data out of that buffer index.
    4. Uses CAS to change READING $\rightarrow$ FREE.

3. Rust-Specific Details

  • UnsafeCell<T>: In Rust, you cannot mutate a value inside a shared reference (&self) unless it is wrapped in a type that provides "interior mutability." Since we are handling synchronization ourselves via atomics, UnsafeCell is the most efficient way to tell the compiler: "I know what I'm doing; allow me to mutate this memory."
  • Ordering::AcqRel: This ensures that memory writes performed by the producer before the state change are visible to the consumer after they see the state change.
  • T: Copy: The C++ code uses buffers[idx] = data, which performs a copy. In Rust, I added the Copy trait bound for simplicity. If T is a large struct, you would use std::ptr::copy_nonoverlapping or wrap the data in Box.
  • Alignment: #[repr(align(64))] ensures the structure starts at a cache line boundary, reducing CPU pipeline stalls caused by multiple cores fighting over the same cache line.