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

use std::cell::UnsafeCell;
use std::rc::Rc;
use std::marker::PhantomData;
use std::os::raw::c_void;

use crate::sender_rc::SenderInner;
use crate::priv_traits::*;
use crate::utils::transform_sink_lifetime;
use crate::Stream;

pub struct FlatMap<'a, T1, T2, S: Into<Stream<'a,T2>>, F: FnMut(T1) -> S>{
    func: F,
    sender: Rc<UnsafeCell<SenderInner<'a,T2>>>,
    next_stream: Option<Stream<'a, T2>>,
    _marker: PhantomData<T1>,
}

impl<'a,T1, T2, S: Into<Stream<'a,T2>>, F: FnMut(T1) -> S> FlatMap<'a,T1,T2,S,F>{
    pub fn new(f: F) -> FlatMap<'a,T1,T2,S,F>{
        FlatMap{
            func: f,
            next_stream: None,
            sender: Rc::new(UnsafeCell::new(SenderInner::new())),
            _marker: PhantomData,
        }
    }
}

impl<'a,T1, T2: Clone, S: Into<Stream<'a,T2>>, F: FnMut(T1) -> S> StreamInner<'a,T2> for FlatMap<'a,T1,T2,S,F>{
    #[inline]
    fn push(&mut self, st: Rc<UnsafeCell<dyn SinkInner<'a,T2>+'a>>){
        unsafe{ (*self.sender.get()).push(st) }
    }

    #[inline]
    fn send(&mut self, t: T2){
        unsafe{ (*self.sender.get()).send(t) }
    }
}

impl<'a,'b,T1, T2, S: Into<Stream<'a,T2>>, F: FnMut(T1) -> S> Remove for FlatMap<'a,T1,T2,S,F>{
    fn remove_raw(&mut self, st: *const c_void){
        unsafe{ (*self.sender.get()).remove_raw(st) }
    }
}

impl<'a,T1, T2: Clone, S: Into<Stream<'a,T2>>, F: FnMut(T1) -> S> SinkInner<'a, T1> for FlatMap<'a,T1,T2,S,F>{
    #[inline]
    fn call(&mut self, t: T1) -> Result<(),()>{
        let mut stream: Stream<'a,T2> = (self.func)(t).into();
        unsafe{ stream.push(transform_sink_lifetime(self.sender.clone())) };
        self.next_stream = Some(stream.into());
        Ok(())
    }

    #[inline]
    fn is_alive(&self) -> bool {
        true
    }
}