From 9b9feeb180f9ba776018f50fc3f487417d8cfa23 Mon Sep 17 00:00:00 2001
From: Michael Schlottke-Lakemper <michael@sloede.com>
Date: Tue, 16 Jan 2024 13:36:12 +0100
Subject: [PATCH] Add `GetCryptoContext`

---
 src/OpenFHE.jl     |  3 +++
 src/convenience.jl | 19 +++++++++++++++++++
 test/test_ckks.jl  |  4 ++++
 3 files changed, 26 insertions(+)

diff --git a/src/OpenFHE.jl b/src/OpenFHE.jl
index 950ba84..c861006 100644
--- a/src/OpenFHE.jl
+++ b/src/OpenFHE.jl
@@ -31,6 +31,9 @@ export SetMultiplicativeDepth, SetScalingModSize, SetBatchSize, SetSecretKeyDist
 # FHECKKSRNS
 export GetBootstrapDepth
 
+# CryptoObject
+export GetCryptoContext
+
 # CryptoContext
 export CryptoContext, GenCryptoContext, Enable, GetRingDimension, KeyGen, EvalMultKeyGen,
        EvalRotateKeyGen, MakeCKKSPackedPlaintext, Encrypt, EvalAdd, EvalSub, EvalMult,
diff --git a/src/convenience.jl b/src/convenience.jl
index b183a87..183979d 100644
--- a/src/convenience.jl
+++ b/src/convenience.jl
@@ -204,6 +204,25 @@ function EvalBootstrap(context::CxxWrap.CxxWrapCore.CxxRef{OpenFHE.CryptoContext
     EvalBootstrap(context, ciphertext, num_iterations, precision)
 end
 
+"""
+    GetCryptoContext(object::Union{Ciphertext})
+
+Return a the crypto context for a an `object` that is is a subtype of `CryptoObject`.
+
+Currently, this is only implemented for [`Ciphertext`](@ref).
+
+See also: [`CryptoContext`](@ref), [`Ciphertext`](@ref)
+"""
+function GetCryptoContext(object::Union{Ciphertext})
+    # Note: Due to limitations of CxxWrap.jl when wrapping mutually dependent template
+    # types, openfhe-julia introduces a non-template proxy type as an additional level of
+    # indirection.
+    # We thus need to first dereference `Ciphertext` to get to the `CiphertextImpl`, then
+    # call the proxy method to get the `CryptoContextProxy`, and finally obtain the actual
+    # `CryptoContext` from it
+    GetCryptoContext(GetCryptoContextProxy(object[]))
+end
+
 """
     Decrypt(crypto_context::CryptoContext, ciphertext::Ciphertext, private_key::PrivateKey, plaintext::Plaintext)
     Decrypt(crypto_context::CryptoContext, private_key::PrivateKey, ciphertext::Ciphertext, plaintext::Plaintext)
diff --git a/test/test_ckks.jl b/test/test_ckks.jl
index 4199e93..e054604 100644
--- a/test/test_ckks.jl
+++ b/test/test_ckks.jl
@@ -135,6 +135,10 @@ end
     @test GetRealPackedValue(result_dec) ≈ circshift(x1, -shift)
 end
 
+@testset verbose=true showtiming=true "GetCryptoContext" begin
+    @test GetCryptoContext(c1) isa CryptoContext
+end
+
 end # @testset "test_ckks.jl"
 
 end # module