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

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

use crate::Sender;
use crate::priv_traits::*;
use crate::{stream, Stream};


pub enum Either<T1, T2>{
    Left(T1),
    Right(T2),
    Both(T1, T2),
    None,
}

pub struct PartitionMap<'a, T1, T2, F>{
    func: F,
    left: Sender<'a,T1>,
    right: Sender<'a,T2>,
}

impl<'a, T1, T2: 'a, F> PartitionMap<'a,T1,T2,F>{
    pub fn new(f: F) -> (PartitionMap<'a,T1,T2,F>, Stream<'a, T2>){
        let (sender_right, stream_right) = stream();
        let filter = PartitionMap{
            func: f,
            left: Sender::new(),
            right: sender_right,
        };
        (filter, stream_right)
    }
}

impl<'a, T1, T2, F> StreamInner<'a,T1> for PartitionMap<'a,T1,T2,F> {
    #[inline]
    fn push(&mut self, st: Rc<UnsafeCell<dyn SinkInner<'a,T1>+'a>>){
        self.left.push(st)
    }

    #[inline]
    fn send(&mut self, t: T1){
        self.left.send(t)
    }
}

impl<'a,T1, T2, F> Remove for PartitionMap<'a,T1,T2,F>{
    fn remove_raw(&mut self, st: *const c_void){
        self.left.remove_raw(st)
    }
}

impl<'a, T, T1, T2, F> SinkInner<'a, T> for PartitionMap<'a,T1,T2,F>
where F: FnMut(T) -> Either<T1, T2>
{
    #[inline]
    fn call(&mut self, t: T) -> Result<(),()>{
        match (self.func)(t) {
            Either::Left(t1) => self.left.send(t1),
            Either::Right(t2) => self.right.send(t2),
            Either::Both(t1, t2) => {
                self.left.send(t1);
                self.right.send(t2);
            }
            Either::None => (),
        }
        Ok(())
    }

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