jam-std-common 0.1.29

use crate::{BytesCursor, RecordedInput, hash_raw};
use bytes::Bytes;
use castaway::cast;
use codec::{Decode, Encode, Input, MaxEncodedLen};
use jam_types::{Hash, Wrap};
use std::{borrow::Borrow, ops::Deref};

/// `HashOrNoHash` is implemented for hash types (eg `WorkPackageHash`) and also `NoHash`.
///
/// `NoHash` is a unit type, used when we don't care about the hash in question.
pub trait HashOrNoHash {
	/// Returns the hash of `data`.
	fn hash(data: &[u8]) -> Self;
}

#[derive(Clone, Copy, Debug)]
pub struct NoHash;

impl HashOrNoHash for NoHash {
	fn hash(_data: &[u8]) -> Self {
		Self
	}
}

impl<T: From<Hash>> HashOrNoHash for T {
	fn hash(data: &[u8]) -> Self {
		hash_raw(data).into()
	}
}

/// An `Immutable<T, H>` stores an _immutable_ `T` plus various derived values.
///
/// One of the derived values is the hash of the encoded `T`, which has type `H`. `H` may be
/// `NoHash` if we do not care about this hash.
#[derive(Clone, Debug)]
pub struct Immutable<T, H = NoHash> {
	value: T,
	/// `value.encode()`.
	encoded: Bytes,
	/// Hash of `encoded`. This may have type `NoHash`.
	hash: H,
}

impl<T, H: HashOrNoHash> Immutable<T, H> {
	/// Construct an `Immutable<T, H>` from a `T`, its hash and its encoding.
	pub fn from_parts(value: T, hash: H, encoded: Bytes) -> Self {
		Self { value, encoded, hash }
	}

	/// Construct an `Immutable<T, H>` from a `T` and its encoding.
	pub fn from_value_and_encoded(value: T, encoded: Bytes) -> Self {
		let hash = H::hash(&encoded);
		Self { value, encoded, hash }
	}

	/// Returns the encoding of the wrapped value.
	pub fn encoded(&self) -> &Bytes {
		&self.encoded
	}

	/// Returns the length of the encoding of the wrapped value.
	pub fn encoded_len(&self) -> usize {
		self.encoded.len()
	}

	/// Returns the wrapped value.
	pub fn into_value(self) -> T {
		self.value
	}

	/// Returns the encoding of the wrapped value.
	#[allow(dead_code)]
	pub fn into_encoded(self) -> Bytes {
		self.encoded
	}

	/// Returns the wrapped value and its encoding.
	pub fn into_value_and_encoded(self) -> (T, Bytes) {
		(self.value, self.encoded)
	}

	/// Returns the wrapped value, hash and its encoding.
	pub fn into_parts(self) -> (T, Bytes, H) {
		(self.value, self.encoded, self.hash)
	}
}

// The Into<Hash> bound is to avoid exposing the hash() method when H is NoHash
impl<T, H: Into<Hash> + Copy> Immutable<T, H> {
	/// Returns the hash of the encoding of the wrapped value.
	pub fn hash(&self) -> H {
		self.hash
	}
}

// This is provided primarily for tests...
impl<T: Encode, H: HashOrNoHash> From<T> for Immutable<T, H> {
	fn from(value: T) -> Self {
		let encoded = value.encode().into();
		Self::from_value_and_encoded(value, encoded)
	}
}

impl<T: Encode, H: HashOrNoHash> Immutable<T, H> {
	pub fn modify_and_encode(&mut self, f: impl FnOnce(&mut T)) {
		f(&mut self.value);
		self.encoded = self.value.encode().into();
		self.hash = H::hash(&self.encoded);
	}
}

impl<T, H> Deref for Immutable<T, H> {
	type Target = T;

	fn deref(&self) -> &T {
		&self.value
	}
}

impl<T, H> Borrow<T> for Immutable<T, H> {
	fn borrow(&self) -> &T {
		&self.value
	}
}

