Revert "Create one map for all automorphism keys (openfheorg#566)" (o…

…penfheorg#609) This reverts commit 0122238. Co-authored-by: Dmitriy Suponitskiy <dsuponitskiy@dualitytech.com>
gong-cr · Nov 30, 2023 · 1f90f06 · 1f90f06
1 parent d8c6135
commit 1f90f06
Show file tree

Hide file tree

Showing 11 changed files with 526 additions and 343 deletions.
diff --git a/src/pke/include/cryptocontext.h b/src/pke/include/cryptocontext.h
diff --git a/src/pke/include/schemebase/base-advancedshe.h b/src/pke/include/schemebase/base-advancedshe.h
@@ -210,30 +210,33 @@ class AdvancedSHEBase {
    * @param privateKey private key.
    * @return returns the evaluation keys
    */
-    virtual std::shared_ptr<std::map<usint, EvalKey<Element>>> EvalSumKeyGen(
-        const PrivateKey<Element> privateKey) const;
+    virtual std::shared_ptr<std::map<usint, EvalKey<Element>>> EvalSumKeyGen(const PrivateKey<Element> privateKey,
+                                                                             const PublicKey<Element> publicKey) const;
 
     /**
    * Virtual function to generate the automorphism keys for EvalSumRows; works
    * only for packed encoding
    *
    * @param privateKey private key.
+   * @param publicKey public key.
    * @param rowSize size of rows in the matrix
    * @param subringDim subring dimension (set to cyclotomic order if set to 0)
    * @return returns the evaluation keys
    */
     virtual std::shared_ptr<std::map<usint, EvalKey<Element>>> EvalSumRowsKeyGen(const PrivateKey<Element> privateKey,
+                                                                                 const PublicKey<Element> publicKey,
                                                                                  usint rowSize, usint subringDim) const;
 
     /**
    * Virtual function to generate the automorphism keys for EvalSumCols; works
    * only for packed encoding
    *
    * @param privateKey private key.
+   * @param publicKey public key.
    * @return returns the evaluation keys
    */
     virtual std::shared_ptr<std::map<usint, EvalKey<Element>>> EvalSumColsKeyGen(
-        const PrivateKey<Element> privateKey) const;
+        const PrivateKey<Element> privateKey, const PublicKey<Element> publicKey) const;
 
     /**
    * Sums all elements in log (batch size) time - works only with packed

diff --git a/src/pke/include/schemebase/base-leveledshe.h b/src/pke/include/schemebase/base-leveledshe.h
@@ -549,6 +549,21 @@ class LeveledSHEBase {
     virtual std::shared_ptr<std::map<usint, EvalKey<Element>>> EvalAutomorphismKeyGen(
         const PrivateKey<Element> privateKey, const std::vector<usint>& indexList) const;
 
+    /**
+   * Virtual function to generate all isomorphism keys for a given private key
+   *
+   * @param publicKey encryption key for the new ciphertext.
+   * @param origPrivateKey original private key used for decryption.
+   * @param indexList list of automorphism indices to be computed
+   * @return returns the evaluation keys
+   */
+    virtual std::shared_ptr<std::map<usint, EvalKey<Element>>> EvalAutomorphismKeyGen(
+        const PublicKey<Element> publicKey, const PrivateKey<Element> privateKey,
+        const std::vector<usint>& indexList) const {
+        std::string errMsg = "EvalAutomorphismKeyGen is not implemented for this scheme.";
+        OPENFHE_THROW(not_implemented_error, errMsg);
+    }
+
     /**
    * Virtual function for evaluating automorphism of ciphertext at index i
    *
@@ -597,12 +612,14 @@ class LeveledSHEBase {
    * Generates evaluation keys for a list of indices
    * Currently works only for power-of-two and cyclic-group cyclotomics
    *
+   * @param publicKey encryption key for the new ciphertext.
    * @param origPrivateKey original private key used for decryption.
    * @param indexList list of indices to be computed
    * @return returns the evaluation keys
    */
     virtual std::shared_ptr<std::map<usint, EvalKey<Element>>> EvalAtIndexKeyGen(
-        const PrivateKey<Element> privateKey, const std::vector<int32_t>& indexList) const;
+        const PublicKey<Element> publicKey, const PrivateKey<Element> privateKey,
+        const std::vector<int32_t>& indexList) const;
 
     /**
    * Moves i-th slot to slot 0

diff --git a/src/pke/include/schemebase/base-scheme.h b/src/pke/include/schemebase/base-scheme.h
@@ -873,6 +873,10 @@ class SchemeBase {
     virtual std::shared_ptr<std::map<usint, EvalKey<Element>>> EvalAutomorphismKeyGen(
         const PrivateKey<Element> privateKey, const std::vector<usint>& indexList) const;
 
+    virtual std::shared_ptr<std::map<usint, EvalKey<Element>>> EvalAutomorphismKeyGen(
+        const PublicKey<Element> publicKey, const PrivateKey<Element> privateKey,
+        const std::vector<usint>& indexList) const;
+
     virtual Ciphertext<Element> EvalAutomorphism(ConstCiphertext<Element> ciphertext, usint i,
                                                  const std::map<usint, EvalKey<Element>>& evalKeyMap,
                                                  CALLER_INFO_ARGS_HDR) const {
@@ -946,7 +950,8 @@ class SchemeBase {
     }
 
     virtual std::shared_ptr<std::map<usint, EvalKey<Element>>> EvalAtIndexKeyGen(
-        const PrivateKey<Element> privateKey, const std::vector<int32_t>& indexList) const;
+        const PublicKey<Element> publicKey, const PrivateKey<Element> privateKey,
+        const std::vector<int32_t>& indexList) const;
 
     virtual Ciphertext<Element> EvalAtIndex(ConstCiphertext<Element> ciphertext, usint i,
                                             const std::map<usint, EvalKey<Element>>& evalKeyMap) const {
@@ -1178,14 +1183,15 @@ class SchemeBase {
     // Advanced SHE EVAL SUM
     /////////////////////////////////////
 
-    virtual std::shared_ptr<std::map<usint, EvalKey<Element>>> EvalSumKeyGen(
-        const PrivateKey<Element> privateKey) const;
+    virtual std::shared_ptr<std::map<usint, EvalKey<Element>>> EvalSumKeyGen(const PrivateKey<Element> privateKey,
+                                                                             const PublicKey<Element> publicKey) const;
 
     virtual std::shared_ptr<std::map<usint, EvalKey<Element>>> EvalSumRowsKeyGen(const PrivateKey<Element> privateKey,
+                                                                                 const PublicKey<Element> publicKey,
                                                                                  usint rowSize, usint subringDim) const;
 
     virtual std::shared_ptr<std::map<usint, EvalKey<Element>>> EvalSumColsKeyGen(
-        const PrivateKey<Element> privateKey) const;
+        const PrivateKey<Element> privateKey, const PublicKey<Element> publicKey) const;
 
     virtual Ciphertext<Element> EvalSum(ConstCiphertext<Element> ciphertext, usint batchSize,
                                         const std::map<usint, EvalKey<Element>>& evalKeyMap) const {