Keep Context with HeScheme generic instead of Scalar

Author: fboemer

Sources/HomomorphicEncryption/Bfv/Bfv.swift

@@ -16,7 +16,7 @@ import ModularArithmetic
 
 /// Brakerski-Fan-Vercauteren cryptosystem.
 public enum Bfv<T: ScalarType>: HeScheme {
-    public typealias Context = HomomorphicEncryption.Context<T>
+    public typealias Context = HomomorphicEncryption.Context<Self>
 
     public typealias Scalar = T
 

Sources/HomomorphicEncryption/Context.swift

@@ -16,7 +16,16 @@
 ///
 /// HE operations are typically only supported between objects, such as ``Ciphertext``, ``Plaintext``,
 /// ``EvaluationKey``, ``SecretKey``,  with the same context.
-public final class Context<Scalar: ScalarType>: Equatable, Sendable, HeContext {
+public final class Context<Scheme: HeScheme>: Equatable, Sendable, HeContext {
+    public typealias Scalar = Scheme.Scalar
+
+    /// The (row, column) dimension counts for ``EncodeFormat/simd`` encoding.
+    ///
+    /// If the HE scheme does not support ``EncodeFormat/simd`` encoding, returns `nil`.
+    public var simdDimensions: SimdEncodingDimensions? {
+        Scheme.encodeSimdDimensions(for: encryptionParameters)
+    }
+
     /// Encryption parameters.
     public let encryptionParameters: EncryptionParameters<Scalar>
 
@@ -124,15 +133,6 @@ public final class Context<Scalar: ScalarType>: Equatable, Sendable, HeContext {
         lhs === rhs || lhs.encryptionParameters == rhs.encryptionParameters
     }
 
-    /// The (row, column) dimension counts for ``EncodeFormat/simd`` encoding.
-    ///
-    /// If the HE scheme does not support ``EncodeFormat/simd`` encoding, returns `nil`.
-    public func simdDimensions<Scheme: HeScheme>(for _: Scheme.Type) -> SimdEncodingDimensions?
-        where Scheme.Scalar == Scalar
-    {
-        Scheme.encodeSimdDimensions(for: encryptionParameters)
-    }
-
     @inlinable
     public func _getRnsTool(moduliCount: Int) -> _RnsTool<Scalar> {
         precondition(moduliCount <= rnsTools.count && moduliCount > 0, "Invalid number of moduli")

Sources/HomomorphicEncryption/HeScheme.swift

@@ -86,19 +86,18 @@ public struct SimdEncodingDimensions: Codable, Equatable, Hashable, Sendable {
 }
 
 public protocol HeContext: Equatable, Sendable, CustomStringConvertible {
-    associatedtype Scalar: ScalarType
+    associatedtype Scheme: HeScheme where Scheme.Scalar == Self.Scalar
+    typealias Scalar = Scheme.Scalar
 
     var encryptionParameters: EncryptionParameters<Scalar> { get }
     var ciphertextContext: PolyContext<Scalar> { get }
     var plaintextContext: PolyContext<Scalar> { get }
     var secretKeyContext: PolyContext<Scalar> { get }
+    var simdDimensions: SimdEncodingDimensions? { get }
     var simdEncodingMatrix: [Int] { get }
 
     init(encryptionParameters: EncryptionParameters<Scalar>) throws
     func _getRnsTool(moduliCount: Int) throws -> _RnsTool<Scalar>
-
-    func simdDimensions<Scheme: HeScheme>(for _: Scheme.Type) -> SimdEncodingDimensions?
-        where Scheme.Scalar == Scalar
 }
 
 extension HeContext {
@@ -116,6 +115,13 @@ extension HeContext {
     public var bitsPerPlaintext: Int { encryptionParameters.bitsPerPlaintext }
     /// The number of bytes that can be encoded in a single ``Plaintext``.
     public var bytesPerPlaintext: Int { encryptionParameters.bytesPerPlaintext }
+
+    /// The (row, column) dimension counts for ``EncodeFormat/simd`` encoding.
+    ///
+    /// If the HE scheme does not support ``EncodeFormat/simd`` encoding, returns `nil`.
+    public var simdDimensions: SimdEncodingDimensions? {
+        Scheme.encodeSimdDimensions(for: encryptionParameters)
+    }
 }
 
 /// Protocol for HE schemes.

Sources/HomomorphicEncryption/NoOpScheme.swift

@@ -17,7 +17,7 @@
 /// The scheme simply takes the plaintext as a "ciphertext" and
 /// ignores any ciphertext coefficient moduli.
 public enum NoOpScheme: HeScheme {
-    public typealias Context = HomomorphicEncryption.Context<UInt64>
+    public typealias Context = HomomorphicEncryption.Context<Self>
     public typealias Scalar = UInt64
     public typealias CanonicalCiphertextFormat = Coeff
 

Sources/PrivateNearestNeighborSearch/CiphertextMatrix.swift

@@ -66,7 +66,7 @@ public struct CiphertextMatrix<Scheme: HeScheme, Format: PolyFormat>: Equatable,
             throw PnnsError.emptyCiphertextArray
         }
         let encryptionParameters = context.encryptionParameters
-        guard let simdDimensions = context.simdDimensions(for: Scheme.self) else {
+        guard let simdDimensions = context.simdDimensions else {
             throw PnnsError.simdEncodingNotSupported(for: encryptionParameters)
         }
         let expectedCiphertextCount = try PlaintextMatrix<Scheme, Format>.plaintextCount(

Sources/PrivateNearestNeighborSearch/PlaintextMatrix.swift

@@ -121,7 +121,7 @@ public struct PlaintextMatrix<Scheme: HeScheme, Format: PolyFormat>: Equatable,
         }
         let context = plaintexts[0].context
         let encryptionParameters = context.encryptionParameters
-        guard let simdDimensions = context.simdDimensions(for: Scheme.self) else {
+        guard let simdDimensions = context.simdDimensions else {
             throw PnnsError.simdEncodingNotSupported(for: encryptionParameters)
         }
         let expectedPlaintextCount = try PlaintextMatrix.plaintextCount(
@@ -276,7 +276,7 @@ public struct PlaintextMatrix<Scheme: HeScheme, Format: PolyFormat>: Equatable,
                                       values: [Scalar]) throws -> [Scheme.CoeffPlaintext]
     {
         let degree = context.degree
-        guard let simdColumnCount = context.simdDimensions(for: Scheme.self)?.columnCount else {
+        guard let simdColumnCount = context.simdDimensions?.columnCount else {
             throw PnnsError.simdEncodingNotSupported(for: context.encryptionParameters)
         }
 
@@ -334,7 +334,7 @@ public struct PlaintextMatrix<Scheme: HeScheme, Format: PolyFormat>: Equatable,
         values: [Scalar]) throws -> [Plaintext<Scheme, Coeff>]
     {
         let encryptionParameters = context.encryptionParameters
-        guard let simdDimensions = context.simdDimensions(for: Scheme.self) else {
+        guard let simdDimensions = context.simdDimensions else {
             throw PnnsError.simdEncodingNotSupported(for: encryptionParameters)
         }
         guard simdDimensions.rowCount == 2 else {
@@ -411,7 +411,7 @@ public struct PlaintextMatrix<Scheme: HeScheme, Format: PolyFormat>: Equatable,
         values: [Scalar]) throws -> [Scheme.CoeffPlaintext]
     {
         let encryptionParameters = context.encryptionParameters
-        guard let simdDimensions = context.simdDimensions(for: Scheme.self) else {
+        guard let simdDimensions = context.simdDimensions else {
             throw PnnsError.simdEncodingNotSupported(for: encryptionParameters)
         }
         let simdColumnCount = simdDimensions.columnCount

Sources/TestUtilities/HeApiTestUtils.swift

@@ -128,7 +128,7 @@ public enum HeAPITestHelpers {
     }
 
     /// test the evaluation key configuration
-    public static func schemeEvaluationKeyTest(context _: Context<some ScalarType>) throws {
+    public static func schemeEvaluationKeyTest(context _: Context<some HeScheme>) throws {
         do {
             let config = EvaluationKeyConfig()
             #expect(!config.hasRelinearizationKey)

Sources/TestUtilities/PirUtilities/IndexPirTests.swift

@@ -22,7 +22,7 @@ extension PirTestUtils {
         /// Testing client configuration.
         @inlinable
         func generateParameter() throws {
-            let context: Context<UInt64> = try TestUtils.getTestContext()
+            let context: Context<Bfv<UInt64>> = try TestUtils.getTestContext()
             // unevenDimensions: false
             do {
                 let config = try IndexPirConfig(entryCount: 16,

Tests/HomomorphicEncryptionTests/EncryptionParametersTests.swift

@@ -243,7 +243,7 @@ struct EncryptionParametersTests {
 
             // check skipLSBsForDecryption
             do {
-                let context = try Context<UInt64>(encryptionParameters: params)
+                let context = try Context<Bfv<UInt64>>(encryptionParameters: params)
                 let data = TestUtils.getRandomPlaintextData(count: params.polyDegree, in: 0..<params.plaintextModulus)
                 let plaintext: Plaintext<Bfv<UInt64>, Coeff> = try context.encode(values: data, format: .coefficient)
                 let secretKey: SecretKey<Bfv<UInt64>> = try context.generateSecretKey()

Tests/HomomorphicEncryptionTests/HeAPITests.swift

@@ -86,7 +86,7 @@ struct HeAPITests {
 
     @Test
     func testNoOpScheme() async throws {
-        let context: Context<NoOpScheme.Scalar> = try TestUtils.getTestContext()
+        let context: Context<NoOpScheme> = try TestUtils.getTestContext()
         try HeAPITestHelpers.schemeEncodeDecodeTest(context: context, scheme: NoOpScheme.self)
         try HeAPITestHelpers.schemeEncryptDecryptTest(context: context, scheme: NoOpScheme.self)
         try HeAPITestHelpers.schemeEncryptZeroDecryptTest(context: context, scheme: NoOpScheme.self)