use gentian::gentian;
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum ProtocolType {
SS,
TLS,
VMESS,
WS,
TROJAN,
DIRECT,
}
pub trait Protocol {
fn protocol_type(&self) -> ProtocolType;
}
struct ChainStreamBuilderProtocolTypeIter<'a> {
builders: &'a Vec<Box<dyn Protocol>>,
ty: Option<ProtocolType>,
pos: usize,
state: u32,
}
impl<'a> ChainStreamBuilderProtocolTypeIter<'a> {
fn new(
builders: &'a Vec<Box<dyn Protocol>>,
last_builder: &'a Option<Box<dyn Protocol>>,
) -> Self {
let mut ty = None;
if let Some(b) = last_builder {
ty = Some(b.protocol_type());
}
Self {
builders,
ty,
pos: builders.len(),
state: 0,
}
}
}
impl<'a> Iterator for ChainStreamBuilderProtocolTypeIter<'a> {
type Item = ProtocolType;
#[gentian]
#[gentian_attr(ret_val=None)]
fn next(&mut self) -> Option<Self::Item> {
if self.ty.is_some() {
co_yield(self.ty);
}
if self.pos == 0 {
return None;
}
while self.pos != 0 {
self.pos -= 1;
co_yield(Some(self.builders[self.pos].protocol_type()));
}
return None;
}
fn size_hint(&self) -> (usize, Option<usize>) {
if self.ty.is_some() {
return (0, Some(self.pos + 1));
}
(0, Some(self.pos))
}
}
macro_rules! impl_protocol {
($name:tt,$type:expr) => {
struct $name;
impl Protocol for $name {
fn protocol_type(&self) -> ProtocolType {
$type
}
}
};
}
impl_protocol!(Trojan, ProtocolType::TROJAN);
impl_protocol!(Shadowsocks, ProtocolType::SS);
impl_protocol!(Vmess, ProtocolType::VMESS);
impl_protocol!(Direct, ProtocolType::DIRECT);
impl_protocol!(Tls, ProtocolType::TLS);
impl_protocol!(Ws, ProtocolType::WS);
#[test]
fn test_iter_impl() {
use ProtocolType::{DIRECT, SS, TLS, TROJAN, VMESS, WS};
let builders: Vec<Box<dyn Protocol>> = vec![
Box::new(Vmess),
Box::new(Tls),
Box::new(Trojan),
Box::new(Shadowsocks),
Box::new(Ws),
Box::new(Vmess),
Box::new(Direct),
];
let last_builder: Option<Box<dyn Protocol>> = Some(Box::new(Trojan));
let expected = vec![TROJAN, DIRECT, VMESS, WS, SS, TROJAN, TLS, VMESS];
let my_iter = ChainStreamBuilderProtocolTypeIter::new(&builders, &last_builder);
for (real, expected) in my_iter.zip(expected.into_iter()) {
assert_eq!(real, expected);
}
let last_builder = None;
let expected = vec![DIRECT, VMESS, WS, SS, TROJAN, TLS, VMESS];
let my_iter = ChainStreamBuilderProtocolTypeIter::new(&builders, &last_builder);
for (real, expected) in my_iter.zip(expected.into_iter()) {
assert_eq!(real, expected);
}
}