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 }
@@ -993,7 +1004,7 @@ extension Ciphertext {
     ///
     /// 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() async throws where Format == Scheme.CanonicalCiphertextFormat {
         try await Scheme.modSwitchDownToSingleAsync(&self)
@@ -1024,3 +1035,80 @@ extension Ciphertext where Format == Scheme.CanonicalCiphertextFormat {
         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()
         }
     }
 }

Sources/PrivateNearestNeighborSearch/CiphertextMatrix.swift

@@ -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()
             }