transfer/src/ext.rs

use std::{
    mem,
    iter::{
	self, 
	ExactSizeIterator,
	FusedIterator,
    },
    slice,
    fmt,
};

pub use std::{
    marker::{
	Send, Sync, Unpin,
    },
    borrow::{ Borrow, BorrowMut },
    convert::{
	Infallible,
	TryFrom,
	TryInto
    },
};
pub use tokio::{
    io::{
	AsyncWriteExt,
	AsyncReadExt,
    },
    task::JoinHandle,
};

/// Make this type act as a reference to itself.
///
/// Implements `AsRef<T>` for type `T`.
#[macro_export] macro_rules! ref_self {
    ($type:ty) => {
	impl AsRef<$type> for $type
	{
	    #[inline] fn as_ref(&self) -> &$type
	    {
		self
	    }
	}

    }
}


#[derive(Debug, Clone)]
pub struct HexStringIter<I>(I, [u8; 2]);

impl<I: Iterator<Item = u8>> HexStringIter<I>
{
    /// Write this hex string iterator to a formattable buffer
    pub fn consume<F>(self, f: &mut F) -> fmt::Result
    where F: std::fmt::Write
    {
	if self.1[0] != 0 {
	    write!(f, "{}", self.1[0] as char)?;
	}
	if self.1[1] != 0 {
	    write!(f, "{}", self.1[1] as char)?;
	}

	for x in self.0 {
	    write!(f, "{:02x}", x)?;
	}
	
	Ok(())
    }

    /// Consume into a string
    pub fn into_string(self) -> String
    {
	let mut output = match self.size_hint() {
	    (0, None) => String::new(),
	    (_, Some(x)) |
	    (x, None) => String::with_capacity(x),
	};
	self.consume(&mut output).unwrap();
	output
    }
}

pub trait HexStringIterExt<I>: Sized
{
    fn into_hex(self) -> HexStringIter<I>;
}

pub type HexStringSliceIter<'a> = HexStringIter<iter::Copied<slice::Iter<'a, u8>>>;

pub trait HexStringSliceIterExt
{
    fn hex(&self) -> HexStringSliceIter<'_>;
}

impl<S> HexStringSliceIterExt for S
where S: AsRef<[u8]>
{
    fn hex(&self) -> HexStringSliceIter<'_>
    {
	self.as_ref().iter().copied().into_hex()
    }
}

impl<I: IntoIterator<Item=u8>> HexStringIterExt<I::IntoIter> for I
{
    #[inline] fn into_hex(self) -> HexStringIter<I::IntoIter> {
	HexStringIter(self.into_iter(), [0u8; 2])
    }
}

impl<I: Iterator<Item = u8>> Iterator for HexStringIter<I>
{
    type Item = char;
    fn next(&mut self) -> Option<Self::Item>
    {
	match self.1 {
	    [_, 0] => {
		use std::io::Write;
		write!(&mut self.1[..], "{:02x}", self.0.next()?).unwrap();

		Some(mem::replace(&mut self.1[0], 0) as char)
	    },
	    [0, _] => Some(mem::replace(&mut self.1[1], 0) as char),
	    _ => unreachable!(),
	}
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
	let (l, h) = self.0.size_hint();

	(l * 2, h.map(|x| x*2))
    }
}

impl<I: Iterator<Item = u8> + ExactSizeIterator> ExactSizeIterator for HexStringIter<I>{}
impl<I: Iterator<Item = u8> + FusedIterator> FusedIterator for HexStringIter<I>{}

impl<I: Iterator<Item = u8>> From<HexStringIter<I>> for String
{
    fn from(from: HexStringIter<I>) -> Self
    {
	from.into_string()
    }
}

impl<I: Iterator<Item = u8> + Clone> fmt::Display for HexStringIter<I>
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result
    {
	self.clone().consume(f)
    }
}

pub trait CollectArrayExt<T>: Sized
{
    /// Collect an iterator into an array.
    ///
    /// If the iterator has more elements than `N`, the rest are discarded.
    ///
    /// # Panics
    /// If the iterator has **less** elements than `N`.
    fn collect_array<const N: usize>(self) -> [T; N];
}

impl<I> CollectArrayExt<I::Item> for I
where I: Iterator
{
    fn collect_array<const N: usize>(self) -> [I::Item; N] {
	use std::mem::MaybeUninit;
	
	// SAFETY: This pattern is safe. The array elements are still maybeuninit.
	let mut out = unsafe { MaybeUninit::<[MaybeUninit::<I::Item>; N]>::uninit().assume_init() };
	let mut init_to = 0;
	
	if N == 0 {
	    // SAFETY: This is valid, [I::Item; N] is 0 sized. (i uhh think...)
	    return unsafe { MaybeUninit::<[I::Item; N]>::uninit().assume_init() };
	}

	let res = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
	    #[cold]
	    #[inline(never)]
	    fn _panic_bad_size(exp: usize, got: usize) -> !
	    {
		panic!("tried to collect into array of size {}, when iterator is only {} elements", exp, got)
	    }
	    init_to = out.iter_mut().zip(self)
		.map(|(o, i)| *o = MaybeUninit::new(i)).count();
	    match init_to
	    {
		n if n == N => (),
		got => _panic_bad_size(N, got),
	    }
	}));

	match res {
	    Ok(()) => {
		// SAFETY: Transmuting MaybeUninit<T> to T is fine.
		// All elements are initialised by this point
		unsafe {
		    
		    #[inline(always)] unsafe fn assume_init_array<T, const N: usize>(array: [MaybeUninit<T>; N]) -> [T; N]
		    {
			//std::intrinsics::assert_inhabited::<[T; N]>();
			(&array as *const _ as *const [T; N]).read()
		    }
		    //MaybeUninit::array_assume_init(out)
		    assume_init_array(out)
		}
	    },
	    Err(e) => {
		// Drop all initialised elements before resuming unwind.
		unsafe {
		    std::ptr::drop_in_place(&mut out[..init_to] as *mut [MaybeUninit<I::Item>] as *mut [I::Item]);
		}
		std::panic::resume_unwind(e)
	    },
	}
    }
}

#[macro_export] macro_rules! prog1 {
    ($first:expr, $($rest:expr);+ $(;)?) => {
	($first, $( $rest ),+).0
    }
}

#[macro_export] macro_rules! r#if {
    ($if:expr, $then:expr, $else:expr) => {
	if $if { $then } else { $else }
    }
}