[ckks] : fixed bug in MultiplyByDiagMatrixBSGS that was modifying inp…

…uts (#159)
tuneinsight · Nov 25, 2021 · c761a16 · c761a16
1 parent bd2ac2f
commit c761a16
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 20 deletions.
diff --git a/ckks/ckks_test.go b/ckks/ckks_test.go
@@ -1117,7 +1117,7 @@ func testLinearTransform(testContext *testParams, t *testing.T) {
 
 		eval := testContext.evaluator.WithKey(rlwe.EvaluationKey{Rlk: testContext.rlk, Rtks: rotKey})
 
-		res := eval.LinearTransformNew(ciphertext1, ptDiagMatrix)[0]
+		eval.LinearTransform(ciphertext1, ptDiagMatrix, []*Ciphertext{ciphertext1})
 
 		tmp := make([]complex128, params.Slots())
 		copy(tmp, values1)
@@ -1133,7 +1133,7 @@ func testLinearTransform(testContext *testParams, t *testing.T) {
 			values1[i] += tmp[(i+15)%params.Slots()]
 		}
 
-		verifyTestVectors(testContext.params, testContext.encoder, testContext.decryptor, values1, res, testContext.params.LogSlots(), 0, t)
+		verifyTestVectors(testContext.params, testContext.encoder, testContext.decryptor, values1, ciphertext1, testContext.params.LogSlots(), 0, t)
 	})
 
 	t.Run(GetTestName(testContext.params, "LinearTransform/Naive/"), func(t *testing.T) {
@@ -1160,7 +1160,7 @@ func testLinearTransform(testContext *testParams, t *testing.T) {
 
 		eval := testContext.evaluator.WithKey(rlwe.EvaluationKey{Rlk: testContext.rlk, Rtks: rotKey})
 
-		res := eval.LinearTransformNew(ciphertext1, ptDiagMatrix)[0]
+		eval.LinearTransform(ciphertext1, ptDiagMatrix, []*Ciphertext{ciphertext1})
 
 		tmp := make([]complex128, params.Slots())
 		copy(tmp, values1)
@@ -1169,7 +1169,7 @@ func testLinearTransform(testContext *testParams, t *testing.T) {
 			values1[i] += tmp[(i-1+params.Slots())%params.Slots()]
 		}
 
-		verifyTestVectors(testContext.params, testContext.encoder, testContext.decryptor, values1, res, testContext.params.LogSlots(), 0, t)
+		verifyTestVectors(testContext.params, testContext.encoder, testContext.decryptor, values1, ciphertext1, testContext.params.LogSlots(), 0, t)
 	})
 }
 

diff --git a/ckks/linear_transform.go b/ckks/linear_transform.go
@@ -395,14 +395,9 @@ func (eval *evaluator) MultiplyByDiagMatrix(ctIn *Ciphertext, matrix PtDiagMatri
 	ksRes0QP := eval.Pool[3]
 	ksRes1QP := eval.Pool[4]
 
-	var ctInTmp0, ctInTmp1 *ring.Poly
-	if ctIn != ctOut {
-		ring.CopyValuesLvl(levelQ, ctIn.Value[0], eval.ctxpool.Value[0])
-		ring.CopyValuesLvl(levelQ, ctIn.Value[1], eval.ctxpool.Value[1])
-		ctInTmp0, ctInTmp1 = eval.ctxpool.Value[0], eval.ctxpool.Value[1]
-	} else {
-		ctInTmp0, ctInTmp1 = ctIn.Value[0], ctIn.Value[1]
-	}
+	ring.CopyValuesLvl(levelQ, ctIn.Value[0], eval.ctxpool.Value[0])
+	ring.CopyValuesLvl(levelQ, ctIn.Value[1], eval.ctxpool.Value[1])
+	ctInTmp0, ctInTmp1 := eval.ctxpool.Value[0], eval.ctxpool.Value[1]
 
 	ringQ.MulScalarBigintLvl(levelQ, ctInTmp0, ringP.ModulusBigint, ct0TimesP) // P*c0
 
@@ -499,14 +494,9 @@ func (eval *evaluator) MultiplyByDiagMatrixBSGS(ctIn *Ciphertext, matrix PtDiagM
 
 	index, rotations := bsgsIndex(matrix.Vec, 1<<matrix.LogSlots, matrix.N1)
 
-	var ctInTmp0, ctInTmp1 *ring.Poly
-	if ctIn == ctOut {
-		ring.CopyValuesLvl(levelQ, ctIn.Value[0], eval.ctxpool.Value[0])
-		ring.CopyValuesLvl(levelQ, ctIn.Value[1], eval.ctxpool.Value[1])
-		ctInTmp0, ctInTmp1 = eval.ctxpool.Value[0], eval.ctxpool.Value[1]
-	} else {
-		ctInTmp0, ctInTmp1 = ctIn.Value[0], ctIn.Value[1]
-	}
+	ring.CopyValuesLvl(levelQ, ctIn.Value[0], eval.ctxpool.Value[0])
+	ring.CopyValuesLvl(levelQ, ctIn.Value[1], eval.ctxpool.Value[1])
+	ctInTmp0, ctInTmp1 := eval.ctxpool.Value[0], eval.ctxpool.Value[1]
 
 	// Pre-rotates ciphertext for the baby-step giant-step algorithm, does not divide by P yet
 	ctInRotQP := eval.RotateHoistedNoModDownNew(levelQ, rotations, ctInTmp0, eval.PoolDecompQP)