1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

//! Integration with the generational-box crate for copy state management.
//!
//! Each scope in dioxus has a single [Owner]

use std::{
    any::{Any, TypeId},
    cell::RefCell,
};

use generational_box::{AnyStorage, Owner, SyncStorage, UnsyncStorage};

use crate::{innerlude::current_scope_id, ScopeId};

/// Run a closure with the given owner.
///
/// This will override the default owner for the current component.
pub fn with_owner<S: AnyStorage, F: FnOnce() -> R, R>(owner: Owner<S>, f: F) -> R {
    let old_owner = set_owner(Some(owner));
    let result = f();
    set_owner(old_owner);
    result
}

/// Set the owner for the current thread.
fn set_owner<S: AnyStorage>(owner: Option<Owner<S>>) -> Option<Owner<S>> {
    let id = TypeId::of::<S>();
    if id == TypeId::of::<SyncStorage>() {
        SYNC_OWNER.with(|cell| {
            std::mem::replace(
                &mut *cell.borrow_mut(),
                owner.map(|owner| {
                    *(Box::new(owner) as Box<dyn Any>)
                        .downcast::<Owner<SyncStorage>>()
                        .unwrap()
                }),
            )
            .map(|owner| *(Box::new(owner) as Box<dyn Any>).downcast().unwrap())
        })
    } else {
        UNSYNC_OWNER.with(|cell| {
            std::mem::replace(
                &mut *cell.borrow_mut(),
                owner.map(|owner| {
                    *(Box::new(owner) as Box<dyn Any>)
                        .downcast::<Owner<UnsyncStorage>>()
                        .unwrap()
                }),
            )
            .map(|owner| *(Box::new(owner) as Box<dyn Any>).downcast().unwrap())
        })
    }
}

thread_local! {
    static SYNC_OWNER: RefCell<Option<Owner<SyncStorage>>> = const { RefCell::new(None) };
    static UNSYNC_OWNER: RefCell<Option<Owner<UnsyncStorage>>> = const { RefCell::new(None) };
}

/// Returns the current owner. This owner will be used to drop any `Copy` state that is created by the `generational-box` crate.
///
/// If an owner has been set with `with_owner`, that owner will be returned. Otherwise, the owner from the current scope will be returned.
pub fn current_owner<S: AnyStorage>() -> Owner<S> {
    let id = TypeId::of::<S>();
    let override_owner = if id == TypeId::of::<SyncStorage>() {
        SYNC_OWNER.with(|cell| {
            let owner = cell.borrow();

            owner.clone().map(|owner| {
                *(Box::new(owner) as Box<dyn Any>)
                    .downcast::<Owner<S>>()
                    .unwrap()
            })
        })
    } else {
        UNSYNC_OWNER.with(|cell| {
            cell.borrow().clone().map(|owner| {
                *(Box::new(owner) as Box<dyn Any>)
                    .downcast::<Owner<S>>()
                    .unwrap()
            })
        })
    };
    if let Some(owner) = override_owner {
        return owner;
    }

    // Otherwise get the owner from the current scope
    current_scope_id().expect("in a virtual dom").owner()
}

impl ScopeId {
    /// Get the owner for the current scope.
    pub fn owner<S: AnyStorage>(self) -> Owner<S> {
        match self.has_context() {
            Some(rt) => rt,
            None => {
                let owner = S::owner();
                self.provide_context(owner)
            }
        }
    }
}