cranpose_core/

state.rs

use std::{
    any::{Any, TypeId},
    cell::{Cell, RefCell},
    fmt,
    hash::Hash,
    marker::PhantomData,
    ops::Deref,
    rc::{Rc, Weak as RcWeak},
    sync::{Arc, Mutex, MutexGuard, PoisonError, Weak},
};

use smallvec::SmallVec;

use crate::{
    RecomposeScope, RecomposeScopeInner, RuntimeHandle, ScopeId, StateId,
    collections::map::{HashMap, HashSet},
    debug_trace::debug_record_scope_invalidation,
    runtime,
    snapshot_id_set::{SnapshotId, SnapshotIdSet},
    snapshot_pinning::lowest_pinned_snapshot,
    snapshot_v2::{
        AnySnapshot, GlobalSnapshot, advance_global_snapshot, allocate_record_id, current_snapshot,
    },
    with_current_composer_opt,
};

pub(crate) const PREEXISTING_SNAPSHOT_ID: SnapshotId = 1;

const INVALID_SNAPSHOT_ID: SnapshotId = 0;

const SNAPSHOT_ID_MAX: SnapshotId = usize::MAX;

#[derive(Clone, Copy, Eq, PartialEq, Hash, Debug, Default)]
pub struct ObjectId(pub(crate) usize);

impl ObjectId {
    pub(crate) fn new<T: ?Sized + 'static>(object: &Arc<T>) -> Self {
        Self(Arc::as_ptr(object) as *const () as usize)
    }

    #[inline]
    pub(crate) fn as_usize(self) -> usize {
        self.0
    }
}

pub struct StateRecord {
    snapshot_id: Cell<SnapshotId>,
    tombstone: Cell<bool>,
    next: Cell<Option<Rc<StateRecord>>>,
    value: RefCell<Option<Box<dyn RecordValue>>>,
}

pub(crate) trait RecordValue: Any {
    fn clone_boxed(&self) -> Box<dyn RecordValue>;
    fn assign_into(&self, target: &mut Option<Box<dyn RecordValue>>);
}

impl<T: Any + Clone> RecordValue for T {
    fn clone_boxed(&self) -> Box<dyn RecordValue> {
        Box::new(self.clone())
    }

    fn assign_into(&self, target: &mut Option<Box<dyn RecordValue>>) {
        match target
            .as_mut()
            .and_then(|value| (&mut **value as &mut dyn Any).downcast_mut::<T>())
        {
            Some(slot) => slot.clone_from(self),
            None => *target = Some(Box::new(self.clone())),
        }
    }
}

#[derive(Debug)]
struct StateReadFailure {
    state_id: ObjectId,
    snapshot_id: SnapshotId,
    fresh_snapshot_id: SnapshotId,
    fresh_invalid: SnapshotIdSet,
    record_chain: Vec<(SnapshotId, bool)>,
}

impl std::fmt::Display for StateReadFailure {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Reading a state that was created after the snapshot was taken or in a snapshot that has not yet been applied\n\
             state={:?}, snapshot_id={}, fresh_snapshot_id={}, fresh_invalid={:?}\n\
             record_chain={:?}",
            self.state_id,
            self.snapshot_id,
            self.fresh_snapshot_id,
            self.fresh_invalid,
            self.record_chain
        )
    }
}

#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub(crate) enum StateRecordValueError {
    Missing,
}

impl StateRecord {
    pub(crate) fn new<T: Any + Clone>(
        snapshot_id: SnapshotId,
        value: T,
        next: Option<Rc<StateRecord>>,
    ) -> Rc<Self> {
        Self::with_boxed(snapshot_id, Box::new(value), next)
    }

    fn with_boxed(
        snapshot_id: SnapshotId,
        value: Box<dyn RecordValue>,
        next: Option<Rc<StateRecord>>,
    ) -> Rc<Self> {
        Rc::new(Self {
            snapshot_id: Cell::new(snapshot_id),
            tombstone: Cell::new(false),
            next: Cell::new(next),
            value: RefCell::new(Some(value)),
        })
    }

    #[inline]
    pub(crate) fn snapshot_id(&self) -> SnapshotId {
        self.snapshot_id.get()
    }

    #[inline]
    pub(crate) fn set_snapshot_id(&self, id: SnapshotId) {
        self.snapshot_id.set(id);
    }

    #[inline]
    pub(crate) fn next(&self) -> Option<Rc<StateRecord>> {
        self.next.take().inspect(|record| {
            self.next.set(Some(Rc::clone(record)));
        })
    }

    #[inline]
    pub(crate) fn set_next(&self, next: Option<Rc<StateRecord>>) {
        self.next.set(next);
    }

    #[inline]
    pub(crate) fn is_tombstone(&self) -> bool {
        self.tombstone.get()
    }

    #[inline]
    pub(crate) fn set_tombstone(&self, tombstone: bool) {
        self.tombstone.set(tombstone);
    }

    /// Takes the value out, for the caller to drop once the operation that
    /// cleared it is done: its destructor may read or write state.
    #[must_use = "drop the cleared value once the record chain is consistent"]
    pub(crate) fn clear_value(&self) -> Option<Box<dyn RecordValue>> {
        self.value.borrow_mut().take()
    }

    /// Marks the record dead and retires its value, for an operation that
    /// holds a [`RetiredValuesScope`].
    pub(crate) fn tombstone_retiring_value(&self) {
        self.set_tombstone(true);
        if let Some(value) = self.clear_value() {
            retire_record_value(value);
        }
    }

    /// Stores `new_value` and returns the value it displaced, for the caller
    /// to drop once the write is published: its destructor may read or write
    /// state.
    #[must_use = "drop the displaced value once the write is published"]
    pub(crate) fn replace_value<T: Any + Clone>(&self, new_value: T) -> Option<T> {
        let mut value = self.value.borrow_mut();
        if let Some(slot) = value
            .as_mut()
            .and_then(|current| (&mut **current as &mut dyn Any).downcast_mut::<T>())
        {
            return Some(std::mem::replace(slot, new_value));
        }
        let placeholder = value.replace(Box::new(new_value));
        drop(value);
        drop(placeholder);
        None
    }

    fn cloned_value(&self) -> Option<Box<dyn RecordValue>> {
        self.value
            .borrow()
            .as_ref()
            .map(|value| (**value).clone_boxed())
    }

    fn cloned_value_of(&self, expected: TypeId) -> Option<Box<dyn RecordValue>> {
        self.value
            .borrow()
            .as_ref()
            .filter(|value| (&***value as &dyn Any).type_id() == expected)
            .map(|value| (**value).clone_boxed())
    }

    fn set_boxed_value(&self, value: Box<dyn RecordValue>) {
        let displaced = self.value.borrow_mut().replace(value);
        if let Some(displaced) = displaced {
            retire_record_value(displaced);
        }
    }

    pub(crate) fn with_value<T: Any, R>(&self, f: impl FnOnce(&T) -> R) -> R {
        self.try_with_value(f)
            .unwrap_or_else(|| panic!("StateRecord value missing or wrong type"))
    }

    pub(crate) fn try_with_value<T: Any, R>(&self, f: impl FnOnce(&T) -> R) -> Option<R> {
        let guard = self.value.borrow();
        let value = guard
            .as_ref()
            .and_then(|boxed| (&**boxed as &dyn Any).downcast_ref::<T>())?;
        Some(f(value))
    }

    #[cfg(test)]
    pub(crate) fn clear_for_reuse(&self) {
        drop(self.clear_value());
    }

    pub(crate) fn assign_value(&self, source: &StateRecord) -> Result<(), StateRecordValueError> {
        if std::ptr::eq(self, source) {
            return match self.value.borrow().is_some() {
                true => Ok(()),
                false => Err(StateRecordValueError::Missing),
            };
        }
        let source_value = source.value.borrow();
        let value = source_value
            .as_ref()
            .ok_or(StateRecordValueError::Missing)?;
        (**value).assign_into(&mut self.value.borrow_mut());
        Ok(())
    }
}

impl Drop for StateRecord {
    fn drop(&mut self) {
        let mut next = self.next.take();
        while let Some(node) = next {
            match Rc::try_unwrap(node) {
                Ok(record) => {
                    next = record.next.take();
                }
                Err(_) => {
                    break;
                }
            }
        }
    }
}

