sealy 0.2.0

use std::ffi::c_void;
use std::fmt::Debug;
use std::ptr::null_mut;
use std::sync::atomic::AtomicPtr;
use std::sync::atomic::Ordering;

use crate::bindgen;
use crate::error::*;
use crate::try_seal;
use crate::Ciphertext;
use crate::Context;
use crate::PublicKey;
use crate::SecretKey;

/// A SEAL array storing a number of polynomials. In particular, type implements
/// methods that can convert from other types (like public keys) and converts
/// them into the same non-NTT RNS format. Enables conversion of RNS format to
/// multiprecision and back losslessly.
pub struct PolynomialArray {
	handle: AtomicPtr<c_void>,
}

impl PolynomialArray {
	/// Creates a new empty polynomial array. Use an encoder to populate with a value.
	pub fn new() -> Result<Self> {
		let mut handle: *mut c_void = null_mut();

		// By giving an empty pool handle we acquire the global one (first argument to create).
		try_seal!(unsafe { bindgen::PolynomialArray_Create(null_mut(), &mut handle) })?;

		Ok(Self {
			handle: AtomicPtr::new(handle),
		})
	}

	/// Creates a polynomial array from a reference to a ciphertext.
	pub fn new_from_ciphertext(
		context: &Context,
		ciphertext: &Ciphertext,
	) -> Result<Self> {
		let mut handle: *mut c_void = null_mut();

		// By giving an empty pool handle we acquire the global one (first argument to create).
		try_seal!(unsafe {
			bindgen::PolynomialArray_CreateFromCiphertext(
				null_mut(),
				context.get_handle(),
				ciphertext.get_handle(),
				&mut handle,
			)
		})?;

		Ok(Self {
			handle: AtomicPtr::new(handle),
		})
	}

	/// Creates a polynomial array from a reference to a public key.
	pub fn new_from_public_key(
		context: &Context,
		public_key: &PublicKey,
	) -> Result<Self> {
		let mut handle: *mut c_void = null_mut();

		// By giving an empty pool handle we acquire the global one (first argument to create).
		try_seal!(unsafe {
			bindgen::PolynomialArray_CreateFromPublicKey(
				null_mut(),
				context.get_handle(),
				public_key.get_handle(),
				&mut handle,
			)
		})?;

		Ok(Self {
			handle: AtomicPtr::new(handle),
		})
	}

	/// Creates a polynomial array from a reference to a secret key.
	pub fn new_from_secret_key(
		context: &Context,
		secret_key: &SecretKey,
	) -> Result<Self> {
		let mut handle: *mut c_void = null_mut();

		// By giving an empty pool handle we acquire the global one (first argument to create).
		try_seal!(unsafe {
			bindgen::PolynomialArray_CreateFromSecretKey(
				null_mut(),
				context.get_handle(),
				secret_key.get_handle(),
				&mut handle,
			)
		})?;

		Ok(Self {
			handle: AtomicPtr::new(handle),
		})
	}

	/// Returns the handle to the underlying SEAL object.
	pub(crate) unsafe fn get_handle(&self) -> *mut c_void {
		self.handle.load(Ordering::SeqCst)
	}

	/// Has the array data been loaded? When an array is created, it initially
	/// has no data. Once data is loaded this is true. Additionally data can only
	/// be loaded once.
	pub fn is_reserved(&self) -> bool {
		let mut is_reserved: bool = false;

		try_seal!(unsafe {
			bindgen::PolynomialArray_IsReserved(self.get_handle(), &mut is_reserved)
		})
		.expect("Fatal error in PolynomialArray::is_reserved()");

		is_reserved
	}

	/// Is the array in RNS form (true).
	pub fn is_rns(&self) -> bool {
		let mut is_rns: bool = false;

		try_seal!(unsafe { bindgen::PolynomialArray_IsRns(self.get_handle(), &mut is_rns) })
			.expect("Fatal error in PolynomialArray::is_rns()");

		is_rns
	}

	/// Is the array in RNS form (true).
	pub fn is_multiprecision(&self) -> bool {
		!self.is_rns()
	}

	/// Converts the polynomial array into the RNS format regardless of its
	/// current format.
	pub fn to_rns(&self) {
		try_seal!(unsafe { bindgen::PolynomialArray_ToRns(self.get_handle()) })
			.expect("Fatal error in PolynomialArray::to_rns()");
	}

