From 4f30bf538d7e3ffb63b6190be8318f7d01c923df Mon Sep 17 00:00:00 2001
From: Christian Grigis <christian.grigis@epfl.ch>
Date: Thu, 28 Nov 2019 21:30:26 +0100
Subject: [PATCH] Add ckgProtocolContext for dbfv

---
 dbfv/dbfv_benchmark_test.go |  2 +-
 dbfv/dbfv_test.go           |  4 ++--
 dbfv/publickey_gen.go       | 33 ++++++++++++++++++++++++++++++---
 examples/dbfv/pir/pir.go    |  2 +-
 examples/dbfv/psi/psi.go    |  2 +-
 5 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/dbfv/dbfv_benchmark_test.go b/dbfv/dbfv_benchmark_test.go
index 26894d3c..ff8c3e1a 100644
--- a/dbfv/dbfv_benchmark_test.go
+++ b/dbfv/dbfv_benchmark_test.go
@@ -39,7 +39,7 @@ func benchPublicKeyGen(b *testing.B) {
 		}
 
 		p := new(Party)
-		p.CKGProtocol = NewCKGProtocol(bfvContext)
+		p.CKGProtocol = NewCKGProtocol(&parameters)
 		p.s = sk0Shards[0].Get()
 		p.s1 = p.AllocateShares()
 
diff --git a/dbfv/dbfv_test.go b/dbfv/dbfv_test.go
index c4816955..01001673 100644
--- a/dbfv/dbfv_test.go
+++ b/dbfv/dbfv_test.go
@@ -136,7 +136,7 @@ func testPublicKeyGen(t *testing.T) {
 			ckgParties := make([]*Party, parties)
 			for i := uint64(0); i < parties; i++ {
 				p := new(Party)
-				p.CKGProtocol = NewCKGProtocol(bfvContext)
+				p.CKGProtocol = NewCKGProtocol(&parameters)
 				p.s = sk0Shards[i].Get()
 				p.s1 = p.AllocateShares()
 				ckgParties[i] = p
@@ -732,7 +732,7 @@ func Test_Marshalling(t *testing.T) {
 	Ciphertext := bfv.NewRandomCiphertext(1, ringCtx)
 
 	t.Run(fmt.Sprintf("CPK/N=%d/limbQ=%d/limbsP=%d", contextQ.N, len(contextQ.Modulus), len(contextPKeys.Modulus)), func(t *testing.T) {
-		keygenProtocol := NewCKGProtocol(bfvCtx)
+		keygenProtocol := NewCKGProtocol(params)
 		KeyGenShareBefore := keygenProtocol.AllocateShares()
 		keygenProtocol.GenShare(sk.Get(), crs, KeyGenShareBefore)
 		//now we marshall it
diff --git a/dbfv/publickey_gen.go b/dbfv/publickey_gen.go
index 48c664a4..a557dccd 100644
--- a/dbfv/publickey_gen.go
+++ b/dbfv/publickey_gen.go
@@ -6,6 +6,31 @@ import (
 	"github.com/ldsec/lattigo/ring"
 )
 
+type ckgProtocolContext struct {
+	// Ternary and Gaussian samplers
+	gaussianSampler *ring.KYSampler
+
+	contextKeys *ring.Context
+}
+
+func newCkgProtocolContext(params *bfv.Parameters) *ckgProtocolContext {
+	n := params.N
+
+	contextKeys := ring.NewContext()
+	contextKeys.SetParameters(n, append(params.Qi, params.KeySwitchPrimes...))
+	err := contextKeys.GenNTTParams()
+	if err != nil {
+		panic(err)
+	}
+
+	gaussianSampler := contextKeys.NewKYSampler(params.Sigma, int(6*params.Sigma))
+
+	return &ckgProtocolContext{
+		contextKeys:     contextKeys,
+		gaussianSampler: gaussianSampler,
+	}
+}
+
 // CKGProtocol is the structure storing the parameters and state for a party in the collective key generation protocol.
 type CKGProtocol struct {
 	context         *ring.Context
@@ -28,10 +53,12 @@ func (share *CKGShare) UnmarshalBinary(data []byte) error {
 }
 
 // NewCKGProtocol creates a new CKGProtocol instance
-func NewCKGProtocol(bfvCtx *bfv.Context) *CKGProtocol {
+func NewCKGProtocol(params *bfv.Parameters) *CKGProtocol {
+	context := newCkgProtocolContext(params)
+
 	ckg := new(CKGProtocol)
-	ckg.context = bfvCtx.ContextKeys()
-	ckg.gaussianSampler = bfvCtx.GaussianSampler()
+	ckg.context = context.contextKeys
+	ckg.gaussianSampler = context.gaussianSampler
 	return ckg
 }
 
diff --git a/examples/dbfv/pir/pir.go b/examples/dbfv/pir/pir.go
index ad9a250c..f023eb59 100644
--- a/examples/dbfv/pir/pir.go
+++ b/examples/dbfv/pir/pir.go
@@ -102,7 +102,7 @@ func main() {
 	colSk.Set(bfvctx.ContextKeys().NewPoly())
 
 	// Instantiation of each of the protocols needed for the pir example
-	ckg := dbfv.NewCKGProtocol(bfvctx)       // public key generation
+	ckg := dbfv.NewCKGProtocol(params)       // public key generation
 	rkg := dbfv.NewEkgProtocol(params)       // relineariation key generation
 	rtg := dbfv.NewRotKGProtocol(params)     // rotation keys generation
 	cks := dbfv.NewCKSProtocol(params, 3.19) // collective public-key re-encryption
diff --git a/examples/dbfv/psi/psi.go b/examples/dbfv/psi/psi.go
index 14269928..a2a41279 100644
--- a/examples/dbfv/psi/psi.go
+++ b/examples/dbfv/psi/psi.go
@@ -81,7 +81,7 @@ func main() {
 		expRes[i] = 1
 	}
 
-	ckg := dbfv.NewCKGProtocol(bfvctx)
+	ckg := dbfv.NewCKGProtocol(params)
 	rkg := dbfv.NewEkgProtocol(params)
 	pcks := dbfv.NewPCKSProtocol(params, 3.19)