jam-std-common 0.1.29

use crate::{CountedOutput, RecordedInput};
use arrayvec::ArrayVec;
use codec::{ConstEncodedLen, Decode, Encode, Error, Input, MaxEncodedLen, Output};

/// `T` wrapper with a constant encoded length (`N`) and an injective decode implementation.
///
/// Note that the decode implementation will _not_ be injective if `T::decode` accepts
/// non-canonical encodings.
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum InjectiveDecode<T, const N: usize> {
	Valid(T),
	Invalid([u8; N]),
}

impl<T, const N: usize> InjectiveDecode<T, N> {
	pub fn valid(&self) -> Option<&T> {
		match self {
			Self::Valid(value) => Some(value),
			Self::Invalid(_) => None,
		}
	}
}

impl<T: MaxEncodedLen, const N: usize> Encode for InjectiveDecode<T, N> {
	fn size_hint(&self) -> usize {
		N
	}

	fn encode_to<O: Output + ?Sized>(&self, output: &mut O) {
		match self {
			Self::Valid(value) => {
				let mut counted_output = CountedOutput::new(output);
				debug_assert!(T::max_encoded_len() <= N);
				value.encode_to(&mut counted_output);
				// Pad with zeroes
				while counted_output.num_bytes_written() < N {
					counted_output.push_byte(0);
				}
				debug_assert_eq!(counted_output.num_bytes_written(), N);
			},
			Self::Invalid(encoded) => output.write(encoded),
		}
	}

	fn using_encoded<R, F: FnOnce(&[u8]) -> R>(&self, f: F) -> R {
		match self {
			Self::Valid(_) => f(&self.encode()),
			Self::Invalid(encoded) => f(encoded),
		}
	}

	fn encoded_size(&self) -> usize {
		N
	}
}

impl<T: MaxEncodedLen, const N: usize> MaxEncodedLen for InjectiveDecode<T, N> {
	fn max_encoded_len() -> usize {
		N
	}
}

impl<T: MaxEncodedLen, const N: usize> ConstEncodedLen for InjectiveDecode<T, N> {}

impl<T: MaxEncodedLen + Decode, const N: usize> Decode for InjectiveDecode<T, N> {
	fn decode<I: Input>(input: &mut I) -> Result<Self, Error> {
		let mut recorded_input = RecordedInput::with_record(input, ArrayVec::<u8, N>::new());

		debug_assert!(T::max_encoded_len() <= N);
		let mut value = T::decode(&mut recorded_input);
		if recorded_input.failed() {
			value?;
			// Expect error to have been forwarded by Decode impl...
			debug_assert!(false);
			return Err("Unforwarded input failure".into())
		}

		// Consume padding bytes; these must be zero
		while !recorded_input.record().is_full() {
			if (recorded_input.read_byte()? != 0) && value.is_ok() {
				value = Err("Non-zero padding".into());
			}
		}

		match value {
			Ok(value) => Ok(Self::Valid(value)),
			Err(_) =>
				Ok(Self::Invalid(recorded_input.into_record().into_inner().expect("Filled above"))),
		}
	}

	fn skip<I: Input>(input: &mut I) -> Result<(), Error> {
		for _ in 0..N {
			input.read_byte()?;
		}
		Ok(())
	}

	fn encoded_fixed_size() -> Option<usize> {
		Some(N)
	}
}

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

	#[test]
	fn encode_valid_pads_with_zeroes() {
		let x = InjectiveDecode::<u16, 4>::Valid(0x0201);
		assert_eq!(&x.encode(), &[0x01, 0x02, 0x00, 0x00]);
		assert_eq!(x.encoded_size(), 4);
		assert_eq!(x.size_hint(), 4);
	}

	#[test]
	fn encode_invalid_writes_raw_bytes() {
		let x = InjectiveDecode::<u16, 4>::Invalid([1, 2, 3, 4]);
		assert_eq!(&x.encode(), &[1, 2, 3, 4]);
	}

	#[test]
	fn decode_valid_with_zero_padding() {
		let encoded = [0x01, 0x02, 0x00, 0x00];
		let decoded = InjectiveDecode::<u16, 4>::decode(&mut &encoded[..]).unwrap();
		assert_eq!(decoded, InjectiveDecode::Valid(0x0201));
	}

	#[test]
	fn decode_nonzero_padding_is_invalid() {
		// The inner u16 decodes fine, but a non-zero padding byte makes the whole value invalid,
		// capturing the raw bytes verbatim
		let encoded = [0x01, 0x02, 0x00, 0x05];
		let decoded = InjectiveDecode::<u16, 4>::decode(&mut &encoded[..]).unwrap();
		assert_eq!(decoded, InjectiveDecode::Invalid(encoded));
	}

	#[test]
	fn decode_invalid_value_is_captured() {
		// 0x02 is not a valid bool encoding, so the value decode fails and the raw bytes are
		// captured rather than the error being propagated
		let encoded = [0x02, 0x00];
		let decoded = InjectiveDecode::<bool, 2>::decode(&mut &encoded[..]).unwrap();
		assert_eq!(decoded, InjectiveDecode::Invalid(encoded));
	}

	#[test]
	fn decode_is_injective() {
		// Re-encoding any decoded value must reproduce the original bytes exactly
		for a in 0..=4u8 {
			for b in 0..=4u8 {
				for c in 0..=1u8 {
					for d in 0..=1u8 {
						let encoded = [a, b, c, d];
						let decoded = InjectiveDecode::<u16, 4>::decode(&mut &encoded[..]).unwrap();
						assert_eq!(&decoded.encode(), &encoded);
					}
				}
			}
		}
	}

	#[test]
	fn decode_fails_on_truncated_input() {
		let encoded = [0x01];
		assert!(InjectiveDecode::<u16, 4>::decode(&mut &encoded[..]).is_err());
	}

	#[test]
	fn skip() {
		let mut input = [1, 2, 3, 4, 5, 6].as_slice();
		InjectiveDecode::<u16, 4>::skip(&mut input).unwrap();
		assert_eq!(input, &[5, 6]);
		assert!(InjectiveDecode::<u16, 4>::skip(&mut input).is_err());
	}
}