false-bottom 0.3.4

// SPDX-License-Identifier: GPL-3.0-or-later
use crypto_bigint::{ArrayEncoding, Uint, generic_array::GenericArray};
use crate::{FbError, PrimeField};

/* PACKING SCHEME
 * This scheme transforms byte chunks into field elements and vice versa.
 * Every x byte chunks are combined to form a field element n.
 * In case the number formed is n >= P where P is the right bound of the prime field,
 * the values P-1 (i.e LIMIT) followed by n-LIMIT are appended to the output vec.
 * While unpacking, LIMIT is used as a signaling element to extract n
 * from the successive element after adding back LIMIT.
 */

pub(crate) trait Packing
where
	Self: Sized
{
	const LIMIT: Self;
	fn pack(inp: &[u8]) -> Vec<Self>;
	fn unpack(inp: Vec<Self>) -> Result<Vec<u8>, FbError>;
}

impl<const LIMBS: usize> Packing for Uint<LIMBS>
where
	Uint<LIMBS>: ArrayEncoding + PrimeField,
{
	const LIMIT: Uint<LIMBS> = Uint::<LIMBS>::PRIME.wrapping_sub(&Uint::<LIMBS>::ONE);
	fn pack(inp: &[u8]) -> Vec<Uint<LIMBS>> {
		let mut out: Vec<_> = inp.chunks(Uint::<LIMBS>::BYTES)
			.flat_map(|inp_chunk| {
				let mut out_chunk = GenericArray::<u8, <Uint<LIMBS> as ArrayEncoding>::ByteSize>::default();
				out_chunk[..inp_chunk.len()].copy_from_slice(inp_chunk);
				let mut out_uint = Uint::from_le_byte_array(out_chunk);
				if out_uint >= Self::LIMIT {
					out_uint = out_uint.wrapping_sub(&Self::LIMIT);
					vec![Self::LIMIT, out_uint]
				} else {
					vec![out_uint]
				}
			})
			.collect();
		let inp_chunk_last = inp.chunks_exact(Uint::<LIMBS>::BYTES)
			.remainder();
		let pad_uint = Uint::from_u8((Uint::<LIMBS>::BYTES - inp_chunk_last.len()) as u8);
		out.push(pad_uint);

		out
	}

	fn unpack(inp: Vec<Uint<LIMBS>>) -> Result<Vec<u8>, FbError> {
		let mut pad_len: usize = inp.last()
			.ok_or(FbError::InvalidKey)?
			.to_be_byte_array()[Uint::<LIMBS>::BYTES - 1] as usize;
		if pad_len >= Uint::<LIMBS>::BYTES || pad_len < 1 {
			return Err(FbError::InvalidKey);
		}
		pad_len += Uint::<LIMBS>::BYTES;
		let mut extract = false;
		let mut out: Vec<u8> = inp.into_iter()
			.filter_map(|el| {
				if extract {
					extract = false;
					let el = el.wrapping_add(&Self::LIMIT);
					Some(el.to_le_byte_array())
				} else if el == Self::LIMIT {
					extract = true;
					None
				} else {
					Some(el.to_le_byte_array())
				}
			})
			.flatten()
			.collect();
		let trunc_len = out.len().checked_sub(pad_len)
			.ok_or(FbError::InvalidKey)?;
		out.truncate(trunc_len);

		Ok(out)
	}
}