Moved the logic to deduplicate the indices in key generation operations for automorphism keys to a lower level

Author: dsuponitskiy-duality

src/pke/include/cryptocontext.h

@@ -230,21 +230,6 @@ class CryptoContextImpl : public Serializable {
                                                value2);
     }
 
-    /**
-    * @brief Gets indices that do not have automorphism keys for the given secret key tag in the key map
-    *
-    * @param keyTag secret key tag
-    * @param indexList array of specific indices to check the key map against
-    * @return indices that do not have automorphism keys associated with
-    */
-    static std::set<uint32_t> GetEvalAutomorphismNoKeyIndices(const std::string& keyTag,
-                                                              const std::set<uint32_t>& indices) {
-        std::set<uint32_t> existingIndices{CryptoContextImpl<Element>::GetExistingEvalAutomorphismKeyIndices(keyTag)};
-        // if no index found for the given keyTag, then the entire set "indices" is returned
-        return (existingIndices.empty()) ? indices :
-                                           CryptoContextImpl<Element>::GetUniqueValues(existingIndices, indices);
-    }
-
     /**
     * @brief Gets automorphism keys for a specific secret key tag and an array of specific indices
     *
@@ -2250,13 +2235,7 @@ class CryptoContextImpl : public Serializable {
         if (!indexList.size())
             OPENFHE_THROW("Input index vector is empty");
 
-        // Do not generate duplicate keys that have been already generated and added to the static storage (map)
-        std::set<uint32_t> allIndices(indexList.begin(), indexList.end());
-        std::set<uint32_t> indicesToGenerate{
-            CryptoContextImpl<Element>::GetEvalAutomorphismNoKeyIndices(privateKey->GetKeyTag(), allIndices)};
-
-        std::vector<uint32_t> newIndices(indicesToGenerate.begin(), indicesToGenerate.end());
-        auto evalKeys = GetScheme()->EvalAutomorphismKeyGen(privateKey, newIndices);
+        auto evalKeys = GetScheme()->EvalAutomorphismKeyGen(privateKey, indexList);
         CryptoContextImpl<Element>::InsertEvalAutomorphismKey(evalKeys, privateKey->GetKeyTag());
 
         return evalKeys;
@@ -4015,6 +3994,21 @@ class CryptoContextImpl : public Serializable {
         GetScheme()->SetSwkFC(FHEWtoCKKSswk);
     }
 
+    /**
+    * @brief Gets indices that do not have automorphism keys for the given secret key tag in the key map
+    *
+    * @param keyTag secret key tag
+    * @param indexList array of specific indices to check the key map against
+    * @return indices that do not have automorphism keys associated with
+    */
+    static std::set<uint32_t> GetEvalAutomorphismNoKeyIndices(const std::string& keyTag,
+                                                              const std::set<uint32_t>& indices) {
+        std::set<uint32_t> existingIndices{CryptoContextImpl<Element>::GetExistingEvalAutomorphismKeyIndices(keyTag)};
+        // if no index found for the given keyTag, then the entire set "indices" is returned
+        return (existingIndices.empty()) ? indices :
+                                           CryptoContextImpl<Element>::GetUniqueValues(existingIndices, indices);
+    }
+
     /**
     * @brief Returns automorphism indices for all existing evaluation keys.
     * @param keyTag      Secret key tag.

src/pke/lib/schemebase/base-leveledshe.cpp

@@ -347,9 +347,16 @@ void LeveledSHEBase<Element>::RelinearizeInPlace(Ciphertext<Element>& ciphertext
 template <class Element>
 std::shared_ptr<std::map<uint32_t, EvalKey<Element>>> LeveledSHEBase<Element>::EvalAutomorphismKeyGen(
     const PrivateKey<Element> privateKey, const std::vector<uint32_t>& indexList) const {
+
+    // Do not generate duplicate keys that have been already generated and added to the static storage (map)
+    std::set<uint32_t> allIndices(indexList.begin(), indexList.end());
+    std::set<uint32_t> indicesToGenerate{
+        CryptoContextImpl<Element>::GetEvalAutomorphismNoKeyIndices(privateKey->GetKeyTag(), allIndices)};
+    std::vector<uint32_t> newIndices(indicesToGenerate.begin(), indicesToGenerate.end());
+
     // we already have checks on higher level?
-    //  auto it = std::find(indexList.begin(), indexList.end(), 2 * n - 1);
-    //  if (it != indexList.end())
+    //  auto it = std::find(newIndices.begin(), newIndices.end(), 2 * n - 1);
+    //  if (it != newIndices.end())
     //    OPENFHE_THROW("conjugation is disabled");
 
     const auto cc = privateKey->GetCryptoContext();
@@ -359,26 +366,26 @@ std::shared_ptr<std::map<uint32_t, EvalKey<Element>>> LeveledSHEBase<Element>::E
     const uint32_t M = s.GetCyclotomicOrder();
 
     // we already have checks on higher level?
-    //  if (indexList.size() > N - 1)
+    //  if (newIndices.size() > N - 1)
     //    OPENFHE_THROW("size exceeds the ring dimension");
 
-    // create and initialize the key map (key is a value from indexList, EvalKey is nullptr). in this case
+    // create and initialize the key map (key is a value from newIndices, EvalKey is nullptr). in this case
     // we should be able to assign values to the map without using "omp critical" as all evalKeys' elements would
     // have already been created
     auto evalKeys = std::make_shared<std::map<uint32_t, EvalKey<Element>>>();
-    for (auto indx : indexList)
+    for (auto indx : newIndices)
         (*evalKeys)[indx];
 
-    const uint32_t sz = indexList.size();
+    const uint32_t sz = newIndices.size();
 #pragma omp parallel for
     for (uint32_t i = 0; i < sz; ++i) {
-        auto index = NativeInteger(indexList[i]).ModInverse(M).ConvertToInt<uint32_t>();
+        auto index = NativeInteger(newIndices[i]).ModInverse(M).ConvertToInt<uint32_t>();
         std::vector<uint32_t> vec(N);
         PrecomputeAutoMap(N, index, &vec);
 
         auto privateKeyPermuted = std::make_shared<PrivateKeyImpl<Element>>(cc);
         privateKeyPermuted->SetPrivateElement(s.AutomorphismTransform(index, vec));
-        (*evalKeys)[indexList[i]] = cc->GetScheme()->KeySwitchGen(privateKey, privateKeyPermuted);
+        (*evalKeys)[newIndices[i]] = cc->GetScheme()->KeySwitchGen(privateKey, privateKeyPermuted);
     }
 
     return evalKeys;