impl<T: Decode, H: HashOrNoHash> Decode for Immutable<T, H> {
	fn decode<I: Input>(input: &mut I) -> Result<Self, codec::Error> {
		let (value, encoded) = match cast!(input, &mut BytesCursor) {
			Ok(input) => {
				let mut encoded = input.remaining().clone();
				let value = T::decode(input)?;
				encoded.truncate(encoded.len() - input.infallible_remaining_len());
				(value, encoded)
			},
			Err(input) => {
				let mut input = RecordedInput::new(input);
				let value = T::decode(&mut input)?;
				debug_assert!(!input.failed());
				let encoded: Vec<_> = input.into_record();
				(value, encoded.into())
			},
		};
		let hash = H::hash(&encoded);
		Ok(Self { value, encoded, hash })
	}

	fn skip<I: Input>(input: &mut I) -> Result<(), codec::Error> {
		T::skip(input)
	}

	fn encoded_fixed_size() -> Option<usize> {
		T::encoded_fixed_size()
	}
}

impl<T, H> Encode for Immutable<T, H> {
	fn size_hint(&self) -> usize {
		self.encoded.len()
	}

	fn using_encoded<R, F: FnOnce(&[u8]) -> R>(&self, f: F) -> R {
		f(&self.encoded)
	}

	fn encoded_size(&self) -> usize {
		self.encoded.len()
	}
}

impl<T: MaxEncodedLen, H> MaxEncodedLen for Immutable<T, H> {
	fn max_encoded_len() -> usize {
		T::max_encoded_len()
	}
}

// These traits are derived to avoid running into issues with Rust's poor handling of derive with
// generic parameters; see https://github.com/rust-lang/rust/issues/26925
#[derive(Clone, Debug)]
pub struct ImmutableWrap;

impl Wrap for ImmutableWrap {
	type Wrap<T> = Immutable<T>;
}

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

	#[derive(Encode, Decode, Debug, PartialEq, Eq)]
	struct Inner(Bytes);

	#[derive(Encode, Decode, Debug, PartialEq, Eq)]
	struct WrappedOuter<W: Wrap> {
		a: W::Wrap<Inner>,
		b: W::Wrap<Inner>,
	}

	type Outer = WrappedOuter<NoWrap>;
	type ImmutableOuter = Immutable<WrappedOuter<ImmutableWrap>, Hash>;

	fn contains(outer: &[u8], inner: &[u8]) -> bool {
		let outer = outer.as_ptr_range();
		let inner = inner.as_ptr_range();
		(outer.start <= inner.start) && (inner.end <= outer.end)
	}

	#[test]
	fn basic_operation() {
		let outer =
			Outer { a: Inner(Bytes::from_static(b"polka")), b: Inner(Bytes::from_static(b"jam!")) };
		let encoded_outer = outer.encode();
		assert_eq!(outer, Outer::decode_all(&mut encoded_outer.as_slice()).unwrap());

		let encoded_outer = Bytes::from(encoded_outer);
		let immutable_outer: ImmutableOuter =
			BytesCursor::from(encoded_outer.clone()).decode_all().unwrap();
		assert_eq!(immutable_outer.a.0, outer.a.0);
		assert_eq!(immutable_outer.b.0, outer.b.0);
		assert_eq!(immutable_outer.encoded(), &encoded_outer);
		assert_eq!(immutable_outer.hash(), hash_raw(&encoded_outer));
		assert_eq!(immutable_outer.a.encoded(), &outer.a.encode());
		assert_eq!(immutable_outer.b.encoded(), &outer.b.encode());
		assert!(contains(&encoded_outer, immutable_outer.encoded()));
		assert!(contains(&encoded_outer, immutable_outer.a.encoded()));
		assert!(contains(&encoded_outer, immutable_outer.b.encoded()));
		assert!(contains(&encoded_outer, &immutable_outer.a.0));
		assert!(contains(&encoded_outer, &immutable_outer.b.0));
	}
}