	/// Converts the polynomial array into the multiprecision format regardless
	/// of its current format.
	pub fn to_multiprecision(&self) {
		try_seal!(unsafe { bindgen::PolynomialArray_ToMultiprecision(self.get_handle()) })
			.expect("Fatal error in PolynomialArray::to_multiprecision()");
	}

	/// Read out the raw data inside the vector. This function is not meant to be
	/// called by users; instead users should call `as_multiprecision_u64s` or
	/// `as_rns_u64s`
	pub fn as_u64s(&self) -> Result<Vec<u64>> {
		let mut num_u64: u64 = 0;

		try_seal!(unsafe { bindgen::PolynomialArray_ExportSize(self.get_handle(), &mut num_u64) })?;

		let mut data: Vec<u64> = Vec::with_capacity(num_u64 as usize);

		try_seal!(unsafe {
			let data_ptr = data.as_mut_ptr();

			bindgen::PolynomialArray_PerformExport(self.get_handle(), data_ptr)
		})?;

		unsafe { data.set_len(num_u64 as usize) };

		Ok(data)
	}

	/// This will be in coefficient order; all the limbs with a given coefficient
	/// are stored together in least significant order.
	///
	/// The number of limbs equals the number of moduli in the coefficient
	/// modulus.
	pub fn as_multiprecision_u64s(&self) -> Result<Vec<u64>> {
		// We want to leave the array in the format it was before, so we check
		// if we need to perform a conversion and temporarily do so if required.
		let perform_conversion = self.is_rns();

		if perform_conversion {
			self.to_multiprecision();
		}

		let result = self.as_u64s();

		if perform_conversion {
			self.to_rns();
		}

		result
	}

	/// This will be in modulus order; all the values associated with a given
	/// moduli are stored together.
	///
	/// The number of limbs equals the number of moduli in the coefficient
	/// modulus.
	pub fn as_rns_u64s(&self) -> Result<Vec<u64>> {
		// We want to leave the array in the format it was before, so we check
		// if we need to perform a conversion and temporarily do so if required.
		let perform_conversion = self.is_multiprecision();

		if perform_conversion {
			self.to_rns();
		}

		let result = self.as_u64s();

		if perform_conversion {
			self.to_multiprecision();
		}

		result
	}

	/// Returns the number of polynomials stored in the `PolynomialArray`.
	pub fn num_polynomials(&self) -> u64 {
		let mut size: u64 = 0;

		try_seal!(unsafe { bindgen::PolynomialArray_PolySize(self.get_handle(), &mut size) })
			.expect("Fatal error in PolynomialArray::num_polynomials");

		size
	}

	/// Returns the number of coefficients in each polynomial in the `PolynomialArray`.
	pub fn poly_modulus_degree(&self) -> u64 {
		let mut size: u64 = 0;

		try_seal!(unsafe {
			bindgen::PolynomialArray_PolyModulusDegree(self.get_handle(), &mut size)
		})
		.expect("Fatal error in PolynomialArray::poly_modulus_degree");

		size
	}

	/// Returns how many moduli are in the coefficient modulus set.
	pub fn coeff_modulus_size(&self) -> u64 {
		let mut size: u64 = 0;

		try_seal!(unsafe {
			bindgen::PolynomialArray_CoeffModulusSize(self.get_handle(), &mut size)
		})
		.expect("Fatal error in PolynomialArray::coeff_modulus_size");

		size
	}

	/// Reduces the polynomial array by dropping the last modulus in the modulus
	/// set.
	pub fn drop_modulus(&self) -> Result<Self> {
		if self.coeff_modulus_size() == 1 {
			return Err(Error::ModulusChainTooSmall);
		}

		let mut handle: *mut c_void = null_mut();

		try_seal!(unsafe { bindgen::PolynomialArray_Drop(self.get_handle(), &mut handle) })
			.expect("Fatal error in PolynomialArray::coeff_modulus_size");

		Ok(Self {
			handle: AtomicPtr::new(handle),
		})
	}
}

impl Clone for PolynomialArray {
	fn clone(&self) -> Self {
		let mut handle = null_mut();

		try_seal!(unsafe { bindgen::PolynomialArray_Copy(self.get_handle(), &mut handle) })
			.expect("Fatal error: Failed to clone polynomial array");

		Self {
			handle: AtomicPtr::new(handle),
		}
	}
}

impl PartialEq for PolynomialArray {
	fn eq(
		&self,
		other: &Self,
	) -> bool {
		self.as_rns_u64s() == other.as_rns_u64s()
	}
}

