289 return a vector of rotation indices

examples/pke/simple-ckks-bootstrapping.py

@@ -52,6 +52,9 @@ def simple_bootstrap_example():
     cryptocontext.EvalMultKeyGen(key_pair.secretKey)
     cryptocontext.EvalBootstrapKeyGen(key_pair.secretKey, num_slots)
 
+    key_indices = cryptocontext.GetExistingEvalAutomorphismKeyIndices(key_pair.secretKey.GetKeyTag())
+    print(f"Number of bootstrapping automorphism keys: {len(key_indices)}")
+
     x = [0.25, 0.5, 0.75, 1.0, 2.0, 3.0, 4.0, 5.0]
     encoded_length = len(x)
 

src/include/docstrings/cryptocontext_docs.h

@@ -957,6 +957,19 @@ const char* cc_GetEvalAutomorphismKeyMap_docs = R"pbdoc(
     :rtype: EvalKeyMap
 )pbdoc";
 
+const char* cc_GetExistingEvalAutomorphismKeyIndices_docs = R"pbdoc(
+    Get the automorphism indices of all evaluation keys held for a secret key tag, including
+    the conjugation keys generated by EvalBootstrapKeyGen.
+
+    These are automorphism indices, not the slot offsets passed to EvalRotateKeyGen. Map a
+    signed offset with FindAutomorphismIndex() to look it up here.
+
+    :param keyTag: secret key identifier, obtained with secretKey.GetKeyTag()
+    :type keyTag: str
+    :return: sorted list of unique automorphism indices; empty if no keys exist for the tag
+    :rtype: list[int]
+)pbdoc";
+
 const char* cc_GetEvalSumKeyMap_docs = R"pbdoc(
     Get a map of summation keys (each is composed of several automorphism keys) for a specific secret key tag
     :return: EvalKeyMap: key map
@@ -1576,7 +1589,7 @@ const char* cc_EvalAutomorphismKeyGen_docs = R"pbdoc(
 const char* cc_FindAutomorphismIndex_docs = R"pbdoc(
     Finds an automorphism index for a given vector index using a scheme-specific algorithm
 
-    :param idx: regular vector index
+    :param idx: signed 32-bit rotation offset, as accepted by EvalRotateKeyGen
     :type idx: int
     :return: the automorphism index
     :rtype: int
@@ -1585,7 +1598,7 @@ const char* cc_FindAutomorphismIndex_docs = R"pbdoc(
 const char* cc_FindAutomorphismIndices_docs = R"pbdoc(
     Finds automorphism indices for a given list of vector indices using a scheme-specific algorithm
 
-    :param idxList: list of indices
+    :param idxList: list of signed 32-bit rotation offsets, as accepted by EvalRotateKeyGen
     :type idxList: List[int]
     :return: a list of automorphism indices
     :rtype: List[int]

src/lib/bindings.cpp

@@ -1110,10 +1110,15 @@ void bind_crypto_context(py::module &m) {
             py::arg("ciphertextVec"))
         .def("EvalAddManyInPlace", &CryptoContextImpl<DCRTPoly>::EvalAddManyInPlace,
             py::arg("ciphertextVec"))
