feat: Add Permute

sp301415 · Aug 30, 2024 · 9689cce · 9689cce
1 parent 34da010
commit 9689cce
Show file tree

Hide file tree

Showing 3 changed files with 102 additions and 2 deletions.
diff --git a/math/poly/poly_evaluator.go b/math/poly/poly_evaluator.go
@@ -37,7 +37,7 @@ type Evaluator[T num.Integer] struct {
 
 // evaluationBuffer contains buffer values for Evaluator.
 type evaluationBuffer[T num.Integer] struct {
-	// pOut holds the intermediate polynomial in MulAdd or MulSub type operations.
+	// pOut holds the intermediate polynomial in MulAdd/MulSub and Permute type operations.
 	pOut Poly[T]
 }
 

diff --git a/math/poly/poly_ops.go b/math/poly/poly_ops.go
@@ -234,3 +234,60 @@ func (e *Evaluator[T]) MonomialMulSubAssign(p0 Poly[T], d int, pOut Poly[T]) {
 		}
 	}
 }
+
+// Permute returns p0(X^d).
+func (e *Evaluator[T]) Permute(p0 Poly[T], d int) Poly[T] {
+	pOut := e.NewPoly()
+	e.PermuteAssign(p0, d, pOut)
+	return pOut
+}
+
+// PermuteAssign computes pOut = p0(X^d).
+//
+// p0 and pOut should not overlap. For inplace permutation,
+// use [*Evaluator.PermuteInPlace].
+func (e *Evaluator[T]) PermuteAssign(p0 Poly[T], d int, pOut Poly[T]) {
+	pOut.Clear()
+	for i := 0; i < e.degree; i++ {
+		j := (d * i) & (2*e.degree - 1)
+		if j < e.degree {
+			pOut.Coeffs[j] += p0.Coeffs[i]
+		} else {
+			pOut.Coeffs[j-e.degree] -= p0.Coeffs[i]
+		}
+	}
+}
+
+// PermuteInPlace computes p0 = p0(X^d).
+func (e *Evaluator[T]) PermuteInPlace(p0 Poly[T], d int) {
+	e.PermuteAssign(p0, d, e.buffer.pOut)
+	p0.CopyFrom(e.buffer.pOut)
+}
+
+// PermuteAddAssign computes pOut += p0(X^d).
+//
+// p0 and pOut should not overlap.
+func (e *Evaluator[T]) PermuteAddAssign(p0 Poly[T], d int, pOut Poly[T]) {
+	for i := 0; i < e.degree; i++ {
+		j := (d * i) & (2*e.degree - 1)
+		if j < e.degree {
+			pOut.Coeffs[j] += p0.Coeffs[i]
+		} else {
+			pOut.Coeffs[j-e.degree] -= p0.Coeffs[i]
+		}
+	}
+}
+
+// PermuteSubAssign computes pOut -= p0(X^d).
+//
+// p0 and pOut should not overlap.
+func (e *Evaluator[T]) PermuteSubAssign(p0 Poly[T], d int, pOut Poly[T]) {
+	for i := 0; i < e.degree; i++ {
+		j := (d * i) & (2*e.degree - 1)
+		if j < e.degree {
+			pOut.Coeffs[j] -= p0.Coeffs[i]
+		} else {
+			pOut.Coeffs[j-e.degree] += p0.Coeffs[i]
+		}
+	}
+}
diff --git a/tfhe/glwe_ops.go b/tfhe/glwe_ops.go
@@ -141,7 +141,8 @@ func (e *Evaluator[T]) MonomialMulGLWE(ct0 GLWECiphertext[T], d int) GLWECiphert
 
 // MonomialMulGLWEAssign computes ctOut = X^d * ct0.
 //
-// ct0 and ctOut should not overlap.
+// ct0 and ctOut should not overlap. For inplace multiplication,
+// use [*Evaluator.MonomialMulGLWEInPlace].
 func (e *Evaluator[T]) MonomialMulGLWEAssign(ct0 GLWECiphertext[T], d int, ctOut GLWECiphertext[T]) {
 	for i := 0; i < e.Parameters.glweRank+1; i++ {
 		e.PolyEvaluator.MonomialMulAssign(ct0.Value[i], d, ctOut.Value[i])
@@ -172,3 +173,45 @@ func (e *Evaluator[T]) MonomialMulSubGLWEAssign(ct0 GLWECiphertext[T], d int, ct
 		e.PolyEvaluator.MonomialMulSubAssign(ct0.Value[i], d, ctOut.Value[i])
 	}
 }
+
+// PermuteGLWE returns ctOut = ct0(X^d).
+func (e *Evaluator[T]) PermuteGLWE(ct0 GLWECiphertext[T], d int) GLWECiphertext[T] {
+	ctOut := NewGLWECiphertext(e.Parameters)
+	e.PermuteGLWEAssign(ct0, d, ctOut)
+	return ctOut
+}
+
+// PermuteGLWEAssign computes ctOut = ct0(X^d).
+//
+// ct0 and ctOut should not overlap. For inplace permutation,
+// use [*Evaluator.PermuteGLWEInPlace].
+func (e *Evaluator[T]) PermuteGLWEAssign(ct0 GLWECiphertext[T], d int, ctOut GLWECiphertext[T]) {
+	for i := 0; i < e.Parameters.glweRank+1; i++ {
+		e.PolyEvaluator.PermuteAssign(ct0.Value[i], d, ctOut.Value[i])
+	}
+}
+
+// PermuteGLWEInPlace computes ct0 = ct0(X^d).
+func (e *Evaluator[T]) PermuteGLWEInPlace(ct0 GLWECiphertext[T], d int) {
+	for i := 0; i < e.Parameters.glweRank+1; i++ {
+		e.PolyEvaluator.PermuteInPlace(ct0.Value[i], d)
+	}
+}
+
+// PermuteAddGLWEAssign computes ctOut += ct0(X^d).
+//
+// ct0 and ctOut should not overlap.
+func (e *Evaluator[T]) PermuteAddGLWEAssign(ct0 GLWECiphertext[T], d int, ctOut GLWECiphertext[T]) {
+	for i := 0; i < e.Parameters.glweRank+1; i++ {
+		e.PolyEvaluator.PermuteAddAssign(ct0.Value[i], d, ctOut.Value[i])
+	}
+}
+
+// PermuteSubGLWEAssign computes ctOut -= ct0(X^d).
+//
+// ct0 and ctOut should not overlap.
+func (e *Evaluator[T]) PermuteSubGLWEAssign(ct0 GLWECiphertext[T], d int, ctOut GLWECiphertext[T]) {
+	for i := 0; i < e.Parameters.glweRank+1; i++ {
+		e.PolyEvaluator.PermuteSubAssign(ct0.Value[i], d, ctOut.Value[i])
+	}
+}