struct CurrentRecord {
    head: RefCell<Rc<StateRecord>>,
}

impl CurrentRecord {
    fn new(head: Rc<StateRecord>) -> Self {
        Self {
            head: RefCell::new(head),
        }
    }

    fn clone_head(&self) -> Rc<StateRecord> {
        self.head.borrow().clone()
    }

    fn replace(&self, new_head: Rc<StateRecord>) {
        *self.head.borrow_mut() = new_head;
    }

    fn prepend(&self, record: Rc<StateRecord>) {
        let current_head = self.clone_head();
        record.set_next(Some(current_head));
        self.replace(record);
    }
}

#[inline]
fn record_is_valid_for(
    record: &Rc<StateRecord>,
    snapshot_id: SnapshotId,
    invalid: &SnapshotIdSet,
) -> bool {
    if record.is_tombstone() {
        return false;
    }

    let candidate = record.snapshot_id();
    if candidate == INVALID_SNAPSHOT_ID || candidate > snapshot_id {
        return false;
    }

    candidate == snapshot_id || !invalid.get(candidate)
}

pub(crate) fn readable_record_for(
    head: &Rc<StateRecord>,
    snapshot_id: SnapshotId,
    invalid: &SnapshotIdSet,
) -> Option<Rc<StateRecord>> {
    let mut best: Option<Rc<StateRecord>> = None;
    let mut cursor = Some(Rc::clone(head));

    while let Some(record) = cursor {
        cursor = record.next();
        if record_is_valid_for(&record, snapshot_id, invalid)
            && best
                .as_ref()
                .is_none_or(|current| current.snapshot_id() < record.snapshot_id())
        {
            best = Some(record);
        }
    }

    best
}

pub(crate) fn used_locked(head: &Rc<StateRecord>) -> Option<Rc<StateRecord>> {
    let mut current = Some(Rc::clone(head));
    let mut valid_record: Option<Rc<StateRecord>> = None;

    let reuse_limit = lowest_pinned_snapshot().map_or_else(
        || allocate_record_id().saturating_sub(1),
        |lowest| lowest.saturating_sub(1),
    );

    let invalid = SnapshotIdSet::EMPTY;

    while let Some(record) = current {
        let current_id = record.snapshot_id();

        if current_id == PREEXISTING_SNAPSHOT_ID {
            current = record.next();
            continue;
        }

        if current_id == INVALID_SNAPSHOT_ID {
            return Some(record);
        }

        if record.is_tombstone() && current_id < reuse_limit {
            return Some(record);
        }

        if record_is_valid_for(&record, reuse_limit, &invalid) {
            if let Some(ref existing) = valid_record {
                return Some(if current_id < existing.snapshot_id() {
                    record
                } else {
                    Rc::clone(existing)
                });
            } else {
                valid_record = Some(record.clone());
            }
        }

        current = record.next();
    }

    None
}

pub(crate) fn new_overwritable_record_locked(state: &dyn StateObject) -> Rc<StateRecord> {
    let state_head = state.first_record();

    if let Some(reusable) = used_locked(&state_head) {
        reusable.set_snapshot_id(SNAPSHOT_ID_MAX);
        return reusable;
    }

    let new_record = StateRecord::new(SNAPSHOT_ID_MAX, (), None);

    state.prepend_state_record(Rc::clone(&new_record));

    new_record
}

pub(crate) fn new_overwritable_record_as_head_locked(state: &dyn StateObject) -> Rc<StateRecord> {
    let head = state.first_record();

    if let Some(reusable) = used_locked(&head) {
        reusable.set_snapshot_id(SNAPSHOT_ID_MAX);

        if !Rc::ptr_eq(&head, &reusable) {
            let mut cursor = Some(Rc::clone(&head));
            let mut unlinked = false;

            while let Some(node) = cursor {
                let next = node.next();
                if let Some(next_record) = next {
                    if Rc::ptr_eq(&next_record, &reusable) {
                        node.set_next(reusable.next());
                        unlinked = true;
                        break;
                    }
                    cursor = Some(next_record);
                } else {
                    break;
                }
            }

            if !unlinked {
                debug_assert!(
                    false,
                    "new_overwritable_record_as_head_locked: reusable record not found in chain"
                );
                let new_record = StateRecord::new(SNAPSHOT_ID_MAX, (), None);
                state.prepend_state_record(Rc::clone(&new_record));
                return new_record;
            }

            state.prepend_state_record(Rc::clone(&reusable));
        }

        return reusable;
    }

    let new_record = StateRecord::new(SNAPSHOT_ID_MAX, (), None);
    state.prepend_state_record(Rc::clone(&new_record));
    new_record
}

#[cold]
fn log_merge_value_missing(which: &str, state: ObjectId, record: &StateRecord) {
    log::error!(
        "SnapshotMutableState::merge_records {which} record value missing or wrong type (state {state:?}, {which}_id={})",
        record.snapshot_id()
    );
}

fn install_copy_as_head(
    state: &dyn StateObject,
    source: &StateRecord,
    expected: TypeId,
) -> Option<SnapshotId> {
    let value = source.cloned_value_of(expected)?;
    let new_id = allocate_record_id();
    let record = new_overwritable_record_as_head_locked(state);
    record.set_boxed_value(value);
    record.set_tombstone(false);
    record.set_snapshot_id(new_id);
    advance_global_snapshot(new_id);
    Some(new_id)
}

fn promote_record_locked(
    state: &dyn StateObject,
    child_id: SnapshotId,
    expected: TypeId,
) -> Result<(), &'static str> {
    let mut cursor = Some(state.first_record());
    while let Some(record) = cursor {
        if record.snapshot_id() == child_id {
            if install_copy_as_head(state, &record, expected).is_none() {
                log::error!(
                    "SnapshotMutableState::promote_record child record value missing or wrong type (state {:?}, child_id={})",
                    state.object_id(),
                    child_id
                );
                return Err("child record value missing or wrong type");
            }
            return Ok(());
        }
        cursor = record.next();
    }
    log::error!(
        "SnapshotMutableState::promote_record missing child record (state {:?}, child_id={})",
        state.object_id(),
        child_id
    );
    Err("missing child record")
}

fn commit_merged_record_locked(
    state: &dyn StateObject,
    merged: &StateRecord,
    expected: TypeId,
) -> Result<SnapshotId, &'static str> {
    install_copy_as_head(state, merged, expected).ok_or_else(|| {
        log::error!(
            "SnapshotMutableState::commit_merged_record merged record value missing or wrong type (state {:?}, merged_id={})",
            state.object_id(),
            merged.snapshot_id()
        );
        "merged record value missing or wrong type"
    })
}

std::thread_local! {
    // Record values displaced inside snapshot bookkeeping, and how many
    // operations that hold them back are running.
    static RETIRED_RECORD_VALUES: RefCell<Vec<Box<dyn RecordValue>>> =
        const { RefCell::new(Vec::new()) };
    static RETIRE_DEPTH: Cell<usize> = const { Cell::new(0) };
}

/// Drops `value`, a record value displaced inside snapshot bookkeeping, once
/// the outermost [`RetiredValuesScope`] ends, or now when none is open: its
/// destructor may read or write state.
pub(crate) fn retire_record_value(value: Box<dyn RecordValue>) {
    if RETIRE_DEPTH.with(Cell::get) == 0 {
        drop(value);
    } else {
        RETIRED_RECORD_VALUES.with(|retired| retired.borrow_mut().push(value));
    }
}

/// Holds back the drop of retired record values until the outermost scope
/// ends, so no destructor runs while snapshot bookkeeping is half done.
pub(crate) struct RetiredValuesScope(());

impl RetiredValuesScope {
    pub(crate) fn enter() -> Self {
        RETIRE_DEPTH.with(|depth| depth.set(depth.get() + 1));
        Self(())
    }
}

impl Drop for RetiredValuesScope {
    fn drop(&mut self) {
        let outermost = RETIRE_DEPTH.with(|depth| {
            depth.set(depth.get() - 1);
            depth.get() == 0
        });
        if outermost {
            let retired =
                RETIRED_RECORD_VALUES.with(|retired| std::mem::take(&mut *retired.borrow_mut()));
            drop(retired);
        }
    }
}

