Add _HomomorphicEncryptionExtras

Package.swift

@@ -40,6 +40,9 @@ let package = Package(
         .library(
             name: "HomomorphicEncryptionProtobuf",
             targets: ["HomomorphicEncryptionProtobuf"]),
+        .library(
+            name: "_HomomorphicEncryptionExtras",
+            targets: ["_HomomorphicEncryptionExtras"]),
         .library(
             name: "PrivateInformationRetrieval",
             targets: ["PrivateInformationRetrieval"]),
@@ -96,6 +99,10 @@ let package = Package(
                            .product(name: "SwiftProtobuf", package: "swift-protobuf")],
             exclude: ["generated/README.md"],
             swiftSettings: librarySettings),
+        .target(
+            name: "_HomomorphicEncryptionExtras",
+            dependencies: ["HomomorphicEncryption"],
+            swiftSettings: librarySettings),
         .target(
             name: "PrivateInformationRetrieval",
             dependencies: ["HomomorphicEncryption",
@@ -114,6 +121,7 @@ let package = Package(
             dependencies: [
                 .product(name: "Algorithms", package: "swift-algorithms"),
                 "HomomorphicEncryption",
+                "_HomomorphicEncryptionExtras",
             ],
             swiftSettings: librarySettings),
         .target(
@@ -128,6 +136,7 @@ let package = Package(
             name: "_TestUtilities",
             dependencies: [
                 "HomomorphicEncryption",
+                "_HomomorphicEncryptionExtras",
                 "PrivateInformationRetrieval",
                 "PrivateNearestNeighborSearch",
                 .product(name: "Numerics", package: "swift-numerics"),

Sources/HomomorphicEncryption/Ciphertext.swift

@@ -240,24 +240,6 @@ public struct Ciphertext<Scheme: HeScheme, Format: PolyFormat>: Equatable, Senda
         try Scheme.rotateColumns(of: &self, by: step, using: evaluationKey)
     }
 
-    /// Rotates the columns of the ciphertext by combining multiple rotation steps corresponding to Galois elements
-    /// available in the `evaluationKey`.
-    ///
-    /// - 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 Galois elements which can be
-    /// combined for the desired rotation step.
-    /// - Throws: Error upon failure to rotate ciphertext's columns.
-    /// - seealso: `HeScheme/rotateColumnsMultiStep(of:by:using:)` for an alternative API and more information.
-    @inlinable
-    package mutating func rotateColumnsMultiStep(by step: Int, using evaluationKey: EvaluationKey<Scheme>) throws
-        where Format == Scheme.CanonicalCiphertextFormat
-    {
-        try Scheme.rotateColumnsMultiStep(of: &self, by: step, using: evaluationKey)
-    }
-
     /// Swaps the rows of a ciphertext.
     ///
     /// A plaintext in ``EncodeFormat/simd`` format can be viewed a `2 x (N / 2)` matrix of coefficients.

Sources/HomomorphicEncryption/HeScheme.swift

@@ -1285,88 +1285,6 @@ extension HeScheme {
         try applyGalois(ciphertext: &ciphertext, element: element, using: evaluationKey)
     }
 
-    @inlinable
-    package static func rotateColumnsMultiStep(
-        of ciphertext: inout CanonicalCiphertext,
-        by step: Int,
-        using evaluationKey: EvaluationKey) throws
-    {
-        if step == 0 {
-            return
-        }
-
-        guard let galoisKey = evaluationKey.galoisKey else {
-            throw HeError.missingGaloisKey
-        }
-
-        // Short-circuit to single rotation if possible.
-        let degree = ciphertext.degree
-        let galoisElement = try GaloisElement.rotatingColumns(by: step, degree: degree)
-        if galoisKey.keys.keys.contains(galoisElement) {
-            try ciphertext.rotateColumns(by: step, using: evaluationKey)
-            return
-        }
-
-        let galoisElements = Array(galoisKey.keys.keys)
-        let steps = GaloisElement.stepsFor(elements: galoisElements, degree: degree).values.compactMap(\.self)
-
-        let positiveStep = if step < 0 {
-            step + degree / 2
-        } else {
-            step
-        }
-
-        let plan = try GaloisElement.planMultiStep(supportedSteps: steps, step: positiveStep, degree: degree)
-        guard let plan else {
-            throw HeError.invalidRotationStep(step: step, degree: degree)
-        }
-        for (step, count) in plan {
-            try (0..<count).forEach { _ in try ciphertext.rotateColumns(by: step, using: evaluationKey) }
-        }
-    }
-
-    @inlinable
-    package static func rotateColumnsMultiStepAsync(
-        of ciphertext: inout CanonicalCiphertext,
-        by step: Int,
-        using evaluationKey: EvaluationKey) async throws
-    {
-        if step == 0 {
-            return
-        }
-
-        guard let galoisKey = evaluationKey.galoisKey else {
-            throw HeError.missingGaloisKey
-        }
-
-        // Short-circuit to single rotation if possible.
-        let degree = ciphertext.degree
-        let galoisElement = try GaloisElement.rotatingColumns(by: step, degree: degree)
-        if galoisKey.keys.keys.contains(galoisElement) {
-            try await rotateColumnsAsync(of: &ciphertext, by: step, using: evaluationKey)
-            return
-        }
-
-        let galoisElements = Array(galoisKey.keys.keys)
-        let steps = GaloisElement.stepsFor(elements: galoisElements, degree: degree).values.compactMap(\.self)
-
-        let positiveStep = if step < 0 {
-            step + degree / 2
-        } else {
-            step
-        }
-
-        let plan = try GaloisElement.planMultiStep(supportedSteps: steps, step: positiveStep, degree: degree)
-        guard let plan else {
-            throw HeError.invalidRotationStep(step: step, degree: degree)
-        }
-        for (step, count) in plan {
-            for _ in 0..<count {
-                try await rotateColumnsAsync(of: &ciphertext, by: step, using: evaluationKey)
-            }
-        }
-    }
-
     @inlinable
     // swiftlint:disable:next missing_docs attributes
     public static func modSwitchDownToSingle(_ ciphertext: inout CanonicalCiphertext) throws {

Sources/HomomorphicEncryption/Keys.swift

@@ -50,7 +50,7 @@ extension SecretKey: PolyCollection {
 /// - seealso: ``HeScheme/relinearize(_:using:)`` and ``HeScheme/applyGalois(ciphertext:element:using:)`` for more
 /// details.
 @usableFromInline
-struct KeySwitchKey<Scheme: HeScheme>: Equatable, Sendable {
+package struct KeySwitchKey<Scheme: HeScheme>: Equatable, Sendable {
     /// The context used for key-switching operations.
     @usableFromInline let context: Context<Scheme>
     /// The ciphertexts of the key-switching key.
@@ -73,7 +73,7 @@ extension KeySwitchKey: PolyCollection {
 }
 
 @usableFromInline
-struct RelinearizationKey<Scheme: HeScheme>: Equatable, Sendable {
+package struct RelinearizationKey<Scheme: HeScheme>: Equatable, Sendable {
     @usableFromInline let keySwitchKey: KeySwitchKey<Scheme>
 
     @inlinable
@@ -92,8 +92,8 @@ extension RelinearizationKey: PolyCollection {
 }
 
 @usableFromInline
-struct GaloisKey<Scheme: HeScheme>: Equatable, Sendable {
-    @usableFromInline let keys: [Int: KeySwitchKey<Scheme>]
+package struct GaloisKey<Scheme: HeScheme>: Equatable, Sendable {
+    @usableFromInline package let keys: [Int: KeySwitchKey<Scheme>]
 
     @inlinable
     init(keys: [Int: KeySwitchKey<Scheme>]) {
@@ -118,8 +118,8 @@ extension GaloisKey: PolyCollection {
 ///
 /// Associated with a ``SecretKey``.
 public struct EvaluationKey<Scheme: HeScheme>: Equatable, Sendable {
-    @usableFromInline let galoisKey: GaloisKey<Scheme>?
-    @usableFromInline let relinearizationKey: RelinearizationKey<Scheme>?
+    @usableFromInline package let galoisKey: GaloisKey<Scheme>?
+    @usableFromInline package let relinearizationKey: RelinearizationKey<Scheme>?
 
     /// Returns the configuration for the evaluation key.
     public var config: EvaluationKeyConfig {

Sources/PrivateNearestNeighborSearch/MatrixMultiplication.swift

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+import _HomomorphicEncryptionExtras
 import Algorithms
 import Foundation
 import HomomorphicEncryption

Sources/_HomomorphicEncryptionExtras/Ciphertext.swift

@@ -0,0 +1,35 @@
+// Copyright 2024-2025 Apple Inc. and the Swift Homomorphic Encryption project authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import HomomorphicEncryption
+
+extension Ciphertext {
+    /// Rotates the columns of the ciphertext by combining multiple rotation steps corresponding to Galois elements
+    /// available in the `evaluationKey`.
+    ///
+    /// - 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 Galois elements which can be
+    /// combined for the desired rotation step.
+    /// - Throws: Error upon failure to rotate ciphertext's columns.
+    /// - seealso: `HeScheme/_rotateColumnsMultiStep(of:by:using:)` for an alternative API and more information.
+    @inlinable
+    public mutating func rotateColumnsMultiStep(by step: Int, using evaluationKey: EvaluationKey<Scheme>) throws
+        where Format == Scheme.CanonicalCiphertextFormat
+    {
+        try Scheme.rotateColumnsMultiStep(of: &self, by: step, using: evaluationKey)
+    }
+}

Sources/_HomomorphicEncryptionExtras/HeScheme.swift

@@ -0,0 +1,152 @@
+// Copyright 2024-2025 Apple Inc. and the Swift Homomorphic Encryption project authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import HomomorphicEncryption
+
+extension HeScheme {
+    /// Rotates the columns of a ciphertext throw multiple steps in async way.
+    /// - seealso: `HeScheme/rotateColumnsAsync(of:by:using:)` for the single-step API.
+    @inlinable
+    public static func rotateColumnsMultiStepAsync(
+        of ciphertext: inout CanonicalCiphertext,
+        by step: Int,
+        using evaluationKey: EvaluationKey) async throws
+    {
+        if step == 0 {
+            return
+        }
+
+        guard let galoisKey = evaluationKey.galoisKey else {
+            throw HeError.missingGaloisKey
+        }
+
+        // Short-circuit to single rotation if possible.
+        let degree = ciphertext.degree
+        let galoisElement = try GaloisElement.rotatingColumns(by: step, degree: degree)
+        if galoisKey.keys.keys.contains(galoisElement) {
+            try await rotateColumnsAsync(of: &ciphertext, by: step, using: evaluationKey)
+            return
+        }
+
+        let galoisElements = Array(galoisKey.keys.keys)
+        let steps = GaloisElement.stepsFor(elements: galoisElements, degree: degree).values.compactMap(\.self)
+
+        let positiveStep = if step < 0 {
+            step + degree / 2
+        } else {
+            step
+        }
+
+        let plan = try GaloisElement.planMultiStep(supportedSteps: steps, step: positiveStep, degree: degree)
+        guard let plan else {
+            throw HeError.invalidRotationStep(step: step, degree: degree)
+        }
+        for (step, count) in plan {
+            for _ in 0..<count {
+                try await rotateColumnsAsync(of: &ciphertext, by: step, using: evaluationKey)
+            }
+        }
+    }
+
+    ///  Rotates the columns of a ciphertext through one or more rotations.
+    /// - seealso: `HeScheme/rotateColumns(of:by:using:)` for the single-step API.
+    @inlinable
+    public static func rotateColumnsMultiStep(
+        of ciphertext: inout CanonicalCiphertext,
+        by step: Int,
+        using evaluationKey: EvaluationKey) throws
+    {
+        if step == 0 {
+            return
+        }
+
+        guard let galoisKey = evaluationKey.galoisKey else {
+            throw HeError.missingGaloisKey
+        }
+
+        // Short-circuit to single rotation if possible.
+        let degree = ciphertext.degree
+        let galoisElement = try GaloisElement.rotatingColumns(by: step, degree: degree)
+        if galoisKey.keys.keys.contains(galoisElement) {
+            try ciphertext.rotateColumns(by: step, using: evaluationKey)
+            return
+        }
+
+        let galoisElements = Array(galoisKey.keys.keys)
+        let steps = GaloisElement.stepsFor(elements: galoisElements, degree: degree).values.compactMap(\.self)
+
+        let positiveStep = if step < 0 {
+            step + degree / 2
+        } else {
+            step
+        }
+
+        let plan = try GaloisElement.planMultiStep(supportedSteps: steps, step: positiveStep, degree: degree)
+        guard let plan else {
+            throw HeError.invalidRotationStep(step: step, degree: degree)
+        }
+        for (step, count) in plan {
+            try (0..<count).forEach { _ in try ciphertext.rotateColumns(by: step, using: evaluationKey) }
+        }
+    }
+
+    /// Sum up an array of ciphertexts after rotate their columns one-by-one.
+    ///
+    /// The i-th (starting from 0) ciphertext is rotated by i * `step` steps before adding up.
+    /// - Parameters:
+    ///  - ciphertexts: ciphertexts to be added up.
+    ///  - step: the rotation steps for each ciphertext.
+    /// - evaluationKey: the evaluation key for rotation.
+    /// - Throws: Error upon failure to compute the inverse.
+    @inlinable
+    public static func rotateColumnsAndSumAsync(
+        _ ciphertexts: consuming [CanonicalCiphertext],
+        by step: Int,
+        using evaluationKey: EvaluationKey) async throws -> CanonicalCiphertext
+    {
+        guard var accumulator = ciphertexts.popLast() else {
+            preconditionFailure("No ciphertexts to sum up.")
+        }
+        if ciphertexts.isEmpty {
+            return accumulator
+        }
+
+        for ciphertext in ciphertexts.reversed() {
+            try await rotateColumnsMultiStepAsync(
+                of: &accumulator,
+                by: step,
+                using: evaluationKey)
+            try await addAssignAsync(&accumulator, ciphertext)
+        }
+        return accumulator
+    }
+
+    /// Sum up two ciphertexts after swap the row of second one.
+    ///
+    /// - Parameters:
+    ///  - ciphertexts: ciphertexts to be added up.
+    ///  - step: the rotation steps for each ciphertext.
+    /// - evaluationKey: the evaluation key for rotation.
+    /// - Throws: Error upon failure to compute the inverse.
+    @inlinable
+    public static func swapRowsAndAddAsync(
+        swapping ciphertext0: consuming CanonicalCiphertext,
+        addingTo ciphertext1: consuming CanonicalCiphertext,
+        using evaluationKey: EvaluationKey) async throws -> CanonicalCiphertext
+    {
+        try await swapRowsAsync(of: &ciphertext0, using: evaluationKey)
+        try await addAssignAsync(&ciphertext0, ciphertext1)
+        return ciphertext0
+    }
+}

Sources/_TestUtilities/HeApiTestUtils.swift

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+import _HomomorphicEncryptionExtras
 import HomomorphicEncryption
 import ModularArithmetic
 import Testing