effect_rs/

stm.rs

use crate::core::{Cause, Effect, Exit};
use std::any::Any;
use std::collections::HashMap;
use std::sync::{Arc, RwLock};

type STMRun<E, A> = dyn Fn(&mut Journal) -> STMResult<E, A> + Send + Sync;

/// STM Effect
pub struct STM<E, A> {
    pub(crate) run: Arc<STMRun<E, A>>,
}

impl<E, A> Clone for STM<E, A> {
    fn clone(&self) -> Self {
        Self {
            run: self.run.clone(),
        }
    }
}

pub enum STMResult<E, A> {
    Success(A),
    Failure(E),
    Retry,
}

/// Abstract wrapper around a TRef state to allow locking in the Journal
#[allow(dead_code)]
trait AbstractEntry: Send + Sync {
    fn id(&self) -> u64;
    fn version(&self) -> u64;
    // We need to be able to clone the inner Arc<RwLock> but type erased?
    // Actually, we just need to be able to lock it.
    // The Journal needs to hold `dyn AbstractEntry`?
    // If we hold `Arc<dyn AbstractEntry>`, we can't lock easily because RwLock is generic.
    // Let's invert: TRef<A> implements Entry which holds Arc<RwLock<State<A>>>.
    // We need a non-generic trait to store in Journal list.

    // Lock for reading version
    fn read_version(&self) -> u64;
    // Lock for writing (update value and version)
    fn commit(&self, value: Box<dyn Any + Send + Sync>);

    fn notify(&self);
    fn listen(&self) -> Arc<tokio::sync::Notify>;
}

struct EntryImpl<A> {
    #[allow(dead_code)]
    id: u64,
    inner: Arc<RwLock<TRefState<A>>>,
}

impl<A> AbstractEntry for EntryImpl<A>
where
    A: Send + Sync + 'static,
{
    fn id(&self) -> u64 {
        self.id
    }
    fn version(&self) -> u64 {
        self.inner.read().unwrap().version
    }

    fn read_version(&self) -> u64 {
        self.inner.read().unwrap().version
    }

    fn commit(&self, value: Box<dyn Any + Send + Sync>) {
        if let Ok(typed) = value.downcast::<A>() {
            let mut guard = self.inner.write().unwrap();
            guard.value = *typed;
            guard.version += 1;
            guard.notify.notify_waiters();
        } else {
            panic!("STM Commit Type Mismatch");
        }
    }

    fn notify(&self) {
        self.inner.read().unwrap().notify.notify_waiters();
    }

    fn listen(&self) -> Arc<tokio::sync::Notify> {
        self.inner.read().unwrap().notify.clone()
    }
}

/// Transactional Reference
pub struct TRef<A> {
    id: u64,
    // Store type-erased entry for Journal interactions
    entry: Arc<EntryImpl<A>>,
}

impl<A> Clone for TRef<A> {
    fn clone(&self) -> Self {
        Self {
            id: self.id,
            entry: self.entry.clone(),
        }
    }
}

struct TRefState<A> {
    version: u64,
    value: A,
    notify: Arc<tokio::sync::Notify>,
}

type ReadEntry = (u64, Arc<dyn AbstractEntry>);
type WriteEntry = (Box<dyn Any + Send + Sync>, Arc<dyn AbstractEntry>);

/// Transaction Journal
pub struct Journal {
    // Reads: Map TRefId -> (Version, Entry)
    reads: HashMap<u64, ReadEntry>,
    // Writes: Map TRefId -> (NewValue, Entry)
    writes: HashMap<u64, WriteEntry>,
}

impl Default for Journal {
    fn default() -> Self {
        Self::new()
    }
}

impl Journal {
    pub fn new() -> Self {
        Self {
            reads: HashMap::new(),
            writes: HashMap::new(),
        }
    }
}

static TREF_ID_COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);