/// Marks the records of `state` that no open snapshot can read as free for
/// reuse, retiring their values, and reports whether the state still keeps
/// more than one record.
pub(crate) fn overwrite_unused_records_locked(state: &dyn StateObject) -> bool {
    let mut current = Some(state.first_record());
    let mut valid_record: Option<Rc<StateRecord>> = None;

    let reuse_limit =
        lowest_pinned_snapshot().unwrap_or_else(crate::snapshot_v2::peek_next_snapshot_id);

    let mut retained_records = 0;

    while let Some(record) = current {
        let current_id = record.snapshot_id();

        if current_id == INVALID_SNAPSHOT_ID {
        } else if current_id < reuse_limit {
            match valid_record.as_ref() {
                None => {
                    valid_record = Some(Rc::clone(&record));
                    retained_records += 1;
                }
                Some(valid) => {
                    let record_to_overwrite = if current_id < valid.snapshot_id() {
                        Rc::clone(&record)
                    } else {
                        let to_overwrite = Rc::clone(valid);
                        valid_record = Some(Rc::clone(&record));
                        to_overwrite
                    };
                    record_to_overwrite.set_snapshot_id(INVALID_SNAPSHOT_ID);
                    if let Some(value) = record_to_overwrite.clear_value() {
                        retire_record_value(value);
                    }
                }
            }
        } else {
            retained_records += 1;
        }

        current = record.next();
    }

    retained_records > 1
}

fn active_snapshot() -> AnySnapshot {
    current_snapshot().unwrap_or_else(|| AnySnapshot::Global(GlobalSnapshot::get_or_create()))
}

pub(crate) trait MutationPolicy<T>: Send + Sync {
    fn equivalent(&self, a: &T, b: &T) -> bool;
    fn merge(&self, _previous: &T, _current: &T, _applied: &T) -> Option<T> {
        None
    }
}

pub(crate) struct NeverEqual;

impl<T> MutationPolicy<T> for NeverEqual {
    fn equivalent(&self, _a: &T, _b: &T) -> bool {
        false
    }
}

pub(crate) struct StructuralEqual;

impl<T: PartialEq> MutationPolicy<T> for StructuralEqual {
    fn equivalent(&self, a: &T, b: &T) -> bool {
        a == b
    }
}

pub trait StateObject: Any {
    fn object_id(&self) -> ObjectId;
    fn first_record(&self) -> Rc<StateRecord>;
    fn try_readable_record(
        &self,
        snapshot_id: SnapshotId,
        invalid: &SnapshotIdSet,
    ) -> Option<Rc<StateRecord>>;
    fn readable_record(&self, snapshot_id: SnapshotId, invalid: &SnapshotIdSet) -> Rc<StateRecord>;

    fn prepend_state_record(&self, record: Rc<StateRecord>);

    fn merge_records(
        &self,
        _previous: Rc<StateRecord>,
        _current: Rc<StateRecord>,
        _applied: Rc<StateRecord>,
    ) -> Option<Rc<StateRecord>> {
        None
    }

    fn commit_merged_record(&self, _merged: Rc<StateRecord>) -> Result<SnapshotId, &'static str> {
        Err("StateObject does not support merged record commits")
    }
    fn promote_record(&self, child_id: SnapshotId) -> Result<(), &'static str>;

    fn overwrite_unused_records(&self) -> bool {
        false
    }

    fn observation_lease(&self) -> Option<Rc<dyn Any>> {
        None
    }

    fn as_any(&self) -> &dyn Any;
}

pub(crate) struct SnapshotMutableState<T> {
    head: CurrentRecord,
    policy: Arc<dyn MutationPolicy<T>>,
    id: ObjectId,
    weak_self: Mutex<Option<Weak<Self>>>,
    apply_observers: ApplyObservers,
    read_observation_lease: Rc<()>,
    scope_observation_count: Cell<usize>,
    subscriber_callbacks: RefCell<Vec<Rc<dyn Fn()>>>,
}

#[inline(always)]
fn assert_record_chain(
    head: &CurrentRecord,
    id: ObjectId,
    caller: &str,
    snapshot_context: Option<SnapshotId>,
) {
    if should_check_chain_integrity() {
        check_chain_integrity(&head.clone_head(), id, caller, snapshot_context);
    }
}

enum Merge<'a, T> {
    KeepCurrent,
    Value(Option<T>),
    Missing(&'static str, &'a Rc<StateRecord>),
}

type ApplyObservers = Mutex<Vec<Box<dyn Fn() + 'static>>>;

fn lock_apply_observers(
    observers: &ApplyObservers,
) -> MutexGuard<'_, Vec<Box<dyn Fn() + 'static>>> {
    observers.lock().unwrap_or_else(PoisonError::into_inner)
}

fn notify_applied(observers: &ApplyObservers) {
    for observer in lock_apply_observers(observers).iter() {
        observer();
    }
}

fn finish_apply(
    observers: &ApplyObservers,
    head: &CurrentRecord,
    id: ObjectId,
    caller: &str,
    snapshot_id: SnapshotId,
) {
    notify_applied(observers);
    assert_record_chain(head, id, caller, Some(snapshot_id));
}

// What a write does that does not depend on the state's value type, so it
// is compiled once instead of once per state type.
struct WriteTarget<'a> {
    head: &'a CurrentRecord,
    id: ObjectId,
    state: &'a dyn StateObject,
}

struct GlobalWrite {
    record: Rc<StateRecord>,
    new_id: SnapshotId,
    written: Option<Arc<dyn StateObject>>,
}

impl WriteTarget<'_> {
    fn assert_chain_integrity(&self, caller: &str, snapshot_context: Option<SnapshotId>) {
        assert_record_chain(self.head, self.id, caller, snapshot_context);
    }

    fn begin_global_write(
        &self,
        snapshot: &AnySnapshot,
        global: &GlobalSnapshot,
        snapshot_id: SnapshotId,
        written: Option<Arc<dyn StateObject>>,
    ) -> GlobalWrite {
        assert!(
            !global.has_pending_children(),
            "SnapshotMutableState::set attempted global write while pending children {:?} exist (state {:?}, snapshot_id={})",
            global.pending_children(),
            self.id,
            snapshot_id
        );
        if let Some(state) = &written {
            snapshot.record_write(Arc::clone(state));
        }
        mark_update_write(self.id);
        let new_id = allocate_record_id();
        let record = new_overwritable_record_as_head_locked(self.state);
        GlobalWrite {
            record,
            new_id,
            written,
        }
    }

    fn finish_global_write(
        &self,
        global: &GlobalSnapshot,
        write: GlobalWrite,
        snapshot_id: SnapshotId,
    ) {
        // Published only now that its value is stored: nothing may read the
        // record while the store runs.
        write.record.set_snapshot_id(write.new_id);
        write.record.set_tombstone(false);
        advance_global_snapshot(write.new_id);
        self.assert_chain_integrity("set(global-push)", Some(snapshot_id));

        let mut cleared: SmallVec<[Box<dyn RecordValue>; 2]> = SmallVec::new();
        if !global.has_pending_children() {
            let mut cursor = write.record.next();
            while let Some(node) = cursor {
                if !node.is_tombstone() && node.snapshot_id() != PREEXISTING_SNAPSHOT_ID {
                    cleared.extend(node.clear_value());
                    node.set_tombstone(true);
                }
                cursor = node.next();
            }
            self.assert_chain_integrity("set(global-tombstone)", Some(snapshot_id));
        }

        if let Some(modified) = write.written.as_ref() {
            crate::snapshot_v2::notify_apply_observers(
                std::slice::from_ref(modified),
                write.new_id,
            );
        }
        drop(cleared);
    }

    fn begin_child_write(
        &self,
        snapshot: &AnySnapshot,
        snapshot_id: SnapshotId,
        invalid: &SnapshotIdSet,
        written: Option<Arc<dyn StateObject>>,
    ) -> Rc<StateRecord> {
        if let Some(state) = written {
            snapshot.record_write(state);
        }
        mark_update_write(self.id);
        self.writable_record(snapshot_id, invalid)
    }

    fn writable_record(&self, snapshot_id: SnapshotId, invalid: &SnapshotIdSet) -> Rc<StateRecord> {
        let Some(readable) = readable_record_for(&self.head.clone_head(), snapshot_id, invalid)
        else {
            let current_head = self.head.clone_head();
            let refreshed = readable_record_for(&current_head, snapshot_id, invalid);
            let source = refreshed.unwrap_or_else(|| current_head.clone());

            let cloned_value = source
                .cloned_value()
                .unwrap_or_else(|| panic!("StateRecord value missing or wrong type"));
            let new_head = StateRecord::with_boxed(snapshot_id, cloned_value, Some(current_head));
            self.head.replace(new_head.clone());
            self.assert_chain_integrity("writable_record(recover)", Some(snapshot_id));
            return new_head;
        };

        if readable.snapshot_id() == snapshot_id {
            return readable;
        }

        let refreshed = {
            let current_head = self.head.clone_head();
            let refreshed = readable_record_for(&current_head, snapshot_id, invalid).unwrap_or_else(
                || {
                    panic!(
                        "SnapshotMutableState::writable_record failed to locate refreshed readable record (state {:?}, snapshot_id={}, invalid={:?})",
                        self.id, snapshot_id, invalid
                    )
                },
            );

            if refreshed.snapshot_id() == snapshot_id {
                return refreshed;
            }

            Rc::clone(&refreshed)
        };

        let overwritable = new_overwritable_record_locked(self.state);
        if let Err(error) = overwritable.assign_value(&refreshed) {
            log::error!(
                "snapshot writable record could not copy refreshed value for state {:?}: {:?}",
                self.id,
                error
            );
        }
        overwritable.set_snapshot_id(snapshot_id);
        overwritable.set_tombstone(false);

        self.assert_chain_integrity("writable_record(reuse)", Some(snapshot_id));

        overwritable
    }
}

