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
103
104
105

use std::{
    pin::Pin,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
    task::{Context, Poll, Wake, Waker},
};

use atomic_waker::AtomicWaker;
use futures_core::Stream;

use crate::AsyncComponent;

#[derive(Debug)]
pub struct ComponentStream<C> {
    cx: StateContext,
    component: C,
}

impl<C: AsyncComponent> ComponentStream<C> {
    /// Create new [`ComponentStream`]
    pub fn new(func: impl FnOnce(&StateContext) -> C) -> Self {
        let cx = StateContext::new();

        let component = func(&cx);
        Self { cx, component }
    }

    pub fn component(&self) -> &C {
        &self.component
    }

    pub fn component_mut(&mut self) -> &mut C {
        &mut self.component
    }
}

impl<C: AsyncComponent> Stream for ComponentStream<C> {
    type Item = ();

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<()>> {
        self.cx.inner.waker.register(cx.waker());

        self.component.update_component();
        if self.cx.inner.updated.swap(false, Ordering::Relaxed) {
            Poll::Ready(Some(()))
        } else {
            Poll::Pending
        }
    }
}

impl<C: AsyncComponent> Unpin for ComponentStream<C> {}

#[derive(Debug, Clone)]
pub struct StateContext {
    inner: Arc<Inner>,
    waker: Waker,
}

impl StateContext {
    pub(crate) fn new() -> Self {
        let inner = Arc::new(Inner::default());
        let waker = Waker::from(inner.clone());

        StateContext { inner, waker }
    }

    /// Signal context to wake
    pub fn signal(&self) {
        self.waker.wake_by_ref();
    }

    /// Returns [`Context`] which can be used for polling future
    pub fn task_context<'a>(&'a self) -> Context<'a> {
        Context::from_waker(&self.waker)
    }
}

#[derive(Debug)]
struct Inner {
    updated: AtomicBool,
    waker: AtomicWaker,
}

impl Wake for Inner {
    fn wake(self: Arc<Self>) {
        self.wake_by_ref()
    }

    fn wake_by_ref(self: &Arc<Self>) {
        self.updated.store(true, Ordering::Relaxed);
        self.waker.wake()
    }
}

impl Default for Inner {
    fn default() -> Self {
        Self {
            updated: AtomicBool::new(true),
            waker: Default::default(),
        }
    }
}