impl Drop for PolynomialArray {
	fn drop(&mut self) {
		try_seal!(unsafe { bindgen::PolynomialArray_Destroy(self.get_handle()) })
			.expect("Internal error in PolynomialArray::drop()");
	}
}

impl Debug for PolynomialArray {
	fn fmt(
		&self,
		f: &mut std::fmt::Formatter<'_>,
	) -> std::fmt::Result {
		f.debug_struct("PolynomialArray")
			.field("handle", &self.handle)
			.finish()
	}
}

#[cfg(test)]
mod tests {
	use crate::{
		AsymmetricComponents, BFVEncoder, BFVEncryptionParametersBuilder,
		CoefficientModulusFactory, DegreeType, Encryptor, KeyGenerator, Modulus,
		PlainModulusFactory, Plaintext, SecurityLevel,
	};

	use super::*;

	#[test]
	fn can_create_and_destroy_static_polynomial_array() {
		let poly_array = PolynomialArray::new().unwrap();

		std::mem::drop(poly_array);
	}

	#[test]
	fn can_create_polynomial_from_ciphertext() {
		let params = BFVEncryptionParametersBuilder::new()
			.set_poly_modulus_degree(DegreeType::D8192)
			.set_coefficient_modulus(
				CoefficientModulusFactory::build(DegreeType::D8192, &[50, 30, 30, 50, 50]).unwrap(),
			)
			.set_plain_modulus_u64(1234)
			.build()
			.unwrap();

		let ctx = Context::new(&params, false, SecurityLevel::TC128).unwrap();

		let ciphertext = Ciphertext::new().unwrap();
		let poly_array = PolynomialArray::new_from_ciphertext(&ctx, &ciphertext).unwrap();

		assert!(poly_array.is_reserved())
	}

	#[test]
	fn polynomial_array_initially_not_reserved() {
		let poly_array = PolynomialArray::new().unwrap();
		assert!(!poly_array.is_reserved());
	}

	fn generate_ciphertext_example() -> (
		Context,
		Vec<Modulus>,
		PublicKey,
		Ciphertext,
		PolynomialArray,
		PolynomialArray,
		Plaintext,
	) {
		let coeff_modulus =
			CoefficientModulusFactory::build(DegreeType::D8192, &[50, 30, 30, 50, 50]).unwrap();
		let params = BFVEncryptionParametersBuilder::new()
			.set_poly_modulus_degree(DegreeType::D8192)
			.set_coefficient_modulus(coeff_modulus.clone())
			.set_plain_modulus(PlainModulusFactory::batching(DegreeType::D8192, 20).unwrap())
			.build()
			.unwrap();

		let ctx = Context::new(&params, false, SecurityLevel::TC128).unwrap();
		let gen = KeyGenerator::new(&ctx).unwrap();

		let encoder = BFVEncoder::new(&ctx).unwrap();

		let mut data = vec![];

		for i in 0..encoder.get_slot_count() {
			data.push(i as u64)
		}

		let plaintext = encoder.encode_u64(&data).unwrap();

		let public_key = gen.create_public_key();
		let secret_key = gen.secret_key();

		let encryptor =
			Encryptor::with_public_and_secret_key(&ctx, &public_key, &secret_key).unwrap();

		let (
			ciphertext,
			AsymmetricComponents {
				u,
				e,
				r,
			},
		) = encryptor.encrypt_return_components(&plaintext).unwrap();

		(ctx, coeff_modulus, public_key, ciphertext, u, e, r)
	}