fn check_chain_integrity(
    head: &Rc<StateRecord>,
    id: ObjectId,
    caller: &str,
    snapshot_context: Option<SnapshotId>,
) {
    let mut cursor = Some(Rc::clone(head));
    let mut seen: HashSet<usize> = HashSet::default();
    let mut ids = Vec::new();

    while let Some(record) = cursor {
        let addr = Rc::as_ptr(&record) as usize;
        assert!(
            seen.insert(addr),
            "SnapshotMutableState::{caller} detected duplicate/cycle at record {:p} for state {id:?} (snapshot_context={snapshot_context:?}, chain_ids={ids:?})",
            Rc::as_ptr(&record),
        );
        ids.push(record.snapshot_id());
        cursor = record.next();
    }

    assert!(
        !ids.is_empty(),
        "SnapshotMutableState::{caller} finished integrity scan with empty id list for state {id:?} (snapshot_context={snapshot_context:?})"
    );
}

fn should_check_chain_integrity() -> bool {
    #[cfg(debug_assertions)]
    {
        true
    }

    #[cfg(not(debug_assertions))]
    {
        crate::env_flag!("CRANPOSE_ASSERT_STATE_CHAIN")
    }
}

impl<T: Clone + 'static> SnapshotMutableState<T> {
    fn record_chain_debug(&self) -> Vec<(SnapshotId, bool)> {
        let mut chain_ids = Vec::new();
        let mut cursor = Some(self.first_record());
        while let Some(record) = cursor {
            chain_ids.push((record.snapshot_id(), record.is_tombstone()));
            cursor = record.next();
        }
        chain_ids
    }

    fn readable_record_for_active_snapshot(&self) -> Result<Rc<StateRecord>, StateReadFailure> {
        let snapshot = active_snapshot();
        snapshot.record_read(self);

        let snapshot_id = snapshot.snapshot_id();
        if let Some(record) =
            snapshot.with_invalid(|invalid| self.readable_for(snapshot_id, invalid))
        {
            return Ok(record);
        }

        let fresh_snapshot = active_snapshot();
        let fresh_id = fresh_snapshot.snapshot_id();
        let fresh_invalid = fresh_snapshot.invalid();

        if let Some(record) = self.readable_for(fresh_id, &fresh_invalid) {
            return Ok(record);
        }

        let global = GlobalSnapshot::get_or_create();
        let global_id = global.snapshot_id();
        let global_invalid = global.invalid();

        if let Some(record) = self.readable_for(global_id, &global_invalid) {
            return Ok(record);
        }

        Err(StateReadFailure {
            state_id: self.id,
            snapshot_id,
            fresh_snapshot_id: fresh_id,
            fresh_invalid,
            record_chain: self.record_chain_debug(),
        })
    }

    fn readable_for(
        &self,
        snapshot_id: SnapshotId,
        invalid: &SnapshotIdSet,
    ) -> Option<Rc<StateRecord>> {
        let head = self.first_record();
        readable_record_for(&head, snapshot_id, invalid)
    }

    fn is_equivalent_to(&self, readable: Option<Rc<StateRecord>>, new_value: &T) -> bool {
        readable.is_some_and(|record| {
            record.with_value(|current: &T| self.policy.equivalent(current, new_value))
        })
    }

    pub(crate) fn new_in_arc(initial: T, policy: Arc<dyn MutationPolicy<T>>) -> Arc<Self> {
        let snapshot = active_snapshot();
        let snapshot_id = snapshot.snapshot_id();

        let tail = StateRecord::new(PREEXISTING_SNAPSHOT_ID, initial.clone(), None);
        let head = StateRecord::new(snapshot_id, initial, Some(tail));

        let mut state = Arc::new(Self {
            head: CurrentRecord::new(head),
            policy,
            id: ObjectId::default(),
            weak_self: Mutex::new(None),
            apply_observers: Mutex::new(Vec::new()),
            read_observation_lease: Rc::new(()),
            scope_observation_count: Cell::new(0),
            subscriber_callbacks: RefCell::new(Vec::new()),
        });

        let id = ObjectId::new(&state);
        if let Some(state_inner) = Arc::get_mut(&mut state) {
            state_inner.id = id;
        }

        *state.lock_weak_self() = Some(Arc::downgrade(&state));

        state
    }

    pub(crate) fn add_apply_observer(&self, observer: Box<dyn Fn() + 'static>) {
        self.lock_apply_observers().push(observer);
    }

    fn acquire_observation_lease(&self) -> Option<Rc<dyn Any>> {
        let was_empty = !self.has_subscribers();
        let lease = Rc::clone(&self.read_observation_lease);
        if was_empty {
            notify_subscriber_callbacks(&self.subscriber_callbacks);
        }
        Some(lease)
    }

    fn add_scope_observer(&self) -> bool {
        let was_empty = !self.has_subscribers();
        let count = self
            .scope_observation_count
            .get()
            .checked_add(1)
            .expect("state scope observation count overflow");
        self.scope_observation_count.set(count);
        was_empty
    }

    fn remove_scope_observers(&self, count: usize) {
        if count == 0 {
            return;
        }
        let remaining = self
            .scope_observation_count
            .get()
            .checked_sub(count)
            .expect("state scope observation count underflow");
        self.scope_observation_count.set(remaining);
    }

    fn has_subscribers(&self) -> bool {
        Rc::strong_count(&self.read_observation_lease) > 1 || self.scope_observation_count.get() > 0
    }

    fn subscriber_callback(&self, callback: Rc<dyn Fn()>, notify: bool) {
        self.subscriber_callbacks
            .borrow_mut()
            .push(Rc::clone(&callback));
        if notify {
            callback();
        }
    }

    fn notify_subscribers(&self) {
        notify_subscriber_callbacks(&self.subscriber_callbacks);
    }

    fn lock_weak_self(&self) -> MutexGuard<'_, Option<Weak<Self>>> {
        self.weak_self
            .lock()
            .unwrap_or_else(PoisonError::into_inner)
    }

    fn lock_apply_observers(&self) -> MutexGuard<'_, Vec<Box<dyn Fn() + 'static>>> {
        lock_apply_observers(&self.apply_observers)
    }

    fn finish_apply(&self, caller: &str, snapshot_id: SnapshotId) {
        finish_apply(
            &self.apply_observers,
            &self.head,
            self.id,
            caller,
            snapshot_id,
        );
    }

    fn upgrade_self(&self) -> Option<Arc<Self>> {
        self.lock_weak_self().as_ref().and_then(Weak::upgrade)
    }

    fn written_state(&self) -> Option<Arc<dyn StateObject>> {
        self.upgrade_self()
            .map(|state| -> Arc<dyn StateObject> { state })
    }

    #[inline]
    pub(crate) fn id(&self) -> ObjectId {
        self.id
    }

    pub(crate) fn try_with_value<R>(&self, f: impl FnOnce(&T) -> R) -> Option<R> {
        let record = self.readable_record_for_active_snapshot().ok()?;
        record.try_with_value(f)
    }

    pub(crate) fn try_get(&self) -> Option<T> {
        self.try_with_value(Clone::clone)
    }

    pub(crate) fn with_value<R>(&self, f: impl FnOnce(&T) -> R) -> R {
        let record = self
            .readable_record_for_active_snapshot()
            .unwrap_or_else(|failure| panic!("{failure}"));
        record.with_value(f)
    }

    pub(crate) fn get(&self) -> T {
        self.with_value(Clone::clone)
    }

    pub(crate) fn set(&self, new_value: T) -> bool {
        #[cfg(debug_assertions)]
        {
            let in_handler = crate::in_event_handler();
            let in_snapshot = crate::in_applied_snapshot();
            if in_handler && !in_snapshot {
                log::warn!(
                    target: "cranpose::state",
                    "State modified in event handler without run_in_mutable_snapshot; \
                     this can make updates invisible to other contexts. Wrap the handler \
                     in run_in_mutable_snapshot() or dispatch_ui_event(). State: {:?}",
                    self.id
                );
            }
        }

        let snapshot = active_snapshot();
        let snapshot_id = snapshot.snapshot_id();
        let target = WriteTarget {
            head: &self.head,
            id: self.id,
            state: self,
        };

        match &snapshot {
            AnySnapshot::Global(global) => {
                let readable =
                    snapshot.with_invalid(|invalid| self.readable_for(snapshot_id, invalid));
                if self.is_equivalent_to(readable, &new_value) {
                    return false;
                }
                let write =
                    target.begin_global_write(&snapshot, global, snapshot_id, self.written_state());
                let displaced = write.record.replace_value(new_value);
                target.finish_global_write(global, write, snapshot_id);
                drop(displaced);
            }
            AnySnapshot::Mutable(_)
            | AnySnapshot::NestedMutable(_)
            | AnySnapshot::TransparentMutable(_) => {
                let invalid = snapshot.invalid();
                if self.is_equivalent_to(self.readable_for(snapshot_id, &invalid), &new_value) {
                    return false;
                }
                let record = target.begin_child_write(
                    &snapshot,
                    snapshot_id,
                    &invalid,
                    self.written_state(),
                );
                let displaced = record.replace_value(new_value);
                target.assert_chain_integrity("set(child-writable)", Some(snapshot_id));
                drop(displaced);
            }
            AnySnapshot::Readonly(_)
            | AnySnapshot::NestedReadonly(_)
            | AnySnapshot::TransparentReadonly(_) => {
                panic!("Cannot write to a read-only snapshot");
            }
        }

        true
    }
}

