text: // Copyright (c) Microsoft Corporation. All rights reserved.
  // Licensed under the MIT license.
  
  #include "seal/context.h"
  #include "seal/decryptor.h"
  #include "seal/encryptor.h"
  #include "seal/evaluator.h"
  #include "seal/keygenerator.h"
  #include "seal/valcheck.h"
  #include "gtest/gtest.h"
  
  using namespace seal;
  using namespace seal::util;
  using namespace std;
  
  namespace sealtest
  {
      TEST(KeyGeneratorTest, BFVKeyGeneration)
      {
          EncryptionParameters parms(scheme_type::bfv);
          {
              parms.set_poly_modulus_degree(64);
              parms.set_plain_modulus(65537);
              parms.set_coeff_modulus(CoeffModulus::Create(64, { 60 }));
              SEALContext context(parms, false, sec_level_type::none);
              KeyGenerator keygen(context);
  
              ASSERT_THROW(auto evk = keygen.create_relin_keys(), logic_error);
              ASSERT_THROW(auto galk = keygen.create_galois_keys(), logic_error);
          }
          {
              parms.set_poly_modulus_degree(64);
              parms.set_plain_modulus(65537);
              parms.set_coeff_modulus(CoeffModulus::Create(64, { 60, 60 }));
              SEALContext context(parms, false, sec_level_type::none);
              KeyGenerator keygen(context);
  
              RelinKeys evk;
              keygen.create_relin_keys(evk);
              ASSERT_TRUE(evk.parms_id() == context.key_parms_id());
              ASSERT_EQ(1ULL, evk.key(2).size());
              for (auto &a : evk.data())
              {
                  for (auto &b : a)
                  {
                      ASSERT_FALSE(b.data().is_transparent());
                  }
              }
              ASSERT_TRUE(is_valid_for(evk, context));
  
              GaloisKeys galks;
              keygen.create_galois_keys(galks);
              for (auto &a : galks.data())
              {
                  for (auto &b : a)
                  {
                      ASSERT_FALSE(b.data().is_transparent());
                  }
              }
              ASSERT_TRUE(is_valid_for(galks, context));
  
              ASSERT_TRUE(galks.parms_id() == context.key_parms_id());
              ASSERT_EQ(1ULL, galks.key(3).size());
              ASSERT_EQ(10ULL, galks.size());
  
              keygen.create_galois_keys(vector<uint32_t>{ 1, 3, 5, 7 }, galks);
              ASSERT_TRUE(galks.parms_id() == context.key_parms_id());
              ASSERT_TRUE(galks.has_key(1));
              ASSERT_TRUE(galks.has_key(3));
              ASSERT_TRUE(galks.has_key(5));
              ASSERT_TRUE(galks.has_key(7));
              ASSERT_FALSE(galks.has_key(9));
              ASSERT_FALSE(galks.has_key(127));
              ASSERT_EQ(1ULL, galks.key(1).size());
              ASSERT_EQ(1ULL, galks.key(3).size());
              ASSERT_EQ(1ULL, galks.key(5).size());
              ASSERT_EQ(1ULL, galks.key(7).size());
              ASSERT_EQ(4ULL, galks.size());
  
              keygen.create_galois_keys(vector<uint32_t>{ 1 }, galks);
              ASSERT_TRUE(galks.parms_id() == context.key_parms_id());
              ASSERT_TRUE(galks.has_key(1));
              ASSERT_FALSE(galks.has_key(3));
              ASSERT_FALSE(galks.has_key(127));
              ASSERT_EQ(1ULL, galks.key(1).size());
              ASSERT_EQ(1ULL, galks.size());
  
              keygen.create_galois_keys(vector<uint32_t>{ 127 }, galks);
              ASSERT_TRUE(galks.parms_id() == context.key_parms_id());
              ASSERT_FALSE(galks.has_key(1));
              ASSERT_TRUE(galks.has_key(127));
              ASSERT_EQ(1ULL, galks.key(127).size());
              ASSERT_EQ(1ULL, galks.size());
          }
          {
              parms.set_poly_modulus_degree(256);
              parms.set_plain_modulus(65537);
              parms.set_coeff_modulus(CoeffModulus::Create(256, { 60, 30, 30 }));
  
              SEALContext context(parms, false, sec_level_type::none);
              KeyGenerator keygen(context);
  
              RelinKeys evk;
              keygen.create_relin_keys(evk);
              ASSERT_TRUE(evk.parms_id() == context.key_parms_id());
              ASSERT_EQ(2ULL, evk.key(2).size());
  
              for (auto &a : evk.data())
              {
                  for (auto &b : a)
                  {
                      ASSERT_FALSE(b.data().is_transparent());
                  }
              }
              ASSERT_TRUE(is_valid_for(evk, context));
  
              GaloisKeys galks;
              keygen.create_galois_keys(galks);
              for (auto &a : galks.data())
              {
                  for (auto &b : a)
                  {
                      ASSERT_FALSE(b.data().is_transparent());
                  }
              }
              ASSERT_TRUE(is_valid_for(galks, context));
  
              ASSERT_TRUE(galks.parms_id() == context.key_parms_id());
              ASSERT_EQ(2ULL, galks.key(3).size());
              ASSERT_EQ(14ULL, galks.size());
  
              keygen.create_galois_keys(vector<uint32_t>{ 1, 3, 5, 7 }, galks);
              ASSERT_TRUE(galks.parms_id() == context.key_parms_id());
              ASSERT_TRUE(galks.has_key(1));
              ASSERT_TRUE(galks.has_key(3));
              ASSERT_TRUE(galks.has_key(5));
              ASSERT_TRUE(galks.has_key(7));
              ASSERT_FALSE(galks.has_key(9));
              ASSERT_FALSE(galks.has_key(511));
              ASSERT_EQ(2ULL, galks.key(1).size());
              ASSERT_EQ(2ULL, galks.key(3).size());
              ASSERT_EQ(2ULL, galks.key(5).size());
              ASSERT_EQ(2ULL, galks.key(7).size());
              ASSERT_EQ(4ULL, galks.size());
  
              keygen.create_galois_keys(vector<uint32_t>{ 1 }, galks);
              ASSERT_TRUE(galks.parms_id() == context.key_parms_id());
              ASSERT_TRUE(galks.has_key(1));
              ASSERT_FALSE(galks.has_key(3));
              ASSERT_FALSE(galks.has_key(511));
              ASSERT_EQ(2ULL, galks.key(1).size());
              ASSERT_EQ(1ULL, galks.size());
  
              keygen.create_galois_keys(vector<uint32_t>{ 511 }, galks);
              ASSERT_TRUE(galks.parms_id() == context.key_parms_id());
              ASSERT_FALSE(galks.has_key(1));
              ASSERT_TRUE(galks.has_key(511));
              ASSERT_EQ(2ULL, galks.key(511).size());
              ASSERT_EQ(1ULL, galks.size());
          }
      }
  
      TEST(KeyGeneratorTest, BGVKeyGeneration)
      {
          EncryptionParameters parms(scheme_type::bgv);
          {
              parms.set_poly_modulus_degree(64);
              parms.set_plain_modulus(65537);
              parms.set_coeff_modulus(CoeffModulus::Create(64, { 60 }));
              SEALContext context(parms, false, sec_level_type::none);
              KeyGenerator keygen(context);
  
              ASSERT_THROW(auto evk = keygen.create_relin_keys(), logic_error);
              ASSERT_THROW(auto galk = keygen.create_galois_keys(), logic_error);
          }
          {
              parms.set_poly_modulus_degree(64);
              parms.set_plain_modulus(65537);
              parms.set_coeff_modulus(CoeffModulus::Create(64, { 60, 60 }));
              SEALContext context(parms, false, sec_level_type::none);
              KeyGenerator keygen(context);
  
              RelinKeys evk;
              keygen.create_relin_keys(evk);
              ASSERT_TRUE(evk.parms_id() == context.key_parms_id());
              ASSERT_EQ(1ULL, evk.key(2).size());
              for (auto &a : evk.data())
              {
                  for (auto &b : a)
                  {
                      ASSERT_FALSE(b.data().is_transparent());
                  }
              }
              ASSERT_TRUE(is_valid_for(evk, context));
  
              GaloisKeys galks;
              keygen.create_galois_keys(galks);
              for (auto &a : galks.data())
              {
                  for (auto &b : a)
                  {
                      ASSERT_FALSE(b.data().is_transparent());
                  }
              }
              ASSERT_TRUE(is_valid_for(galks, context));
  
              ASSERT_TRUE(galks.parms_id() == context.key_parms_id());
              ASSERT_EQ(1ULL, galks.key(3).size());
              ASSERT_EQ(10ULL, galks.size());
  
              keygen.create_galois_keys(vector<uint32_t>{ 1, 3, 5, 7 }, galks);
              ASSERT_TRUE(galks.parms_id() == context.key_parms_id());
              ASSERT_TRUE(galks.has_key(1));
              ASSERT_TRUE(galks.has_key(3));
              ASSERT_TRUE(galks.has_key(5));
              ASSERT_TRUE(galks.has_key(7));
              ASSERT_FALSE(galks.has_key(9));
              ASSERT_FALSE(galks.has_key(127));
              ASSERT_EQ(1ULL, galks.key(1).size());
              ASSERT_EQ(1ULL, galks.key(3).size());
              ASSERT_EQ(1ULL, galks.key(5).size());
              ASSERT_EQ(1ULL, galks.key(7).size());
              ASSERT_EQ(4ULL, galks.size());
  
              keygen.create_galois_keys(vector<uint32_t>{ 1 }, galks);
              ASSERT_TRUE(galks.parms_id() == context.key_parms_id());
              ASSERT_TRUE(galks.has_key(1));
              ASSERT_FALSE(galks.has_key(3));
              ASSERT_FALSE(galks.has_key(127));
              ASSERT_EQ(1ULL, galks.key(1).size());
              ASSERT_EQ(1ULL, galks.size());
  
              keygen.create_galois_keys(vector<uint32_t>{ 127 }, galks);
              ASSERT_TRUE(galks.parms_id() == context.key_parms_id());
              ASSERT_FALSE(galks.has_key(1));
              ASSERT_TRUE(galks.has_key(127));
              ASSERT_EQ(1ULL, galks.key(127).size());
              ASSERT_EQ(1ULL, galks.size());
          }
          {
              parms.set_poly_modulus_degree(256);
              parms.set_plain_modulus(65537);
              parms.set_coeff_modulus(CoeffModulus::Create(256, { 60, 30, 30 }));
  
              SEALContext context(parms, false, sec_level_type::none);
              KeyGenerator keygen(context);
  
              RelinKeys evk;
              keygen.create_relin_keys(evk);
              ASSERT_TRUE(evk.parms_id() == context.key_parms_id());
              ASSERT_EQ(2ULL, evk.key(2).size());
  
              for (auto &a : evk.data())
              {
                  for (auto &b : a)
                  {
                      ASSERT_FALSE(b.data().is_transparent());
                  }
              }
              ASSERT_TRUE(is_valid_for(evk, context));
  
              GaloisKeys galks;
              keygen.create_galois_keys(galks);
              for (auto &a : galks.data())
              {
                  for (auto &b : a)
                  {
                      ASSERT_FALSE(b.data().is_transparent());
                  }
              }
              ASSERT_TRUE(is_valid_for(galks, context));
  
              ASSERT_TRUE(galks.parms_id() == context.key_parms_id());
              ASSERT_EQ(2ULL, galks.key(3).size());
              ASSERT_EQ(14ULL, galks.size());
  
              keygen.create_galois_keys(vector<uint32_t>{ 1, 3, 5, 7 }, galks);
              ASSERT_TRUE(galks.parms_id() == context.key_parms_id());
              ASSERT_TRUE(galks.has_key(1));
              ASSERT_TRUE(galks.has_key(3));
              ASSERT_TRUE(galks.has_key(5));
              ASSERT_TRUE(galks.has_key(7));
              ASSERT_FALSE(galks.has_key(9));
              ASSERT_FALSE(galks.has_key(511));
              ASSERT_EQ(2ULL, galks.key(1).size());
              ASSERT_EQ(2ULL, galks.key(3).size());
              ASSERT_EQ(2ULL, galks.key(5).size());
              ASSERT_EQ(2ULL, galks.key(7).size());
              ASSERT_EQ(4ULL, galks.size());
  
              keygen.create_galois_keys(vector<uint32_t>{ 1 }, galks);
              ASSERT_TRUE(galks.parms_id() == context.key_parms_id());
              ASSERT_TRUE(galks.has_key(1));
              ASSERT_FALSE(galks.has_key(3));
              ASSERT_FALSE(galks.has_key(511));
              ASSERT_EQ(2ULL, galks.key(1).size());
              ASSERT_EQ(1ULL, galks.size());
  
              keygen.create_galois_keys(vector<uint32_t>{ 511 }, galks);
              ASSERT_TRUE(galks.parms_id() == context.key_parms_id());
              ASSERT_FALSE(galks.has_key(1));
              ASSERT_TRUE(galks.has_key(511));
              ASSERT_EQ(2ULL, galks.key(511).size());
              ASSERT_EQ(1ULL, galks.size());
          }
      }
  
      TEST(KeyGeneratorTest, CKKSKeyGeneration)
      {
          EncryptionParameters parms(scheme_type::ckks);
          {
              parms.set_poly_modulus_degree(64);
              parms.set_coeff_modulus(CoeffModulus::Create(64, { 60 }));
              SEALContext context(parms, false, sec_level_type::none);
              KeyGenerator keygen(context);
  
              ASSERT_THROW(auto evk = keygen.create_relin_keys(), logic_error);
              ASSERT_THROW(auto galk = keygen.create_galois_keys(), logic_error);
          }
          {
              parms.set_poly_modulus_degree(64);
              parms.set_coeff_modulus(CoeffModulus::Create(64, { 60, 60 }));
  
              SEALContext context(parms, false, sec_level_type::none);
              KeyGenerator keygen(context);
  
              RelinKeys evk;
              keygen.create_relin_keys(evk);
              ASSERT_TRUE(evk.parms_id() == context.key_parms_id());
              ASSERT_EQ(1ULL, evk.key(2).size());
              for (auto &a : evk.data())
              {
                  for (auto &b : a)
                  {
                      ASSERT_FALSE(b.data().is_transparent());
                  }
              }
              ASSERT_TRUE(is_valid_for(evk, context));
  
              GaloisKeys galks;
              keygen.create_galois_keys(galks);
              for (auto &a : galks.data())
              {
                  for (auto &b : a)
                  {
                      ASSERT_FALSE(b.data().is_transparent());
                  }
              }
              ASSERT_TRUE(is_valid_for(galks, context));
  
              ASSERT_TRUE(galks.parms_id() == context.key_parms_id());
              ASSERT_EQ(1ULL, galks.key(3).size());
              ASSERT_EQ(10ULL, galks.size());
  
              keygen.create_galois_keys(vector<uint32_t>{ 1, 3, 5, 7 }, galks);
              ASSERT_TRUE(galks.parms_id() == context.key_parms_id());
              ASSERT_TRUE(galks.has_key(1));
              ASSERT_TRUE(galks.has_key(3));
              ASSERT_TRUE(galks.has_key(5));
              ASSERT_TRUE(galks.has_key(7));
              ASSERT_FALSE(galks.has_key(9));
              ASSERT_FALSE(galks.has_key(127));
              ASSERT_EQ(1ULL, galks.key(1).size());
              ASSERT_EQ(1ULL, galks.key(3).size());
              ASSERT_EQ(1ULL, galks.key(5).size());
              ASSERT_EQ(1ULL, galks.key(7).size());
              ASSERT_EQ(4ULL, galks.size());
  
              keygen.create_galois_keys(vector<uint32_t>{ 1 }, galks);
              ASSERT_TRUE(galks.parms_id() == context.key_parms_id());
              ASSERT_TRUE(galks.has_key(1));
              ASSERT_FALSE(galks.has_key(3));
              ASSERT_FALSE(galks.has_key(127));
              ASSERT_EQ(1ULL, galks.key(1).size());
              ASSERT_EQ(1ULL, galks.size());
  
              keygen.create_galois_keys(vector<uint32_t>{ 127 }, galks);
              ASSERT_TRUE(galks.parms_id() == context.key_parms_id());
              ASSERT_FALSE(galks.has_key(1));
              ASSERT_TRUE(galks.has_key(127));
              ASSERT_EQ(1ULL, galks.key(127).size());
              ASSERT_EQ(1ULL, galks.size());
          }
          {
              parms.set_poly_modulus_degree(256);
              parms.set_coeff_modulus(CoeffModulus::Create(256, { 60, 30, 30 }));
  
              SEALContext context(parms, false, sec_level_type::none);
              KeyGenerator keygen(context);
  
              RelinKeys evk;
              keygen.create_relin_keys(evk);
              ASSERT_TRUE(evk.parms_id() == context.key_parms_id());
              ASSERT_EQ(2ULL, evk.key(2).size());
              for (auto &a : evk.data())
              {
                  for (auto &b : a)
                  {
                      ASSERT_FALSE(b.data().is_transparent());
                  }
              }
              ASSERT_TRUE(is_valid_for(evk, context));
  
              GaloisKeys galks;
              keygen.create_galois_keys(galks);
              for (auto &a : galks.data())
              {
                  for (auto &b : a)
                  {
                      ASSERT_FALSE(b.data().is_transparent());
                  }
              }
              ASSERT_TRUE(is_valid_for(galks, context));
  
              ASSERT_TRUE(galks.parms_id() == context.key_parms_id());
              ASSERT_EQ(2ULL, galks.key(3).size());
              ASSERT_EQ(14ULL, galks.size());
  
              keygen.create_galois_keys(vector<uint32_t>{ 1, 3, 5, 7 }, galks);
              ASSERT_TRUE(galks.parms_id() == context.key_parms_id());
              ASSERT_TRUE(galks.has_key(1));
              ASSERT_TRUE(galks.has_key(3));
              ASSERT_TRUE(galks.has_key(5));
              ASSERT_TRUE(galks.has_key(7));
              ASSERT_FALSE(galks.has_key(9));
              ASSERT_FALSE(galks.has_key(511));
              ASSERT_EQ(2ULL, galks.key(1).size());
              ASSERT_EQ(2ULL, galks.key(3).size());
              ASSERT_EQ(2ULL, galks.key(5).size());
              ASSERT_EQ(2ULL, galks.key(7).size());
              ASSERT_EQ(4ULL, galks.size());
  
              keygen.create_galois_keys(vector<uint32_t>{ 1 }, galks);
              ASSERT_TRUE(galks.parms_id() == context.key_parms_id());
              ASSERT_TRUE(galks.has_key(1));
              ASSERT_FALSE(galks.has_key(3));
              ASSERT_FALSE(galks.has_key(511));
              ASSERT_EQ(2ULL, galks.key(1).size());
              ASSERT_EQ(1ULL, galks.size());
  
              keygen.create_galois_keys(vector<uint32_t>{ 511 }, galks);
              ASSERT_TRUE(galks.parms_id() == context.key_parms_id());
              ASSERT_FALSE(galks.has_key(1));
              ASSERT_TRUE(galks.has_key(511));
              ASSERT_EQ(2ULL, galks.key(511).size());
              ASSERT_EQ(1ULL, galks.size());
          }
      }
  
      TEST(KeyGeneratorTest, Constructors)
      {
          auto constructors = [](scheme_type scheme) {
              EncryptionParameters parms(scheme);
              parms.set_poly_modulus_degree(128);
              parms.set_plain_modulus(65537);
              parms.set_coeff_modulus(CoeffModulus::Create(128, { 60, 50, 40 }));
              SEALContext context(parms, false, sec_level_type::none);
              Evaluator evaluator(context);
  
              KeyGenerator keygen(context);
              PublicKey pk;
              keygen.create_public_key(pk);
              auto sk = keygen.secret_key();
              RelinKeys rlk;
              keygen.create_relin_keys(rlk);
              GaloisKeys galk;
              keygen.create_galois_keys(galk);
  
              ASSERT_TRUE(is_valid_for(rlk, context));
              ASSERT_TRUE(is_valid_for(galk, context));
  
              Encryptor encryptor(context, pk);
              Decryptor decryptor(context, sk);
              Plaintext pt("1x^2 + 2"), ptres;
              Ciphertext ct;
              encryptor.encrypt(pt, ct);
              evaluator.square_inplace(ct);
              evaluator.relinearize_inplace(ct, rlk);
              decryptor.decrypt(ct, ptres);
              ASSERT_EQ("1x^4 + 4x^2 + 4", ptres.to_string());
  
              KeyGenerator keygen2(context, sk);
              auto sk2 = keygen.secret_key();
              PublicKey pk2;
              keygen2.create_public_key(pk2);
              ASSERT_EQ(sk2.data(), sk.data());
  
              RelinKeys rlk2;
              keygen2.create_relin_keys(rlk2);
              GaloisKeys galk2;
              keygen2.create_galois_keys(galk2);
  
              ASSERT_TRUE(is_valid_for(rlk2, context));
              ASSERT_TRUE(is_valid_for(galk2, context));
  
              Encryptor encryptor2(context, pk2);
              Decryptor decryptor2(context, sk2);
              pt = "1x^2 + 2";
              ptres.set_zero();
              encryptor.encrypt(pt, ct);
              evaluator.square_inplace(ct);
              evaluator.relinearize_inplace(ct, rlk2);
              decryptor.decrypt(ct, ptres);
              ASSERT_EQ("1x^4 + 4x^2 + 4", ptres.to_string());
  
              PublicKey pk3;
              keygen2.create_public_key(pk3);
  
              // There is a small random chance for this to fail
              for (size_t i = 0; i < pk3.data().dyn_array().size(); i++)
              {
                  ASSERT_NE(pk3.data().data()[i], pk2.data().data()[i]);
              }
          };
  
          constructors(scheme_type::bfv);
          constructors(scheme_type::bgv);
      }
  } // namespace sealtest
  