text: // Copyright (c) Microsoft Corporation. All rights reserved.
  // Licensed under the MIT license.
  
  #include "seal/context.h"
  #include "seal/keygenerator.h"
  #include "seal/modulus.h"
  #include "seal/relinkeys.h"
  #include "seal/util/polyarithsmallmod.h"
  #include "seal/util/uintcore.h"
  #include "gtest/gtest.h"
  
  using namespace seal;
  using namespace seal::util;
  using namespace std;
  
  namespace sealtest
  {
      TEST(RelinKeysTest, RelinKeysSaveLoad)
      {
          stringstream stream;
          {
              EncryptionParameters parms(scheme_type::BFV);
              parms.set_poly_modulus_degree(64);
              parms.set_plain_modulus(1 << 6);
              parms.set_coeff_modulus(CoeffModulus::Create(64, { 60, 60 }));
              auto context = SEALContext::Create(parms, false, sec_level_type::none);
              KeyGenerator keygen(context);
  
              RelinKeys keys;
              RelinKeys test_keys;
              keys = keygen.relin_keys_local();
              keys.save(stream);
              test_keys.load(context, stream);
              ASSERT_EQ(keys.size(), test_keys.size());
              ASSERT_TRUE(keys.parms_id() == test_keys.parms_id());
              for (size_t j = 0; j < test_keys.size(); j++)
              {
                  for (size_t i = 0; i < test_keys.key(j + 2).size(); i++)
                  {
                      ASSERT_EQ(keys.key(j + 2)[i].data().size(), test_keys.key(j + 2)[i].data().size());
                      ASSERT_EQ(
                          keys.key(j + 2)[i].data().int_array().size(),
                          test_keys.key(j + 2)[i].data().int_array().size());
                      ASSERT_TRUE(is_equal_uint_uint(
                          keys.key(j + 2)[i].data().data(), test_keys.key(j + 2)[i].data().data(),
                          keys.key(j + 2)[i].data().int_array().size()));
                  }
              }
          }
          {
              EncryptionParameters parms(scheme_type::BFV);
              parms.set_poly_modulus_degree(256);
              parms.set_plain_modulus(1 << 6);
              parms.set_coeff_modulus(CoeffModulus::Create(256, { 60, 50 }));
  
              auto context = SEALContext::Create(parms, false, sec_level_type::none);
              KeyGenerator keygen(context);
  
              RelinKeys keys;
              RelinKeys test_keys;
              keys = keygen.relin_keys_local();
              keys.save(stream);
              test_keys.load(context, stream);
              ASSERT_EQ(keys.size(), test_keys.size());
              ASSERT_TRUE(keys.parms_id() == test_keys.parms_id());
              for (size_t j = 0; j < test_keys.size(); j++)
              {
                  for (size_t i = 0; i < test_keys.key(j + 2).size(); i++)
                  {
                      ASSERT_EQ(keys.key(j + 2)[i].data().size(), test_keys.key(j + 2)[i].data().size());
                      ASSERT_EQ(
                          keys.key(j + 2)[i].data().int_array().size(),
                          test_keys.key(j + 2)[i].data().int_array().size());
                      ASSERT_TRUE(is_equal_uint_uint(
                          keys.key(j + 2)[i].data().data(), test_keys.key(j + 2)[i].data().data(),
                          keys.key(j + 2)[i].data().int_array().size()));
                  }
              }
          }
      }
      TEST(RelinKeysTest, RelinKeysSeededSaveLoad)
      {
          // Returns true if a, b contains the same error.
          auto compare_kswitchkeys = [](const KSwitchKeys &a, const KSwitchKeys &b, const SecretKey &sk,
                                        shared_ptr<SEALContext> context) {
              auto compare_error = [](const Ciphertext &a_ct, const Ciphertext &b_ct, const SecretKey &sk1,
                                      shared_ptr<SEALContext> context1) {
                  auto get_error = [](const Ciphertext &encrypted, const SecretKey &sk2,
                                      shared_ptr<SEALContext> context2) {
                      auto pool = MemoryManager::GetPool();
                      auto &context_data = *context2->get_context_data(encrypted.parms_id());
                      auto &parms = context_data.parms();
                      auto &coeff_modulus = parms.coeff_modulus();
                      size_t coeff_count = parms.poly_modulus_degree();
                      size_t coeff_modulus_size = coeff_modulus.size();
                      size_t rns_poly_uint64_count = util::mul_safe(coeff_count, coeff_modulus_size);
  
                      IntArray<Ciphertext::ct_coeff_type> error;
                      error.resize(rns_poly_uint64_count);
                      auto destination = error.begin();
  
                      auto copy_operand1(util::allocate_uint(coeff_count, pool));
                      for (size_t i = 0; i < coeff_modulus_size; i++)
                      {
                          // Initialize pointers for multiplication
                          const uint64_t *encrypted_ptr = encrypted.data(1) + (i * coeff_count);
                          const uint64_t *secret_key_ptr = sk2.data().data() + (i * coeff_count);
                          uint64_t *destination_ptr = destination + (i * coeff_count);
                          util::set_zero_uint(coeff_count, destination_ptr);
                          util::set_uint_uint(encrypted_ptr, coeff_count, copy_operand1.get());
                          // compute c_{j+1} * s^{j+1}
                          util::dyadic_product_coeffmod(
                              copy_operand1.get(), secret_key_ptr, coeff_count, coeff_modulus[i], copy_operand1.get());
                          // add c_{j+1} * s^{j+1} to destination
                          util::add_poly_poly_coeffmod(
                              destination_ptr, copy_operand1.get(), coeff_count, coeff_modulus[i], destination_ptr);
                          // add c_0 into destination
                          util::add_poly_poly_coeffmod(
                              destination_ptr, encrypted.data() + (i * coeff_count), coeff_count, coeff_modulus[i],
                              destination_ptr);
                      }
                      return error;
                  };
  
                  auto error_a = get_error(a_ct, sk1, context1);
                  auto error_b = get_error(b_ct, sk1, context1);
                  ASSERT_EQ(error_a.size(), error_b.size());
                  ASSERT_TRUE(is_equal_uint_uint(error_a.cbegin(), error_b.cbegin(), error_a.size()));
              };
  
              ASSERT_EQ(a.size(), b.size());
              auto iter_a = a.data().begin();
              auto iter_b = b.data().begin();
              for (; iter_a != a.data().end(); iter_a++, iter_b++)
              {
                  ASSERT_EQ(iter_a->size(), iter_b->size());
                  auto pk_a = iter_a->begin();
                  auto pk_b = iter_b->begin();
                  for (; pk_a != iter_a->end(); pk_a++, pk_b++)
                  {
                      compare_error(pk_a->data(), pk_b->data(), sk, context);
                  }
              }
          };
  
          stringstream stream;
          {
              EncryptionParameters parms(scheme_type::BFV);
              parms.set_poly_modulus_degree(8);
              parms.set_plain_modulus(65537);
              parms.set_coeff_modulus(CoeffModulus::Create(8, { 60, 60 }));
              random_seed_type seed;
              for (auto &i : seed)
              {
                  i = random_uint64();
              }
              auto rng = make_shared<BlakePRNGFactory>(BlakePRNGFactory(seed));
              parms.set_random_generator(rng);
              auto context = SEALContext::Create(parms, false, sec_level_type::none);
              KeyGenerator keygen(context);
              SecretKey secret_key = keygen.secret_key();
  
              keygen.relin_keys().save(stream);
              RelinKeys test_keys;
              test_keys.load(context, stream);
              RelinKeys keys = keygen.relin_keys_local();
              compare_kswitchkeys(keys, test_keys, secret_key, context);
          }
          {
              EncryptionParameters parms(scheme_type::BFV);
              parms.set_poly_modulus_degree(256);
              parms.set_plain_modulus(65537);
              parms.set_coeff_modulus(CoeffModulus::Create(256, { 60, 50 }));
              random_seed_type seed;
              for (auto &i : seed)
              {
                  i = random_uint64();
              }
              auto rng = make_shared<BlakePRNGFactory>(BlakePRNGFactory(seed));
              parms.set_random_generator(rng);
              auto context = SEALContext::Create(parms, false, sec_level_type::none);
              KeyGenerator keygen(context);
              SecretKey secret_key = keygen.secret_key();
  
              keygen.relin_keys().save(stream);
              RelinKeys test_keys;
              test_keys.load(context, stream);
              RelinKeys keys = keygen.relin_keys_local();
              compare_kswitchkeys(keys, test_keys, secret_key, context);
          }
      }
  } // namespace sealtest
  