thread_local! {
    static ACTIVE_UPDATES: RefCell<HashSet<ObjectId>> = RefCell::new(HashSet::default());
    static PENDING_WRITES: RefCell<HashSet<ObjectId>> = RefCell::new(HashSet::default());
}

pub(crate) struct UpdateScope {
    id: ObjectId,
    finished: bool,
}

impl UpdateScope {
    pub(crate) fn new(id: ObjectId) -> Self {
        ACTIVE_UPDATES.with(|active| {
            active.borrow_mut().insert(id);
        });
        PENDING_WRITES.with(|pending| {
            pending.borrow_mut().remove(&id);
        });
        Self {
            id,
            finished: false,
        }
    }

    pub(crate) fn finish(mut self) -> bool {
        self.finished = true;
        ACTIVE_UPDATES.with(|active| {
            active.borrow_mut().remove(&self.id);
        });
        PENDING_WRITES.with(|pending| pending.borrow_mut().remove(&self.id))
    }
}

impl Drop for UpdateScope {
    fn drop(&mut self) {
        if self.finished {
            return;
        }
        ACTIVE_UPDATES.with(|active| {
            active.borrow_mut().remove(&self.id);
        });
        PENDING_WRITES.with(|pending| {
            pending.borrow_mut().remove(&self.id);
        });
    }
}

fn mark_update_write(id: ObjectId) {
    ACTIVE_UPDATES.with(|active| {
        if active.borrow().contains(&id) {
            PENDING_WRITES.with(|pending| {
                pending.borrow_mut().insert(id);
            });
        }
    });
}

impl<T: Clone + 'static> SnapshotMutableState<T> {
    fn try_readable_record(
        &self,
        snapshot_id: SnapshotId,
        invalid: &SnapshotIdSet,
    ) -> Option<Rc<StateRecord>> {
        self.readable_for(snapshot_id, invalid)
    }
}

impl<T: Clone + 'static> StateObject for SnapshotMutableState<T> {
    fn object_id(&self) -> ObjectId {
        self.id
    }

    fn first_record(&self) -> Rc<StateRecord> {
        self.head.clone_head()
    }

    fn try_readable_record(
        &self,
        snapshot_id: SnapshotId,
        invalid: &SnapshotIdSet,
    ) -> Option<Rc<StateRecord>> {
        self.try_readable_record(snapshot_id, invalid)
    }

    fn readable_record(&self, snapshot_id: SnapshotId, invalid: &SnapshotIdSet) -> Rc<StateRecord> {
        self.try_readable_record(snapshot_id, invalid)
            .unwrap_or_else(|| {
                panic!(
                    "SnapshotMutableState::readable_record returned null (state={:?}, snapshot_id={})",
                    self.id, snapshot_id
                )
            })
    }

    fn prepend_state_record(&self, record: Rc<StateRecord>) {
        self.head.prepend(record);
    }

    fn observation_lease(&self) -> Option<Rc<dyn Any>> {
        self.acquire_observation_lease()
    }

    fn merge_records(
        &self,
        previous: Rc<StateRecord>,
        current: Rc<StateRecord>,
        applied: Rc<StateRecord>,
    ) -> Option<Rc<StateRecord>> {
        let outcome = current
            .try_with_value(|current_value: &T| {
                applied
                    .try_with_value(|applied_value: &T| {
                        if self.policy.equivalent(current_value, applied_value) {
                            return Merge::KeepCurrent;
                        }
                        previous
                            .try_with_value(|previous_value: &T| {
                                Merge::Value(self.policy.merge(
                                    previous_value,
                                    current_value,
                                    applied_value,
                                ))
                            })
                            .unwrap_or(Merge::Missing("previous", &previous))
                    })
                    .unwrap_or(Merge::Missing("applied", &applied))
            })
            .unwrap_or(Merge::Missing("current", &current));
        match outcome {
            Merge::KeepCurrent => Some(current),
            Merge::Value(merged) => {
                merged.map(|value| StateRecord::new(applied.snapshot_id(), value, None))
            }
            Merge::Missing(which, record) => {
                log_merge_value_missing(which, self.id, record);
                None
            }
        }
    }

    fn promote_record(&self, child_id: SnapshotId) -> Result<(), &'static str> {
        promote_record_locked(self, child_id, TypeId::of::<T>())?;
        self.finish_apply("promote_record", child_id);
        Ok(())
    }

    fn commit_merged_record(&self, merged: Rc<StateRecord>) -> Result<SnapshotId, &'static str> {
        let new_id = commit_merged_record_locked(self, &merged, TypeId::of::<T>())?;
        self.finish_apply("commit_merged_record", new_id);
        Ok(new_id)
    }

    fn overwrite_unused_records(&self) -> bool {
        overwrite_unused_records_locked(self)
    }

    fn as_any(&self) -> &dyn Any {
        self
    }
}

pub(crate) struct MutableStateInner<T: Clone + 'static> {
    pub(crate) state: Arc<SnapshotMutableState<T>>,
    pub(crate) watchers: RefCell<HashMap<ScopeId, RcWeak<RecomposeScopeInner>>>,
    runtime: RuntimeHandle,
    state_id: Cell<Option<StateId>>,
}

fn notify_subscriber_callbacks(callbacks: &RefCell<Vec<Rc<dyn Fn()>>>) {
    let registered = callbacks.borrow().len();
    for index in 0..registered {
        let callback = callbacks.borrow().get(index).map(Rc::clone);
        if let Some(callback) = callback {
            callback();
        }
    }
}

