Fboemer/more ciphertext async apis (#255)

* Async ciphertext APIs for rotation & mod-switching
* Add ciphertext async APIs for inner product & sum

Author: fboemer

Package.resolved

@@ -102,6 +102,7 @@ let package = Package(
         .target(
             name: "HomomorphicEncryption",
             dependencies: [
+                .product(name: "AsyncAlgorithms", package: "swift-async-algorithms"),
                 .product(name: "Crypto", package: "swift-crypto"),
                 .product(name: "_CryptoExtras", package: "swift-crypto"),
                 "CUtil",

Package.swift

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+public import AsyncAlgorithms
+
 /// Ciphertext type.
 public struct Ciphertext<Scheme: HeScheme, Format: PolyFormat>: Equatable, Sendable {
     public typealias Scalar = Scheme.Scalar
@@ -28,7 +30,7 @@ public struct Ciphertext<Scheme: HeScheme, Format: PolyFormat>: Equatable, Senda
     ///
     /// After a fresh encryption, the ciphertext has ``HeScheme/freshCiphertextPolyCount`` polynomials.
     /// The count may change during the course of HE operations, e.g. increase during ciphertext multiplication,
-    /// or decrease during relinearization ``Ciphertext/relinearize(using:)``.
+    /// or decrease during relinearization ``Ciphertext/relinearize(using:)-41bsm``.
     public var polyCount: Int {
         polys.count
     }
@@ -314,7 +316,7 @@ public struct Ciphertext<Scheme: HeScheme, Format: PolyFormat>: Equatable, Senda
     ///
     /// If the ciphertext already has a single modulus, this is a no-op.
     /// - Throws: Error upon failure to modulus switch.
-    /// - seealso: ``Ciphertext/modSwitchDown()`` for more information and an alternative API.
+    /// - seealso: ``Ciphertext/modSwitchDown()-4an2b`` for more information and an alternative API.
     @inlinable
     public mutating func modSwitchDownToSingle() throws where Format == Scheme.CanonicalCiphertextFormat {
         try Scheme.modSwitchDownToSingle(&self)
@@ -535,22 +537,31 @@ extension Ciphertext where Format == Scheme.CanonicalCiphertextFormat {
 }
 
 extension Collection {
+    /// Sums together the ciphertexts in the collection.
+    /// - Throws: Precondition failure if the collection is empty.
+    /// - Returns: The sum.
     @inlinable
-    func sum<Scheme>() throws -> Element where Element == Ciphertext<Scheme, Eval> {
+    public func sum<Scheme>() throws -> Element where Element == Ciphertext<Scheme, Eval> {
         precondition(!isEmpty)
         // swiftlint:disable:next force_unwrapping
         return try dropFirst().reduce(first!) { try $0 + $1 }
     }
 
+    /// Sums together the ciphertexts in the collection.
+    /// - Throws: Precondition failure if the collection is empty.
+    /// - Returns: The sum.
     @inlinable
-    func sum<Scheme>() throws -> Element where Element == Ciphertext<Scheme, Coeff> {
+    public func sum<Scheme>() throws -> Element where Element == Ciphertext<Scheme, Coeff> {
         precondition(!isEmpty)
         // swiftlint:disable:next force_unwrapping
         return try dropFirst().reduce(first!) { try $0 + $1 }
     }
 
+    /// Sums together the ciphertexts in the collection.
+    /// - Throws: Precondition failure if the collection is empty.
+    /// - Returns: The sum.
     @inlinable
-    func sum<Scheme>() throws -> Element where Element == Ciphertext<Scheme, Scheme.CanonicalCiphertextFormat> {
+    public func sum<Scheme>() throws -> Element where Element == Ciphertext<Scheme, Scheme.CanonicalCiphertextFormat> {
         precondition(!isEmpty)
         // swiftlint:disable:next force_unwrapping
         return try dropFirst().reduce(first!) { try $0 + $1 }
@@ -937,4 +948,167 @@ extension Ciphertext {
         }
         throw HeError.errorCastingPolyFormat(from: Format.self, to: Scheme.CanonicalCiphertextFormat.self)
     }
+
+    // MARK: Async rotations
+
+    /// Asynchronously rotates the columns of a ciphertext.
+    ///
+    /// - Parameters:
+    ///   - step: Number of slots to rotate. Negative values indicate a left rotation, and positive values indicate a
+    /// right rotation. Must have absolute value in `[1, N / 2 - 1]` where `N` is the RLWE ring dimension, given by
+    /// ``EncryptionParameters/polyDegree``.
+    ///   - evaluationKey: Evaluation key to use in the HE computation. Must contain the Galois element associated with
+    /// `step`, see ``GaloisElement/rotatingColumns(by:degree:)``.
+    /// - Throws: failure to rotate ciphertext's columns.
+    /// - seealso: ``HeScheme/rotateColumns(of:by:using:)-7h3fz`` for an alternate API and more information.
+    @inlinable
+    public mutating func rotateColumns(by step: Int,
+                                       using evaluationKey: EvaluationKey<Scheme>) async throws
+        where Format == Scheme.CanonicalCiphertextFormat
+    {
+        try await Scheme.rotateColumnsAsync(of: &self, by: step, using: evaluationKey)
+    }
+
+    /// Asynchronously swaps the rows of a ciphertext.
+    ///
+    /// A plaintext in ``EncodeFormat/simd`` format can be viewed a `2 x (N / 2)` matrix of coefficients.
+    /// For instance, for `N = 8`, given a ciphertext encrypting a plaintext with values
+    /// ```
+    /// [1, 2, 3, 4, 5, 6, 7, 8]
+    /// ```
+    /// calling ``HeScheme/swapRows(of:using:)`` with `step: 1` will yield a ciphertext decrypting to
+    /// ```
+    /// [5, 6, 7, 8, 1, 2, 3, 4]
+    /// ```
+    /// - Parameter evaluationKey: Evaluation key to use in the HE computation. Must contain the Galois element
+    /// associated with `step`, see ``GaloisElement/rotatingColumns(by:degree:)``.
+    /// - Throws: error upon failure to swap the ciphertext's rows.
+    /// - seealso: ``HeScheme/swapRows(of:using:)-50tac`` for an alternate API.
+    @inlinable
+    public mutating func swapRows(using evaluationKey: EvaluationKey<Scheme>) async throws
+        where Format == Scheme.CanonicalCiphertextFormat
+    {
+        try await Scheme.swapRowsAsync(of: &self, using: evaluationKey)
+    }
+
+    /// Asynchronously performs modulus switching on the ciphertext.
+    ///
+    /// - Throws: Error upon failure to mod-switch.
+    /// - seealso: ``HeScheme/modSwitchDown(_:)`` for an alternative API and more information.
+    @inlinable
+    public mutating func modSwitchDown() async throws where Format == Scheme.CanonicalCiphertextFormat {
+        try await Scheme.modSwitchDownAsync(&self)
+    }
+
+    /// Asynchronously performs modulus switching to a single modulus.
+    ///
+    /// If the ciphertext already has a single modulus, this is a no-op.
+    /// - Throws: Error upon failure to modulus switch.
+    /// - seealso: ``Ciphertext/modSwitchDown()-4an2b`` for more information and an alternative API.
+    @inlinable
+    public mutating func modSwitchDownToSingle() async throws where Format == Scheme.CanonicalCiphertextFormat {
+        try await Scheme.modSwitchDownToSingleAsync(&self)
+    }
+}
+
+extension Ciphertext where Format == Scheme.CanonicalCiphertextFormat {
+    /// Asynchronously applies a Galois transformation.
+    ///
+    /// - Parameters:
+    ///   - element: Galois element of the transformation. Must be odd in `[1, 2 * N - 1]` where `N` is the RLWE ring
+    /// dimension, given by ``EncryptionParameters/polyDegree``.
+    ///   - key: Evaluation key. Must contain Galois element `element`.
+    /// - Throws: Error upon failure to apply the Galois transformation.
+    /// - seealso: ``HeScheme/applyGalois(ciphertext:element:using:)`` for an alternative API and more information.
+    @inlinable
+    public mutating func applyGalois(element: Int, using key: EvaluationKey<Scheme>) async throws {
+        try await Scheme.applyGaloisAsync(ciphertext: &self, element: element, using: key)
+    }
+
+    /// Asynchronously Relinearizes the ciphertext.
+    ///
+    /// - Parameter key: Evaluation key to relinearize with. Must contain a `RelinearizationKey`.
+    /// - Throws: Error upon failure to relinearize.
+    /// - seealso: ``HeScheme/relinearize(_:using:)`` for an alternative API and more information.
+    @inlinable
+    public mutating func relinearize(using key: EvaluationKey<Scheme>) async throws {
+        try await Scheme.relinearizeAsync(&self, using: key)
+    }
+}
+
+// MARK: - Async collection extensions
+
+extension Collection {
+    /// Sums together the ciphertexts in the collection.
+    /// - Throws: Precondition failure if the collection is empty.
+    /// - Returns: The sum.
+    @inlinable
+    public func sum<Scheme>() async throws -> Element where Element == Ciphertext<Scheme, Eval> {
+        precondition(!isEmpty)
+        // swiftlint:disable:next force_unwrapping
+        return try await dropFirst().async.reduce(first!) { try await $0 + $1 }
+    }
+
+    /// Sums together the ciphertexts in the collection.
+    /// - Throws: Precondition failure if the collection is empty.
+    /// - Returns: The sum.
+    @inlinable
+    public func sum<Scheme>() async throws -> Element where Element == Ciphertext<Scheme, Coeff> {
+        precondition(!isEmpty)
+        // swiftlint:disable:next force_unwrapping
+        return try await dropFirst().async.reduce(first!) { try await $0 + $1 }
+    }
+
+    /// Sums together the ciphertexts in the collection.
+    /// - Throws: Precondition failure if the collection is empty.
+    /// - Returns: The sum.
+    @inlinable
+    public func sum<Scheme>() async throws -> Element where
+        Element == Ciphertext<Scheme, Scheme.CanonicalCiphertextFormat>
+    {
+        precondition(!isEmpty)
+        // swiftlint:disable:next force_unwrapping
+        return try await dropFirst().async.reduce(first!) { try await $0 + $1 }
+    }
+
+    /// Asynchronously computes an inner product between self and a collection of (optional) plaintexts in ``Eval``
+    /// format.
+    ///
+    /// The inner product encrypts `sum_{i, plaintexts[i] != nil} self[i] * plaintexts[i]`. `plaintexts[i]`
+    /// may be `nil`, which denotes a zero plaintext.
+    /// - Parameter plaintexts: Plaintexts. Must not be empty and have `count` matching `self.count`.
+    /// - Returns: A ciphertext encrypting the inner product.
+    /// - Throws: Error upon failure to compute inner product.
+    @inlinable
+    public func innerProduct<Scheme>(plaintexts: some Collection<Plaintext<Scheme, Eval>?>) async throws -> Element
+        where Element == Ciphertext<Scheme, Eval>
+    {
+        try await Scheme.innerProductAsync(ciphertexts: self, plaintexts: plaintexts)
+    }
+
+    /// Asynchronously computes an inner product between self and a collection of plaintexts in ``Eval`` format.
+    ///
+    /// The inner product encrypts `sum_{i} self[i] * plaintexts[i]`.
+    /// - Parameter plaintexts: Plaintexts. Must not be empty and have `count` matching `self.count`.
+    /// - Returns: A ciphertext encrypting the inner product.
+    /// - Throws: Error upon failure to compute inner product.
+    @inlinable
+    public func innerProduct<Scheme>(plaintexts: some Collection<Plaintext<Scheme, Eval>>) async throws -> Element
+        where Element == Ciphertext<Scheme, Eval>
+    {
+        try await Scheme.innerProductAsync(ciphertexts: self, plaintexts: plaintexts)
+    }
+
+    /// Asynchronously computes an inner product between self and another collection of ciphertexts.
+    ///
+    /// The inner product encrypts `sum_{i} self[i] * ciphertexts[i]`.
+    /// - Parameter ciphertexts: Ciphertexts. Must not be empty and have `count` matching `self.count`.
+    /// - Returns: A ciphertext encrypting the inner product.
+    /// - Throws: Error upon failure to compute inner product.
+    @inlinable
+    public func innerProduct<Scheme>(ciphertexts: some Collection<Element>) async throws -> Element
+        where Element == Ciphertext<Scheme, Scheme.CanonicalCiphertextFormat>
+    {
+        try await Scheme.innerProductAsync(self, ciphertexts)
+    }
 }

Sources/HomomorphicEncryption/Ciphertext.swift

@@ -485,7 +485,7 @@ public protocol HeScheme: Sendable {
     ///   - evaluationKey: Evaluation key to use in the HE computation. Must contain the Galois element associated with
     /// `step`, see ``GaloisElement/rotatingColumns(by:degree:)``.
     /// - Throws: failure to rotate ciphertext's columns.
-    /// - seealso: ``Ciphertext/rotateColumns(by:using:)`` for an alternate API.
+    /// - seealso: ``Ciphertext/rotateColumns(by:using:)-4f3tp`` for an alternate API.
     static func rotateColumns(
         of ciphertext: inout CanonicalCiphertext,
         by step: Int,
@@ -513,7 +513,7 @@ public protocol HeScheme: Sendable {
     ///   - evaluationKey: Evaluation key to use in the HE computation. Must contain the Galois element returned from
     /// ``GaloisElement/swappingRows(degree:)``.
     /// - Throws: error upon failure to swap the ciphertext's rows.
-    /// - seealso: ``Ciphertext/swapRows(using:)`` for an alternate API. ``swapRowsAsync(of:using:)`` for an async
+    /// - seealso: ``Ciphertext/swapRows(using:)-4o179`` for an alternate API. ``swapRowsAsync(of:using:)`` for an async
     /// version of this API
     static func swapRows(of ciphertext: inout CanonicalCiphertext, using evaluationKey: EvaluationKey) throws
 
@@ -862,7 +862,7 @@ public protocol HeScheme: Sendable {
     /// serialization and sending the ciphertext to the secret key owner.
     /// - Parameter ciphertext: Ciphertext; must have > 1 ciphertext modulus.
     /// - Throws: Error upon failure to mod-switch.
-    /// - seealso: ``Ciphertext/modSwitchDown()`` for an alternative API.
+    /// - seealso: ``Ciphertext/modSwitchDown()-4an2b`` for an alternative API.
     /// - seealso: ``modSwitchDownAsync(_:)``  for an async version of this API
     static func modSwitchDown(_ ciphertext: inout CanonicalCiphertext) throws
 
@@ -873,7 +873,7 @@ public protocol HeScheme: Sendable {
     ///
     /// If the ciphertext already has a single modulus, this is a no-op.
     /// - Throws: Error upon failure to modulus switch.
-    /// - seealso: ``Ciphertext/modSwitchDownToSingle()`` for more information and an alternative API.
+    /// - seealso: ``Ciphertext/modSwitchDownToSingle()-3x0dy`` for more information and an alternative API.
     /// - seealso: ``modSwitchDownToSingleAsync(_:)``  for an async version of this API
     static func modSwitchDownToSingle(_ ciphertext: inout CanonicalCiphertext) throws
 

Sources/HomomorphicEncryption/HeScheme.swift

@@ -359,10 +359,7 @@ extension MulPirServer {
                 let startIndex = dataChunk.startIndex + expandedDim0Query.count * columnIndex
                 let endIndex = min(startIndex + expandedDim0Query.count, dataChunk.endIndex)
                 let plaintexts = dataChunk[startIndex..<endIndex]
-                return try await Scheme.innerProductAsync(
-                    ciphertexts: expandedDim0Query,
-                    plaintexts: plaintexts)
-                    .convertToCanonicalFormat()
+                return try await expandedDim0Query.innerProduct(plaintexts: plaintexts).convertToCanonicalFormat()
             })
 
         var queryStartingIndex = expandedRemainingQuery.startIndex
@@ -373,17 +370,16 @@ extension MulPirServer {
                 .async.map { startIndex in
                     let vector0 = expandedRemainingQuery[currentIndex..<currentIndex + dimensionSize]
                     let vector1 = currentResults[startIndex..<startIndex + dimensionSize]
-                    var product = try await Scheme.innerProductAsync(vector0, vector1)
-                    try await Scheme.relinearizeAsync(&product, using: evaluationKey)
+                    var product = try await vector0.innerProduct(ciphertexts: vector1)
+                    try await product.relinearize(using: evaluationKey)
                     return product
                 })
             queryStartingIndex += dimensionSize
         }
 
-        precondition(
-            intermediateResults.count == 1,
-            "There should be only 1 ciphertext in the final result for each chunk")
-        try await Scheme.modSwitchDownToSingleAsync(&intermediateResults[0])
+        precondition(intermediateResults.count == 1,
+                     "There should be only 1 ciphertext in the final result for each chunk")
+        try await intermediateResults[0].modSwitchDownToSingle()
         return try await intermediateResults[0].convertToCoeffFormat()
     }
 

Sources/PrivateInformationRetrieval/IndexPir/MulPir.swift

@@ -90,7 +90,7 @@ extension PirUtilProtocol {
         let applyGaloisCount = 1 << ((targetElement - 1).log2 - (galoisElement - 1).log2)
         var currElement = 1
         for await _ in (0..<applyGaloisCount).async {
-            try await Scheme.applyGaloisAsync(ciphertext: &c1, element: galoisElement, using: evaluationKey)
+            try await c1.applyGalois(element: galoisElement, using: evaluationKey)
             currElement *= galoisElement
             currElement %= (2 * degree)
         }

Sources/PrivateInformationRetrieval/IndexPir/PirUtil.swift

@@ -199,8 +199,8 @@ extension CiphertextMatrix {
     /// - Throws: Error upon failure to modulus switch.
     @inlinable
     public mutating func modSwitchDownToSingle() async throws where Format == Scheme.CanonicalCiphertextFormat {
-        for index in 0..<ciphertexts.count {
-            try await Scheme.modSwitchDownToSingleAsync(&ciphertexts[index])
+        for index in ciphertexts.indices {
+            try await ciphertexts[index].modSwitchDownToSingle()
         }
     }
 }
@@ -337,10 +337,7 @@ extension CiphertextMatrix {
         let rotateCount = simdColumnCount / (copiesInMask * columnCountPowerOfTwo) - 1
         var ciphertextCopyRight = ciphertext
         for await _ in (0..<rotateCount).async {
-            try await Scheme.rotateColumnsAsync(
-                of: &ciphertextCopyRight,
-                by: columnCountPowerOfTwo,
-                using: evaluationKey)
+            try await ciphertextCopyRight.rotateColumns(by: columnCountPowerOfTwo, using: evaluationKey)
             try await ciphertext += ciphertextCopyRight
         }
         // e.g., `ciphertext` now encrypts
@@ -349,7 +346,7 @@ extension CiphertextMatrix {
 
         // Duplicate values to both SIMD rows
         var ciphertextCopy = ciphertext
-        try await Scheme.swapRowsAsync(of: &ciphertextCopy, using: evaluationKey)
+        try await ciphertextCopy.swapRows(using: evaluationKey)
         try await ciphertext += ciphertextCopy
         // e.g., `ciphertext` now encrypts
         // [[3, 4, 3, 4, 3, 4, 3, 4],

Sources/PrivateNearestNeighborSearch/CiphertextMatrix.swift

@@ -176,7 +176,7 @@ extension PlaintextMatrix {
         for step in 0..<babyStepGiantStep.babyStep {
             rotatedStates.append(state)
             if step != babyStepGiantStep.babyStep - 1 {
-                try await Scheme.rotateColumnsAsync(of: &state, by: -1, using: evaluationKey)
+                try await state.rotateColumns(by: -1, using: evaluationKey)
             }
         }
         let rotatedCiphertexts: [Scheme.EvalCiphertext] = try await .init(
@@ -201,10 +201,7 @@ extension PlaintextMatrix {
                 let ciphertexts = rotatedCiphertexts[0..<plaintextRows.count]
 
                 // 2) Compute w_k
-                let innerProduct =
-                    try await Scheme.innerProductAsync(
-                        ciphertexts: ciphertexts,
-                        plaintexts: plaintextRows)
+                let innerProduct = try await ciphertexts.innerProduct(plaintexts: plaintextRows)
                 return try await innerProduct.convertToCanonicalFormat()
             }