jam-std-common 0.1.26

use codec::{Encode, Output};
use jam_types::Hash;
use tiny_keccak::{Hasher, Keccak};

pub fn keccak(data: &[u8]) -> Hash {
	let mut hasher = Keccak::v256();
	hasher.update(data);
	let mut output = [0u8; 32];
	hasher.finalize(&mut output);
	output
}

/// Compute Keccak hash of the concatenation of byte slices.
pub fn keccak_concat<I, T>(items: I) -> Hash
where
	I: IntoIterator<Item = T>,
	T: AsRef<[u8]>,
{
	let mut hasher = Keccak::v256();
	for item in items.into_iter() {
		hasher.update(item.as_ref());
	}
	let mut output = [0u8; 32];
	hasher.finalize(&mut output);
	output
}

pub fn hash_raw(data: &[u8]) -> Hash {
	let mut hasher = Blake2bHasher::new();
	hasher.update(data);
	hasher.into_hash()
}

/// Compute the hash of the concatenation of byte slices.
pub fn hash_raw_concat<I, T>(items: I) -> Hash
where
	I: IntoIterator<Item = T>,
	T: AsRef<[u8]>,
{
	let mut hasher = Blake2bHasher::new();
	for item in items.into_iter() {
		hasher.update(item.as_ref());
	}
	hasher.into_hash()
}

pub fn hash_encoded(what: &impl Encode) -> Hash {
	let mut hasher = Blake2bHasher::new();
	what.encode_to(&mut hasher);
	hasher.into_hash()
}

struct Blake2bHasher {
	state: blake2b_simd::State,
}

impl Blake2bHasher {
	fn new() -> Self {
		let state = blake2b_simd::Params::new().hash_length(32).to_state();
		Self { state }
	}

	fn update(&mut self, data: &[u8]) {
		self.state.update(data);
	}

	fn into_hash(self) -> Hash {
		self.state.finalize().as_bytes().try_into().expect("Hash length set to 32")
	}
}

impl Output for Blake2bHasher {
	fn write(&mut self, bytes: &[u8]) {
		self.update(bytes);
	}
}

/// A wrapper for [`Input`](codec::Input) that hashes the bytes as they are read from the
/// underlying input.
pub struct HashedInput<'a, I: codec::Input> {
	inner: &'a mut I,
	hasher: Blake2bHasher,
}

impl<'a, I: codec::Input> HashedInput<'a, I> {
	/// Create new wrapper from the supplied underlying input.
	pub fn new(inner: &'a mut I) -> Self {
		Self { inner, hasher: Blake2bHasher::new() }
	}

	/// Consume the wrapper returning the hash of all successfully read bytes.
	pub fn into_hash(self) -> Hash {
		self.hasher.into_hash()
	}
}

impl<I: codec::Input> codec::Input for HashedInput<'_, I> {
	fn remaining_len(&mut self) -> Result<Option<usize>, codec::Error> {
		self.inner.remaining_len()
	}

	fn read(&mut self, buf: &mut [u8]) -> Result<(), codec::Error> {
		self.inner.read(buf)?;
		self.hasher.update(buf);
		Ok(())
	}

	fn read_byte(&mut self) -> Result<u8, codec::Error> {
		self.inner.read_byte().inspect(|byte| self.hasher.update(&[*byte]))
	}

	fn ascend_ref(&mut self) {
		self.inner.ascend_ref()
	}

	fn descend_ref(&mut self) -> Result<(), codec::Error> {
		self.inner.descend_ref()
	}

	fn on_before_alloc_mem(&mut self, size: usize) -> Result<(), codec::Error> {
		self.inner.on_before_alloc_mem(size)
	}
}

#[cfg(test)]
mod tests {
	use super::*;
	use codec::Decode;

	#[derive(Encode, Decode, PartialEq, Eq, Debug)]
	struct Dummy {
		x: Vec<Dummy>,
		y: u8,
		z: String,
	}

	#[test]
	fn hashed_input_works() {
		let dummy = Dummy {
			x: vec![Dummy { x: Vec::new(), y: 123, z: "Hello".into() }],
			y: 133,
			z: "world".into(),
		};
		let bytes = dummy.encode();
		let expected_hash = hash_raw(&bytes[..]);
		let mut input = &bytes[..];
		let mut hashed_input = HashedInput::new(&mut input);
		let actual = Dummy::decode(&mut hashed_input).unwrap();
		assert_eq!(dummy, actual);
		let hash = hashed_input.into_hash();
		assert_eq!(expected_hash, hash);
	}

	#[test]
	fn hash_encoded_works() {
		let dummy = Dummy {
			x: vec![Dummy { x: Vec::new(), y: 123, z: "Hello".into() }],
			y: 133,
			z: "world".into(),
		};
		assert_eq!(hash_raw(&dummy.encode()), hash_encoded(&dummy));
		assert_eq!(dummy.using_encoded(hash_raw), hash_encoded(&dummy));
	}

	#[test]
	fn hash_raw_concat_works() {
		assert_eq!(
			hash_raw(&[b"x".as_ref(), b"y", b"z"].concat()),
			hash_raw_concat([b"x", b"y", b"z"])
		);
	}

	#[test]
	fn keccak_concat_works() {
		assert_eq!(
			keccak(&[b"x".as_ref(), b"y", b"z"].concat()),
			keccak_concat([b"x", b"y", b"z"])
		);
	}
}