	#[test]
	fn correct_poly_array_sizes_from_ciphertext() {
		let (ctx, _, public_key, ciphertext, u, e, r) = generate_ciphertext_example();
		let poly_array_public_key =
			PolynomialArray::new_from_public_key(&ctx, &public_key).unwrap();
		let poly_array_ciphertext =
			PolynomialArray::new_from_ciphertext(&ctx, &ciphertext).unwrap();

		let poly_array_ciphertext_encoded = poly_array_ciphertext.as_u64s().unwrap();
		let poly_array_public_key_encoded = poly_array_public_key.as_u64s().unwrap();
		let u_encoded = u.as_u64s().unwrap();
		let e_encoded = e.as_u64s().unwrap();

		assert!(poly_array_public_key.is_reserved());
		assert!(poly_array_ciphertext.is_reserved());
		assert!(u.is_reserved());
		assert!(e.is_reserved());

		// Ciphertext size checks
		assert_eq!(
			poly_array_ciphertext.num_polynomials(),
			ciphertext.num_polynomials()
		);
		assert_eq!(poly_array_ciphertext.poly_modulus_degree(), 8192);
		assert_eq!(
			poly_array_ciphertext.coeff_modulus_size(),
			ciphertext.coeff_modulus_size()
		);

		assert_eq!(
			poly_array_ciphertext_encoded.len() as u64,
			poly_array_ciphertext.num_polynomials()
				* poly_array_ciphertext.poly_modulus_degree()
				* poly_array_ciphertext.coeff_modulus_size()
		);

		// Public key
		assert_eq!(poly_array_public_key.num_polynomials(), 2);
		assert_eq!(poly_array_public_key.poly_modulus_degree(), 8192);
		assert_eq!(poly_array_public_key.coeff_modulus_size(), 4);

		assert_eq!(
			poly_array_public_key_encoded.len() as u64,
			poly_array_public_key.num_polynomials()
				* poly_array_public_key.poly_modulus_degree()
				* poly_array_public_key.coeff_modulus_size()
		);

		// u
		assert_eq!(u.num_polynomials(), 1);
		assert_eq!(u.poly_modulus_degree(), 8192);
		assert_eq!(u.coeff_modulus_size(), 4);
		assert_eq!(
			u_encoded.len() as u64,
			poly_array_public_key.poly_modulus_degree()
				* poly_array_ciphertext.coeff_modulus_size()
		);

		// e
		assert_eq!(e.num_polynomials(), 2);
		assert_eq!(e.poly_modulus_degree(), 8192);
		assert_eq!(e.coeff_modulus_size(), 4);
		assert_eq!(
			e_encoded.len() as u64,
			ciphertext.num_polynomials()
				* poly_array_public_key.poly_modulus_degree()
				* ciphertext.coeff_modulus_size()
		);

		// r
		assert_eq!(r.len(), 8192);
	}

	#[test]
	fn multiprecision_and_back_is_identity() {
		let (ctx, _, _, ciphertext, _, _, _) = generate_ciphertext_example();
		let poly_array = PolynomialArray::new_from_ciphertext(&ctx, &ciphertext).unwrap();

		let poly_array_encoded_original = poly_array.as_u64s().unwrap();

		poly_array.to_multiprecision();
		poly_array.to_rns();

		let poly_array_encoded_round_trip = poly_array.as_u64s().unwrap();

		assert_eq!(poly_array_encoded_original, poly_array_encoded_round_trip);
	}

	#[test]
	fn poly_array_u64_and_bytes_match() {
		let (ctx, _, _, ciphertext, _, _, _) = generate_ciphertext_example();
		let poly_array = PolynomialArray::new_from_ciphertext(&ctx, &ciphertext).unwrap();

		let poly_array_encoded = poly_array.as_u64s().unwrap();

		let cipher_data = (0..poly_array_encoded.len())
			.map(|i| ciphertext.get_data(i).unwrap())
			.collect::<Vec<u64>>();

		assert_eq!(&poly_array_encoded, &cipher_data);
	}

	#[test]
	fn poly_array_clone() {
		let (_, _, _, _, u, e, _) = generate_ciphertext_example();

		let u_clone = u.clone();
		assert_eq!(u, u_clone);
		assert_eq!(u.as_rns_u64s(), u_clone.as_rns_u64s());
		unsafe { assert_ne!(u.get_handle(), u_clone.get_handle()) };

		let e_clone = e.clone();
		assert_eq!(e, e_clone);
		assert_eq!(e.as_rns_u64s(), e_clone.as_rns_u64s());
		unsafe { assert_ne!(e.get_handle(), e_clone.get_handle()) };
	}

	#[test]
	fn poly_array_drop() {
		let (ctx, _, public_key, _, _, _, _) = generate_ciphertext_example();
		let poly_array = PolynomialArray::new_from_public_key(&ctx, &public_key).unwrap();
		let poly_array_lower = poly_array.drop_modulus().unwrap();

		assert_eq!(poly_array_lower.num_polynomials(), 2);
		assert_eq!(poly_array_lower.poly_modulus_degree(), 8192);
		assert_eq!(poly_array_lower.coeff_modulus_size(), 3);

		let poly_array_lower_rns = poly_array_lower.as_rns_u64s().unwrap();
		let poly_array_rns = poly_array
			.as_rns_u64s()
			.unwrap()
			.into_iter()
			.take(poly_array_lower_rns.len())
			.collect::<Vec<u64>>();

		assert_eq!(poly_array_rns, poly_array_lower_rns);
	}
}