fn shrink_watchers_if_sparse(watchers: &mut HashMap<ScopeId, RcWeak<RecomposeScopeInner>>) {
    let len = watchers.len();
    let capacity = watchers.capacity();
    if capacity > len.saturating_mul(4).max(32) {
        watchers.shrink_to_fit();
    }
}

impl<T: Clone + 'static> MutableStateInner<T> {
    pub(crate) fn new_with_policy(
        value: T,
        runtime: RuntimeHandle,
        policy: Arc<dyn MutationPolicy<T>>,
    ) -> Self {
        Self {
            state: SnapshotMutableState::new_in_arc(value, policy),
            watchers: RefCell::new(HashMap::default()),
            runtime,
            state_id: Cell::new(None),
        }
    }

    pub(crate) fn install_snapshot_observer(&self, state_id: StateId) {
        self.state_id.set(Some(state_id));
        let runtime_handle = self.runtime.clone();
        self.state.add_apply_observer(Box::new(move || {
            let runtime = runtime_handle.clone();
            runtime_handle.enqueue_ui_task(Box::new(move || {
                runtime.with_state_arena(|arena| {
                    let _ = arena.with_typed_opt::<T, _>(state_id, |inner| {
                        inner.invalidate_watchers();
                    });
                });
            }));
        }));
    }

    /// Adds `scope` to the state's watchers, saying whether it was new and
    /// whether the state thereby gained its first subscriber.
    ///
    /// Dead watchers still count as subscribers until pruned. They are
    /// pruned whenever no live one is left, so a reader arriving after every
    /// earlier one died still makes the state newly subscribed, and
    /// otherwise only when the map is full, before the insert would grow
    /// it: a state read by thousands of scopes, such as the density, must
    /// not scan them all for every new reader.
    fn register_scope(&self, scope: &RecomposeScope) -> (bool, bool) {
        let mut watchers = self.watchers.borrow_mut();
        let id = scope.id();
        match watchers.get(&id) {
            Some(existing) if existing.strong_count() > 0 => return (false, false),
            Some(_) => {
                watchers.remove(&id);
                self.state.remove_scope_observers(1);
            }
            None => {}
        }
        let any_live = watchers.values().any(|watcher| watcher.strong_count() > 0);
        if !any_live || watchers.len() >= watchers.capacity() {
            let before = watchers.len();
            watchers.retain(|_, watcher| watcher.strong_count() > 0);
            self.state.remove_scope_observers(before - watchers.len());
        }
        watchers.insert(id, scope.downgrade());
        drop(watchers);
        (true, self.state.add_scope_observer())
    }

    fn has_subscribers(&self) -> bool {
        let mut watchers = self.watchers.borrow_mut();
        let before = watchers.len();
        watchers.retain(|_, existing| existing.upgrade().is_some());
        self.state.remove_scope_observers(before - watchers.len());
        self.state.has_subscribers()
    }

    pub(crate) fn unregister_scope(&self, scope_id: ScopeId) {
        let mut watchers = self.watchers.borrow_mut();
        let removed = if watchers
            .get(&scope_id)
            .is_some_and(|weak| weak.upgrade().is_none())
        {
            watchers.remove(&scope_id);
            shrink_watchers_if_sparse(&mut watchers);
            true
        } else {
            false
        };
        drop(watchers);
        self.state.remove_scope_observers(usize::from(removed));
    }

    fn state_id(&self) -> Option<StateId> {
        self.state_id.get()
    }

    fn invalidate_watchers(&self) {
        let (watchers, removed_count): (Vec<RecomposeScope>, usize) = {
            let mut watchers = self.watchers.borrow_mut();
            let before = watchers.len();
            let mut live = Vec::with_capacity(watchers.len());
            watchers.retain(|_, scope| {
                if let Some(inner) = scope.upgrade() {
                    live.push(RecomposeScope { inner });
                    true
                } else {
                    false
                }
            });
            let removed_count = before - watchers.len();
            shrink_watchers_if_sparse(&mut watchers);
            (live, removed_count)
        };
        self.state.remove_scope_observers(removed_count);

        for watcher in watchers {
            debug_record_scope_invalidation(
                watcher.id(),
                self.state_id.get(),
                std::any::type_name::<T>(),
            );
            if let Some(state_id) = self.state_id.get() {
                watcher.invalidate_from_state(state_id);
            } else {
                watcher.invalidate();
            }
        }
    }
}

impl<T: Clone + 'static> Drop for MutableStateInner<T> {
    fn drop(&mut self) {
        self.state
            .remove_scope_observers(self.watchers.get_mut().len());
    }
}

fn register_current_state_scope<T: Clone + 'static>(inner: &MutableStateInner<T>) {
    let Some(Some(scope)) =
        with_current_composer_opt(super::composer::Composer::current_state_invalidation_scope)
    else {
        return;
    };
    let (registered, became_subscribed) = inner.register_scope(&scope);
    if registered {
        if let Some(state_id) = inner.state_id() {
            scope.record_state_subscription(state_id);
        }
        if became_subscribed {
            inner.state.notify_subscribers();
        }
    }
}

trait StateArenaHandle<T: Clone + 'static> {
    fn state_id(&self) -> StateId;
    fn runtime_id(&self) -> runtime::RuntimeId;

    fn runtime_handle(&self) -> RuntimeHandle {
        runtime::runtime_handle_by_id(self.runtime_id())
            .unwrap_or_else(|| panic!("runtime {:?} dropped", self.runtime_id()))
    }

    fn runtime_handle_opt(&self) -> Option<RuntimeHandle> {
        runtime::runtime_handle_by_id(self.runtime_id())
    }

    fn with_inner<R>(&self, f: impl FnOnce(&MutableStateInner<T>) -> R) -> R {
        runtime::with_state_arena_by_id(self.runtime_id(), |arena| {
            arena.with_typed::<T, R>(self.state_id(), f)
        })
        .unwrap_or_else(|| panic!("runtime {:?} dropped", self.runtime_id()))
    }

    fn try_with_inner<R>(&self, f: impl FnOnce(&MutableStateInner<T>) -> R) -> Option<R> {
        runtime::with_state_arena_by_id(self.runtime_id(), |arena| {
            arena.with_typed_opt::<T, R>(self.state_id(), f)
        })?
    }

    fn read_subscribed<R>(&self, read: impl FnOnce(&MutableStateInner<T>) -> R) -> R {
        self.with_inner(|inner| {
            let result = read(inner);
            register_current_state_scope(inner);
            result
        })
    }
}

/// Cheap copyable read-only view of a state cell.
pub struct State<T: Clone + 'static> {
    id: StateId,
    runtime_id: runtime::RuntimeId,
    _marker: PhantomData<fn() -> T>,
}

/// Cheap copyable mutable view of a state cell.
///
/// Ownership lives elsewhere: a composition slot, an [`OwnedMutableState`], or
/// the runtime for states created with [`crate::mutableStateOf`] /
/// [`MutableState::with_runtime`].
pub struct MutableState<T: Clone + 'static> {
    id: StateId,
    runtime_id: runtime::RuntimeId,
    _marker: PhantomData<fn() -> T>,
}

/// Owning state handle for reclaimable state cells.
#[derive(Clone)]
pub struct OwnedMutableState<T: Clone + 'static> {
    state: MutableState<T>,
    _lease: Rc<runtime::StateHandleLease>,
    _marker: PhantomData<fn() -> T>,
}

impl<T: Clone + 'static> PartialEq for State<T> {
    fn eq(&self, other: &Self) -> bool {
        self.state_id() == other.state_id() && self.runtime_id() == other.runtime_id()
    }
}

impl<T: Clone + 'static> Eq for State<T> {}

impl<T: Clone + 'static> PartialEq for MutableState<T> {
    fn eq(&self, other: &Self) -> bool {
        self.state_id() == other.state_id() && self.runtime_id() == other.runtime_id()
    }
}

impl<T: Clone + 'static> Eq for MutableState<T> {}

impl<T: Clone + 'static> Copy for State<T> {}

impl<T: Clone + 'static> Clone for State<T> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<T: Clone + 'static> Copy for MutableState<T> {}

impl<T: Clone + 'static> Clone for MutableState<T> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<T: Clone + 'static> StateArenaHandle<T> for State<T> {
    fn state_id(&self) -> StateId {
        self.id
    }

    fn runtime_id(&self) -> runtime::RuntimeId {
        self.runtime_id
    }
}