impl<A> TRef<A>
where
    A: Send + Sync + Clone + 'static,
{
    pub fn new(value: A) -> Self {
        let id = TREF_ID_COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        let notify = Arc::new(tokio::sync::Notify::new());
        let inner = Arc::new(RwLock::new(TRefState {
            version: 0,
            value,
            notify,
        }));
        Self {
            id,
            entry: Arc::new(EntryImpl { id, inner }),
        }
    }

    pub fn get(&self) -> STM<(), A> {
        let entry = self.entry.clone();
        let id = self.id;
        STM {
            run: Arc::new(move |journal| {
                if let Some((val, _)) = journal.writes.get(&id) {
                    if let Some(typed_val) = val.downcast_ref::<A>() {
                        return STMResult::Success(typed_val.clone());
                    } else {
                        panic!("TRef type mismatch");
                    }
                }

                // Read via Entry
                // We need to access value directly. Entry trait only exposes generic ops?
                // We shouldn't lock via trait if we have typed access here.
                // But we must assume 'entry' is sufficient.
                // Actually, `self.entry.inner` is accessible here because we are in `get` closure
                // which captures `self.entry` (EntryImpl<A>).
                let guard = entry.inner.read().unwrap();
                journal
                    .reads
                    .entry(id)
                    .or_insert((guard.version, entry.clone())); // Store generic entry
                STMResult::Success(guard.value.clone())
            }),
        }
    }

    pub fn set(&self, value: A) -> STM<(), ()> {
        let entry = self.entry.clone();
        let id = self.id;
        STM {
            run: Arc::new(move |journal| {
                journal
                    .writes
                    .insert(id, (Box::new(value.clone()), entry.clone()));
                STMResult::Success(())
            }),
        }
    }
}

// STM Combinators