-        .def("FindAutomorphismIndex", &CryptoContextImpl<DCRTPoly>::FindAutomorphismIndex,
+        .def("FindAutomorphismIndex", [](const CryptoContextImpl<DCRTPoly>& self, int32_t idx) {
+                return self.FindAutomorphismIndex(static_cast<uint32_t>(idx));
+            },
             py::arg("idx"),
             py::doc(cc_FindAutomorphismIndex_docs))
-        .def("FindAutomorphismIndices", &CryptoContextImpl<DCRTPoly>::FindAutomorphismIndices,
+        .def("FindAutomorphismIndices", [](const CryptoContextImpl<DCRTPoly>& self, const std::vector<int32_t>& idxList) {
+                // The C++ API carries signed rotation offsets in uint32_t values.
+                return self.FindAutomorphismIndices(std::vector<uint32_t>(idxList.begin(), idxList.end()));
+            },
             py::arg("idxList"),
             py::doc(cc_FindAutomorphismIndices_docs))
         .def("GetEvalSumKeyMap",
@@ -1148,6 +1153,12 @@ void bind_crypto_context(py::module &m) {
         .def_static("GetEvalAutomorphismKeyMap", &CryptoContextImpl<DCRTPoly>::GetEvalAutomorphismKeyMapPtr,
             py::arg("keyTag") = "",
             py::doc(cc_GetEvalAutomorphismKeyMap_docs))
+        .def_static("GetExistingEvalAutomorphismKeyIndices", [](const std::string& keyTag) {
+                const auto indices = CryptoContextImpl<DCRTPoly>::GetExistingEvalAutomorphismKeyIndices(keyTag);
+                return std::vector<uint32_t>(indices.begin(), indices.end());
+            },
+            py::arg("keyTag") = "",
+            py::doc(cc_GetExistingEvalAutomorphismKeyIndices_docs))
         .def_static("SerializeEvalMultKey", [](const std::string &filename, const SerType::SERBINARY &sertype, std::string keyTag = "") {
                 std::ofstream outfile(filename, std::ios::out | std::ios::binary);
                 bool res = CryptoContextImpl<DCRTPoly>::SerializeEvalMultKey<SerType::SERBINARY>(outfile, sertype, keyTag);

tests/test_ckks.py

@@ -55,3 +55,76 @@ def test_add_two_numbers(ckks_context):
     raw_added = [a + b for (a, b) in zip(*raw)]
     total = sum(abs(a - b) for (a, b) in zip(raw_added, final_added))
     assert total < 1e-3
+
+
+def test_existing_eval_automorphism_key_indices(ckks_context):
+    _, cc, existing_keys = ckks_context
+    keys = cc.KeyGen()
+    key_tag = keys.secretKey.GetKeyTag()
+    assert cc.GetExistingEvalAutomorphismKeyIndices(keyTag=key_tag) == []
+
+    cc.EvalRotateKeyGen(keys.secretKey, [1, -2])
+    # CKKS rotation offsets map to powers of 5 modulo the cyclotomic order.
+    cyclotomic_order = 2 * cc.GetRingDimension()
+    expected = sorted(pow(5, offset, cyclotomic_order) for offset in [1, -2])
+    indices = cc.GetExistingEvalAutomorphismKeyIndices(key_tag)
+    assert isinstance(indices, list)
+    assert indices == expected
+    assert fhe.CryptoContext.GetExistingEvalAutomorphismKeyIndices(key_tag) == expected
+    # The documented way to go from a slot offset to an entry in this list. This also
+    # cross-checks the 5^offset formula above for both rotation directions.
+    for offset in [1, -2]:
+        assert cc.FindAutomorphismIndex(offset) == pow(5, offset, cyclotomic_order)
+        assert cc.FindAutomorphismIndex(offset) in indices
+    # keyTag defaults to "", which matches no key map.
+    assert cc.GetExistingEvalAutomorphismKeyIndices() == []
+
+    # Repeated key generation must not duplicate indices or affect another tag.
+    cc.EvalRotateKeyGen(keys.secretKey, [1, 3])
+    updated = sorted(expected + [pow(5, 3, cyclotomic_order)])
+    assert cc.GetExistingEvalAutomorphismKeyIndices(key_tag) == updated
+    assert cc.GetExistingEvalAutomorphismKeyIndices(existing_keys.secretKey.GetKeyTag()) == expected
+    assert indices == expected
+
+
+def test_find_automorphism_indices_signed_offsets(ckks_context):
+    _, cc, _ = ckks_context
+    cyclotomic_order = 2 * cc.GetRingDimension()
+    offsets = [0, 1, -2, 3, -1, -2]
+    expected = [pow(5, offset, cyclotomic_order) for offset in offsets]
+    assert cc.FindAutomorphismIndices(idxList=offsets) == expected
+    assert [cc.FindAutomorphismIndex(idx=offset) for offset in offsets] == expected
+    assert cc.FindAutomorphismIndices([]) == []
+    for offset in [-(2**31) - 1, 2**31]:
+        with pytest.raises(TypeError):
+            cc.FindAutomorphismIndex(offset)
+        with pytest.raises(TypeError):
+            cc.FindAutomorphismIndices([1, offset])
+
+
+def test_existing_bootstrap_key_indices():
+    parameters = fhe.CCParamsCKKSRNS()
+    parameters.SetSecurityLevel(fhe.HEStd_NotSet)
+    parameters.SetRingDim(512)
+    parameters.SetSecretKeyDist(fhe.UNIFORM_TERNARY)
+    level_budget = [2, 2]
+    depth = fhe.FHECKKSRNS.GetBootstrapDepth(level_budget, fhe.UNIFORM_TERNARY)
+    parameters.SetMultiplicativeDepth(depth + 2)
+    parameters.SetScalingTechnique(fhe.FIXEDAUTO)
+    parameters.SetScalingModSize(78 if fhe.get_native_int() == 128 else 59)
+    parameters.SetFirstModSize(89 if fhe.get_native_int() == 128 else 60)
+    cc = fhe.GenCryptoContext(parameters)
+    for feature in [fhe.PKE, fhe.KEYSWITCH, fhe.LEVELEDSHE, fhe.ADVANCEDSHE, fhe.FHE]:
+        cc.Enable(feature)
+
+    slots = 8
+    cc.EvalBootstrapSetup(level_budget, [0, 0], slots)
+    keys = cc.KeyGen()
+    key_tag = keys.secretKey.GetKeyTag()
+    assert cc.GetExistingEvalAutomorphismKeyIndices(key_tag) == []
+    cc.EvalBootstrapKeyGen(keys.secretKey, slots)
+    indices = cc.GetExistingEvalAutomorphismKeyIndices(key_tag)
+    assert isinstance(indices, list)
+    assert indices == sorted(set(indices))
+    assert len(indices) > 1
+    assert 2 * cc.GetRingDimension() - 1 in indices  # Conjugation key.