impl<T: Clone + 'static> State<T> {
    pub fn is_alive(&self) -> bool {
        self.try_with_inner(|_| ()).is_some()
    }

    pub fn try_with<R>(&self, f: impl FnOnce(&T) -> R) -> Option<R> {
        self.try_with_inner(|inner| inner.state.try_with_value(f))?
    }

    pub fn try_value(&self) -> Option<T> {
        self.try_with_inner(|inner| inner.state.try_get())?
    }

    /// Reads a copy of the value through `f` and subscribes the current
    /// scope. `f` may write this state; see [`Self::read`] for a read that
    /// borrows instead of copying.
    pub fn with<R>(&self, f: impl FnOnce(&T) -> R) -> R {
        let value = self.read_subscribed(|inner| inner.state.get());
        f(&value)
    }

    /// Reads the value in place through `f` and subscribes the current
    /// scope. Nothing is cloned, so a read of a large value costs nothing
    /// beyond `f`; in return `f` borrows the stored value and must not write
    /// this state.
    pub fn read<R>(&self, f: impl FnOnce(&T) -> R) -> R {
        self.read_subscribed(|inner| inner.state.with_value(f))
    }

    pub fn value(&self) -> T {
        self.read_subscribed(|inner| inner.state.get())
    }

    pub fn get(&self) -> T {
        self.value()
    }

    pub fn has_subscribers(&self) -> bool {
        self.with_inner(MutableStateInner::has_subscribers)
    }

    /// Runs `callback` each time this state gains its first subscriber, and
    /// right away when it already has one.
    ///
    /// The state owns `callback` for as long as the state lives, so nothing
    /// has to be kept alive on the caller's side.
    pub fn on_subscriber(&self, callback: impl Fn() + 'static) {
        let callback: Rc<dyn Fn()> = Rc::new(callback);
        self.with_inner(|inner| {
            inner
                .state
                .subscriber_callback(callback, inner.has_subscribers());
        });
    }

    /// Counts this state as subscribed for as long as the returned guard
    /// lives, without observing it from any scope.
    ///
    /// Snapshot observers subscribe reads made under composition or the
    /// draw phase, and producers gated on [`State::has_subscribers`] — an
    /// infinite transition, for one — stop themselves when the last such
    /// subscriber leaves. A consumer that reads the value from a polling
    /// context (a frame effect stepping it into a derived state, a
    /// background task) is invisible to that accounting and would starve
    /// the producer it depends on. Holding this guard keeps the producer
    /// alive; dropping it releases the count.
    pub fn subscription_hold(&self) -> StateSubscriptionHold {
        StateSubscriptionHold {
            _lease: self
                .try_with_inner(|inner| inner.state.observation_lease())
                .flatten(),
        }
    }
}

/// Keeps a [`State`] counted as subscribed while alive; see
/// [`State::subscription_hold`].
pub struct StateSubscriptionHold {
    _lease: Option<Rc<dyn Any>>,
}

impl<T: Clone + 'static> StateArenaHandle<T> for MutableState<T> {
    fn state_id(&self) -> StateId {
        self.id
    }

    fn runtime_id(&self) -> runtime::RuntimeId {
        self.runtime_id
    }
}

impl<T: Clone + 'static> MutableState<T> {
    pub fn with_runtime(value: T, runtime: RuntimeHandle) -> Self {
        runtime.alloc_persistent_state(value)
    }

    fn from_parts(id: StateId, runtime_id: runtime::RuntimeId) -> Self {
        Self {
            id,
            runtime_id,
            _marker: PhantomData,
        }
    }

    pub(crate) fn from_lease(lease: &Rc<runtime::StateHandleLease>) -> Self {
        Self::from_parts(lease.id(), lease.runtime().id())
    }

    pub fn is_alive(&self) -> bool {
        self.try_with_inner(|_| ()).is_some()
    }

    pub fn try_with<R>(&self, f: impl FnOnce(&T) -> R) -> Option<R> {
        self.try_with_inner(|inner| inner.state.try_with_value(f))?
    }

    pub fn try_value(&self) -> Option<T> {
        self.try_with_inner(|inner| inner.state.try_get())?
    }

    pub fn as_state(&self) -> State<T> {
        State {
            id: self.id,
            runtime_id: self.runtime_id,
            _marker: PhantomData,
        }
    }

    pub fn try_retain(&self) -> Option<OwnedMutableState<T>> {
        let lease = self
            .runtime_handle_opt()?
            .retain_state_lease(self.state_id())?;
        Some(OwnedMutableState {
            state: *self,
            _lease: lease,
            _marker: PhantomData,
        })
    }

    pub fn retain(&self) -> OwnedMutableState<T> {
        self.try_retain()
            .unwrap_or_else(|| panic!("state {:?} is no longer alive", self.state_id()))
    }

    /// Reads a copy of the value through `f` and subscribes the current
    /// scope. `f` may write this state; see [`Self::read`] for a read that
    /// borrows instead of copying.
    pub fn with<R>(&self, f: impl FnOnce(&T) -> R) -> R {
        let value = self.read_subscribed(|inner| inner.state.get());
        f(&value)
    }

    /// Reads the value in place through `f` and subscribes the current
    /// scope. Nothing is cloned, so a read of a large value costs nothing
    /// beyond `f`; in return `f` borrows the stored value and must not write
    /// this state.
    pub fn read<R>(&self, f: impl FnOnce(&T) -> R) -> R {
        self.read_subscribed(|inner| inner.state.with_value(f))
    }

    pub fn update<R>(&self, f: impl FnOnce(&mut T) -> R) -> R {
        let runtime = self.runtime_handle();
        runtime.assert_ui_thread();
        runtime.with_state_arena(|arena| {
            arena.with_typed::<T, R>(self.state_id(), |inner| {
                let mut value = inner.state.get();
                let tracker = UpdateScope::new(inner.state.id());
                let result = f(&mut value);
                let wrote_elsewhere = tracker.finish();
                if !wrote_elsewhere && inner.state.set(value) {
                    inner.invalidate_watchers();
                }
                result
            })
        })
    }

    pub fn replace(&self, value: T) {
        let Some(runtime) = self.runtime_handle_opt() else {
            log::debug!(
                "MutableState::replace skipped: runtime {:?} dropped",
                self.runtime_id()
            );
            return;
        };
        runtime.assert_ui_thread();
        let replaced = runtime
            .try_with_state_arena(|arena| {
                arena.with_typed_opt::<T, ()>(self.state_id(), |inner| {
                    if inner.state.set(value) {
                        inner.invalidate_watchers();
                    }
                })
            })
            .flatten();
        if replaced.is_none() {
            log::debug!(
                "MutableState::replace skipped: state cell released (slot={}, gen={})",
                self.state_id().slot(),
                self.state_id().generation(),
            );
        }
    }

    pub fn set_value(&self, value: T) {
        self.replace(value);
    }

    pub fn set(&self, value: T) {
        self.replace(value);
    }

    pub fn value(&self) -> T {
        self.read_subscribed(|inner| inner.state.get())
    }

    pub fn get(&self) -> T {
        self.value()
    }

    pub fn get_non_reactive(&self) -> T {
        self.with_inner(|inner| inner.state.get())
    }

    #[doc(hidden)]
    pub fn runtime_state_id(&self) -> StateId {
        self.state_id()
    }

    #[doc(hidden)]
    pub fn subscribe_current_scope_only(&self) {
        self.subscribe_current_scope();
    }

    fn subscribe_current_scope(&self) {
        self.with_inner(register_current_state_scope::<T>);
    }

    #[cfg(test)]
    pub(crate) fn watcher_count(&self) -> usize {
        self.with_inner(|inner| inner.watchers.borrow().len())
    }

    #[cfg(test)]
    pub(crate) fn watcher_capacity(&self) -> usize {
        self.with_inner(|inner| inner.watchers.borrow().capacity())
    }

    #[cfg(test)]
    pub(crate) fn state_id_for_test(&self) -> StateId {
        self.state_id()
    }

    #[cfg(test)]
    pub(crate) fn subscribe_scope_for_test(&self, scope: &RecomposeScope) {
        self.as_state().subscribe_scope_for_test(scope);
    }
}