impl<E, A> STM<E, A>
where
    E: 'static,
    A: 'static,
{
    pub fn succeed(a: A) -> STM<E, A>
    where
        A: Clone + Send + Sync,
    {
        STM {
            run: Arc::new(move |_| STMResult::Success(a.clone())),
        }
    }

    pub fn fail(e: E) -> STM<E, A>
    where
        E: Clone + Send + Sync,
    {
        STM {
            run: Arc::new(move |_| STMResult::Failure(e.clone())),
        }
    }

    pub fn retry() -> STM<E, A> {
        STM {
            run: Arc::new(|_| STMResult::Retry),
        }
    }

    pub fn map<B>(self, f: impl Fn(A) -> B + Send + Sync + 'static) -> STM<E, B>
    where
        A: Send + Sync,
        B: Send + Sync,
    {
        STM {
            run: Arc::new(move |journal| match (self.run)(journal) {
                STMResult::Success(a) => STMResult::Success(f(a)),
                STMResult::Failure(e) => STMResult::Failure(e),
                STMResult::Retry => STMResult::Retry,
            }),
        }
    }

    pub fn flat_map<B>(self, f: impl Fn(A) -> STM<E, B> + Send + Sync + 'static) -> STM<E, B>
    where
        A: Send + Sync,
        B: Send + Sync,
    {
        STM {
            run: Arc::new(move |journal| match (self.run)(journal) {
                STMResult::Success(a) => {
                    let next = f(a);
                    (next.run)(journal)
                }
                STMResult::Failure(e) => STMResult::Failure(e),
                STMResult::Retry => STMResult::Retry,
            }),
        }
    }

    pub fn map_error<E2>(self, f: impl Fn(E) -> E2 + Send + Sync + 'static) -> STM<E2, A>
    where
        E2: 'static,
    {
        STM {
            run: Arc::new(move |journal| match (self.run)(journal) {
                STMResult::Success(a) => STMResult::Success(a),
                STMResult::Failure(e) => STMResult::Failure(f(e)),
                STMResult::Retry => STMResult::Retry,
            }),
        }
    }
}

// Convert STM to Effect
impl<E, A> STM<E, A>
where
    E: Send + Sync + Clone + 'static,
    A: Send + Sync + Clone + 'static,
{
    pub fn commit<R>(self) -> Effect<R, E, A>
    where
        R: Send + Sync + 'static + Clone,
    {
        Effect {
            inner: Arc::new(move |_, _| {
                // R is ignored for now
                let stm = self.clone();
                Box::pin(async move {
                    loop {
                        let mut journal = Journal::new();
                        let result = (stm.run)(&mut journal);

                        match result {
                            STMResult::Success(val) => {
                                // 1. Gather all unique IDs involved (reads + writes)
                                // But for validation we only care about Reads.
                                // For locking, typical STM locks writes.
                                // Optimistic: Validate Reads, Lock Writes.

                                // Lock Implementation (Simplified Global ordered lock on TRefs?)
                                // We need to lock ALL modified TRefs to prevent write skews.
                                // Basic Strat: Lock everything involved in arbitrary order (risk deadlock)
                                // Or sort by ID to prevent deadlock.

                                let mut write_ids: Vec<u64> =
                                    journal.writes.keys().cloned().collect();
                                write_ids.sort(); // Lock in order

                                // We need to get the Entries.
                                // Actually, we can just grab entries from map.

                                // Lock Phase (Only abstract entries, we can't really "lock" generic traits easily without exposing Mutex)
                                // The AbstractEntry uses RwLock.
                                // We can't hold multiple guards easily in a generic list in implementation.
                                // This is tricky in Rust.

                                // Alternative: Global STM Lock (Mutex).
                                // Very slow but safe.

                                // Let's implement Optimistic Validation logic without holding locks across all operations if possible?
                                // No, commit must be atomic.

                                // Let's use a "Try Commit" approach.
                                // 1. Validate all Reads (check versions == current).
                                //    If any fail, RETRY loop.
                                // 2. If valid, try to Lock all Writes.
                                // 3. Re-validate Reads (because between 1 and 2, things changed).
                                // 4. Update.

                                // Since `AbstractEntry` hides `RwLock`, we implemented `commit` which locks internally.
                                // But we need to lock multiple at once.

                                // Hack: Global STM Mutex to serialize Commits.
                                // Reads can run concurrent. Commits are serialized.
                                // This is a common simplified STM strategy.
                                // Let's do that for now.

                                {
                                    // GLOBAL_COMMIT_LOCK
                                    let _guard = GLOBAL_STM_LOCK.lock().unwrap();

                                    // Validate Reads
                                    let mut valid = true;
                                    for (ver, entry) in journal.reads.values() {
                                        if entry.read_version() != *ver {
                                            valid = false;
                                            break;
                                        }
                                    }

                                    if !valid {
                                        continue; // Retry loop
                                    }

                                    // Commit Writes
                                    for (_, (val, entry)) in journal.writes {
                                        entry.commit(val);
                                    }
                                }

                                return Exit::Success(val);
                            }
                            STMResult::Failure(e) => return Exit::Failure(Cause::Fail(e)),
                            STMResult::Retry => {
                                // Collect notifications
                                let notifiers: Vec<Arc<tokio::sync::Notify>> = journal
                                    .reads
                                    .values()
                                    .map(|(_, entry)| entry.listen())
                                    .collect();

                                if notifiers.is_empty() {
                                    // Retry called but no reads? Deadlock/Infinite loop.
                                    // Just yield or fail.
                                    return Exit::Failure(Cause::Die(Arc::new(
                                        "STM Retry with no dependencies".to_string(),
                                    )));
                                }

                                // Wait for ANY notification
                                // Basic: Wait for first one.
                                // Optimized: Select all.
                                // Tokio Notify doesn't support select_all trivially on a Vec.
                                // We can spawn tasks or just wait on one?
                                // If we wait on just one, we might miss others? No, we just need *one* change to retry.
                                // But if we pick the wrong one (that doesn't change), we sleep forever?
                                // We need to wait until ANY of them fires.

                                // We need to wait until ANY of them fires.

                                // Pin the futures on the heap so they are Unpin for select_all
                                let futures: Vec<_> =
                                    notifiers.iter().map(|n| Box::pin(n.notified())).collect();

                                futures::future::select_all(futures).await;

                                continue;
                            }
                        }
                    }
                })
            }),
        }
    }
}

lazy_static::lazy_static! {
    static ref GLOBAL_STM_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
}