mtfse/src/ext/hashers.rs

use super::*;
use std::hash::{BuildHasherDefault, Hasher};
use smallvec::SmallVec;
use cryptohelpers::sha256;
use ::bytes::Buf;
use cryptohelpers::sha2::{
    Sha256,
    Digest,
};
use std::borrow::BorrowMut;
use tokio::io::AsyncWrite;
use std::{
    pin::Pin,
    task::{Context, Poll},
    io,
};

/// A hasher that takes the first 8 bytes from SHA256 hash as its output.
///
/// # Notes
/// Intended for use for `HashSet` with a SHA256 key.
/// Hashing anything other than a SHA256 hash with this hasher is undefined.
#[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd, Default)]
pub struct Sha256TopHasher(SmallVec<[u8; sha256::SIZE]>);

/// A `BuildHasher` for [`Sha256TopHasher`].
pub type Sha256TopBuildHasher = BuildHasherDefault<Sha256TopHasher>;

impl Sha256TopHasher
{
    /// Create a new hasher
    #[inline] fn new() -> Self
    {
	Self(SmallVec::new())
    }
}

impl Hasher for Sha256TopHasher
{
    #[inline] fn finish(&self) -> u64 {
	let mut bytes = [0u8; std::mem::size_of::<u64>()];
	bytes::memcpy(self.0.as_ref(), &mut bytes[..]);
	u64::from_le_bytes(bytes)
    }
    #[inline] fn write(&mut self, bytes: &[u8]) {
	self.0.extend_from_slice(bytes);
    }
}

/// An `AsyncWrite` implementor that writes it's inputs to a sha256 digest.
#[pin_project]
#[derive(Debug)]
pub struct Sha256Sink<H: BorrowMut<Sha256> = Sha256>
{
    digest: H
}

impl<H: BorrowMut<Sha256>> Sha256Sink<H>
{
    /// Create a new Sha256-computing `AsyncWrite` sink.
    #[inline] pub fn new(digest: H) -> Self
    {
	Self{digest}
    }
    /// Consume into the inner digest
    #[inline] pub fn into_inner(self) -> H
    {
	self.digest
    }
    /// The inner digest
    #[inline] pub fn inner(&self) -> &H
    {
	&self.digest
    }
    /// The inner digest (mutable)
    #[inline] pub fn inner_mut(&mut self) -> &mut H
    {
	&mut self.digest
    }
    
    #[inline(always)] pub fn digest(&self) -> &Sha256
    {
	self.digest.borrow()
    }
    #[inline(always)] pub fn digest_mut(&mut self) -> &mut Sha256
    {
	self.digest.borrow_mut()
    }
}
impl<H: BorrowMut<Sha256>> AsRef<Sha256> for Sha256Sink<H>
{
    fn as_ref(&self) -> &Sha256 {
	self.digest.borrow()
    }
}
impl<H: BorrowMut<Sha256>> AsMut<Sha256> for Sha256Sink<H>
{
    fn as_mut(&mut self) -> &mut Sha256 {
	self.digest.borrow_mut()
    }
}
/*
impl<'a, H: BorrowMut<Sha256>> AsRef<Sha256> for &'a Sha256Sink<H>
    //where H: 'a
{
    fn as_ref(&self) -> &Sha256 {
	self.digest.borrow()
    }
}
impl<'a, H: BorrowMut<Sha256>> AsMut<Sha256> for &'a mut Sha256Sink<H>
   // where H: 'a
{
    fn as_mut(&mut self) -> &mut Sha256 {
	self.digest.borrow_mut()
    }
}
*/
impl<H: BorrowMut<Sha256>> BorrowMut<Sha256> for Sha256Sink<H>
{
    #[inline] fn borrow_mut(&mut self) -> &mut Sha256 {
	self.digest.borrow_mut()
    }
}

impl<H: BorrowMut<Sha256>> Borrow<Sha256> for Sha256Sink<H>
{
    #[inline] fn borrow(&self) -> &Sha256 {
	self.digest.borrow()
    }
}



impl<H: BorrowMut<Sha256>> AsyncWrite for Sha256Sink<H>
{
    fn poll_write(self: Pin<&mut Self>, _: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize, io::Error>> {
	let this = self.project();

	this.digest.borrow_mut().update(buf);
	Poll::Ready(Ok(buf.len()))
    }
    #[inline(always)] fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
	Poll::Ready(Ok(()))
    }
    #[inline(always)] fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
	Poll::Ready(Ok(()))

    }
}

pub trait Sha256HashExt
{
    fn compute_sha256_hash(&self) -> sha256::Sha256Hash;
}

pub trait Sha256HashOwnedExt: Sized
{
    fn into_sha256_hash(self) -> sha256::Sha256Hash;
}


impl<T: ?Sized> Sha256HashExt for T
where T: AsRef<[u8]>
{
    #[inline] fn compute_sha256_hash(&self) -> sha256::Sha256Hash {
	sha256::compute_slice(self.as_ref())
    }
}

impl<T: Buf> Sha256HashOwnedExt for T
{
    #[inline] fn into_sha256_hash(self) -> sha256::Sha256Hash {
	sha256::compute_sync(self.reader()).unwrap()
    }
}