impl<T: Clone + 'static> OwnedMutableState<T> {
    pub fn with_runtime(value: T, runtime: RuntimeHandle) -> Self {
        let lease = runtime.alloc_state(value);
        Self {
            state: MutableState::from_lease(&lease),
            _lease: lease,
            _marker: PhantomData,
        }
    }

    pub fn with_runtime_structural_eq(value: T, runtime: RuntimeHandle) -> Self
    where
        T: PartialEq,
    {
        Self::with_runtime_and_policy(value, runtime, Arc::new(StructuralEqual))
    }

    pub(crate) fn with_runtime_and_policy(
        value: T,
        runtime: RuntimeHandle,
        policy: Arc<dyn MutationPolicy<T>>,
    ) -> Self {
        let lease = runtime.alloc_state_with_policy(value, policy);
        Self {
            state: MutableState::from_lease(&lease),
            _lease: lease,
            _marker: PhantomData,
        }
    }

    pub fn handle(&self) -> MutableState<T> {
        self.state
    }

    pub fn as_state(&self) -> State<T> {
        self.state.as_state()
    }
}

impl<T: Clone + 'static> Deref for OwnedMutableState<T> {
    type Target = MutableState<T>;

    fn deref(&self) -> &Self::Target {
        &self.state
    }
}

#[cfg(test)]
impl<T: Clone + 'static> State<T> {
    pub(crate) fn subscribe_scope_for_test(&self, scope: &RecomposeScope) {
        self.with_inner(|inner| {
            let (registered, became_subscribed) = inner.register_scope(scope);
            if registered {
                if let Some(state_id) = inner.state_id() {
                    scope.record_state_subscription(state_id);
                }
                if became_subscribed {
                    inner.state.notify_subscribers();
                }
            }
        });
    }
}

impl<T: fmt::Debug + Clone + 'static> fmt::Debug for MutableState<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if let Some(value) = self.try_value() {
            f.debug_struct("MutableState")
                .field("value", &value)
                .finish()
        } else {
            f.write_str("MutableState { value: <unavailable> }")
        }
    }
}

#[derive(Clone)]
pub struct SnapshotStateList<T: Clone + 'static> {
    state: OwnedMutableState<Vec<T>>,
}

impl<T: Clone + 'static> SnapshotStateList<T> {
    pub fn with_runtime<I>(values: I, runtime: RuntimeHandle) -> Self
    where
        I: IntoIterator<Item = T>,
    {
        let initial: Vec<T> = values.into_iter().collect();
        Self {
            state: OwnedMutableState::with_runtime(initial, runtime),
        }
    }

    pub fn as_state(&self) -> State<Vec<T>> {
        self.state.as_state()
    }

    pub fn as_mutable_state(&self) -> MutableState<Vec<T>> {
        self.state.handle()
    }

    pub fn len(&self) -> usize {
        self.state.with(Vec::len)
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn to_vec(&self) -> Vec<T> {
        self.state.with(Clone::clone)
    }

    pub fn iter(&self) -> Vec<T> {
        self.to_vec()
    }

    pub fn get(&self, index: usize) -> T {
        self.state.with(|values| values[index].clone())
    }

    pub fn get_opt(&self, index: usize) -> Option<T> {
        self.state.with(|values| values.get(index).cloned())
    }

    pub fn first(&self) -> Option<T> {
        self.get_opt(0)
    }

    pub fn last(&self) -> Option<T> {
        self.state.with(|values| values.last().cloned())
    }

    pub fn push(&self, value: T) {
        self.state.update(|values| values.push(value));
    }

    pub fn extend<I>(&self, iter: I)
    where
        I: IntoIterator<Item = T>,
    {
        self.state.update(|values| values.extend(iter));
    }

    pub fn insert(&self, index: usize, value: T) {
        self.state.update(|values| values.insert(index, value));
    }

    pub fn set(&self, index: usize, value: T) -> T {
        self.state
            .update(|values| std::mem::replace(&mut values[index], value))
    }

    pub fn remove(&self, index: usize) -> T {
        self.state.update(|values| values.remove(index))
    }

    pub fn pop(&self) -> Option<T> {
        self.state.update(Vec::pop)
    }

    pub fn clear(&self) {
        self.state.replace(Vec::new());
    }

    pub fn retain<F>(&self, mut predicate: F)
    where
        F: FnMut(&T) -> bool,
    {
        self.state
            .update(|values| values.retain(|value| predicate(value)));
    }

    pub fn replace_with<I>(&self, iter: I)
    where
        I: IntoIterator<Item = T>,
    {
        self.state.replace(iter.into_iter().collect());
    }
}

impl<T: fmt::Debug + Clone + 'static> fmt::Debug for SnapshotStateList<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let contents = self.to_vec();
        f.debug_struct("SnapshotStateList")
            .field("values", &contents)
            .finish()
    }
}

#[derive(Clone)]
pub struct SnapshotStateMap<K, V>
where
    K: Clone + Eq + Hash + 'static,
    V: Clone + 'static,
{
    state: OwnedMutableState<HashMap<K, V>>,
}

impl<K, V> SnapshotStateMap<K, V>
where
    K: Clone + Eq + Hash + 'static,
    V: Clone + 'static,
{
    pub fn with_runtime<I>(pairs: I, runtime: RuntimeHandle) -> Self
    where
        I: IntoIterator<Item = (K, V)>,
    {
        let map: HashMap<K, V> = pairs.into_iter().collect();
        Self {
            state: OwnedMutableState::with_runtime(map, runtime),
        }
    }

    pub fn as_state(&self) -> State<HashMap<K, V>> {
        self.state.as_state()
    }

    pub fn as_mutable_state(&self) -> MutableState<HashMap<K, V>> {
        self.state.handle()
    }

    pub fn len(&self) -> usize {
        self.state.with(std::collections::HashMap::len)
    }

    pub fn is_empty(&self) -> bool {
        self.state.with(std::collections::HashMap::is_empty)
    }

    pub fn contains_key(&self, key: &K) -> bool {
        self.state.with(|map| map.contains_key(key))
    }

    pub fn get(&self, key: &K) -> Option<V> {
        self.state.with(|map| map.get(key).cloned())
    }

    pub fn to_hash_map(&self) -> HashMap<K, V> {
        self.state.with(Clone::clone)
    }

    pub fn insert(&self, key: K, value: V) -> Option<V> {
        self.state.update(|map| map.insert(key, value))
    }

    pub fn extend<I>(&self, iter: I)
    where
        I: IntoIterator<Item = (K, V)>,
    {
        self.state.update(|map| map.extend(iter));
    }

    pub fn remove(&self, key: &K) -> Option<V> {
        self.state.update(|map| map.remove(key))
    }

    pub fn clear(&self) {
        self.state.replace(HashMap::default());
    }

    pub fn retain<F>(&self, mut predicate: F)
    where
        F: FnMut(&K, &mut V) -> bool,
    {
        self.state.update(|map| map.retain(|k, v| predicate(k, v)));
    }
}

impl<K, V> fmt::Debug for SnapshotStateMap<K, V>
where
    K: Clone + Eq + Hash + fmt::Debug + 'static,
    V: Clone + fmt::Debug + 'static,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let contents = self.to_hash_map();
        f.debug_struct("SnapshotStateMap")
            .field("entries", &contents)
            .finish()
    }
}

pub(crate) struct DerivedState<T: Clone + PartialEq + 'static> {
    compute: Rc<dyn Fn() -> T>,
    pub(crate) state: OwnedMutableState<T>,
}

impl<T: Clone + PartialEq + 'static> DerivedState<T> {
    pub(crate) fn new(runtime: RuntimeHandle, compute: Rc<dyn Fn() -> T>) -> Self {
        let initial = compute();
        Self {
            compute,
            state: OwnedMutableState::with_runtime_structural_eq(initial, runtime),
        }
    }

    pub(crate) fn set_compute(&mut self, compute: Rc<dyn Fn() -> T>) {
        self.compute = compute;
    }

    pub(crate) fn recompute(&self) {
        let value = (self.compute)();
        self.state.set_value(value);
    }
}

impl<T: fmt::Debug + Clone + 'static> fmt::Debug for State<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if let Some(value) = self.try_value() {
            f.debug_struct("State").field("value", &value).finish()
        } else {
            f.write_str("State { value: <unavailable> }")
        }
    }
}

#[cfg(test)]
#[path = "tests/state_tests.rs"]
mod tests;