Make HeScheme context flexible

Benchmarks/RlweBenchmark/RlweBenchmark.swift

@@ -43,7 +43,7 @@ func getRandomPlaintextData<T: ScalarType>(count: Int, in range: Range<T>) -> [T
 
 struct RlweBenchmarkContext<Scheme: HeScheme>: Sendable {
     var encryptionParameters: EncryptionParameters<Scheme.Scalar>
-    var context: Context<Scheme>
+    var context: Scheme.Context
 
     let data: [Scheme.Scalar]
     let signedData: [Scheme.SignedScalar]
@@ -69,7 +69,7 @@ struct RlweBenchmarkContext<Scheme: HeScheme>: Sendable {
                                                                  Scheme.Scalar.self),
                                                              errorStdDev: ErrorStdDev.stdDev32,
                                                              securityLevel: SecurityLevel.quantum128)
-        self.context = try Context(encryptionParameters: encryptionParameters)
+        self.context = try Scheme.Context(encryptionParameters: encryptionParameters)
         self.secretKey = try context.generateSecretKey()
         let columnElement = GaloisElement.swappingRows(degree: polyDegree)
         let rowElement = try GaloisElement.rotatingColumns(by: rotateColumnsStep, degree: polyDegree)
@@ -163,7 +163,7 @@ func contextInitBenchmark<Scheme: HeScheme>(_: Scheme.Type, config: EncryptionPa
             let encryptionParameters = try EncryptionParameters<Scheme.Scalar>(config: config)
             benchmark.startMeasurement()
             for _ in benchmark.scaledIterations {
-                try blackHole(_ = Context<Scheme>(encryptionParameters: encryptionParameters))
+                try blackHole(_ = Scheme.Context(encryptionParameters: encryptionParameters))
             }
         }
     }
@@ -279,8 +279,7 @@ func decodeSignedSimdBenchmark<Scheme: HeScheme>(_: Scheme.Type) -> () -> Void {
             let benchmarkContext: RlweBenchmarkContext<Scheme> = try StaticRlweBenchmarkContext.getBenchmarkContext()
             benchmark.startMeasurement()
             for _ in benchmark.scaledIterations {
-                try blackHole(
-                    benchmarkContext.coeffPlaintext.decode(format: .simd) as [Scheme.SignedScalar])
+                try blackHole(benchmarkContext.coeffPlaintext.decode(format: .simd) as [Scheme.SignedScalar])
             }
         }
     }
@@ -292,7 +291,7 @@ func generateSecretKeyBenchmark<Scheme: HeScheme>(_: Scheme.Type) -> () -> Void
             let benchmarkContext: RlweBenchmarkContext<Scheme> = try StaticRlweBenchmarkContext.getBenchmarkContext()
             benchmark.startMeasurement()
             for _ in benchmark.scaledIterations {
-                try blackHole(benchmarkContext.context.generateSecretKey())
+                try blackHole(benchmarkContext.context.generateSecretKey() as SecretKey<Scheme>)
             }
         }
     }
@@ -320,8 +319,7 @@ func encryptBenchmark<Scheme: HeScheme>(_: Scheme.Type) -> () -> Void {
             let benchmarkContext: RlweBenchmarkContext<Scheme> = try StaticRlweBenchmarkContext.getBenchmarkContext()
             benchmark.startMeasurement()
             for _ in benchmark.scaledIterations {
-                try blackHole(
-                    benchmarkContext.coeffPlaintext.encrypt(using: benchmarkContext.secretKey))
+                try blackHole(benchmarkContext.coeffPlaintext.encrypt(using: benchmarkContext.secretKey))
             }
         }
     }
@@ -333,9 +331,7 @@ func decryptBenchmark<Scheme: HeScheme>(_: Scheme.Type) -> () -> Void {
             let benchmarkContext: RlweBenchmarkContext<Scheme> = try StaticRlweBenchmarkContext.getBenchmarkContext()
             benchmark.startMeasurement()
             for _ in benchmark.scaledIterations {
-                try blackHole(
-                    benchmarkContext.evalCiphertext.decrypt(
-                        using: benchmarkContext.secretKey))
+                try blackHole(benchmarkContext.evalCiphertext.decrypt(using: benchmarkContext.secretKey))
             }
         }
     }

Snippets/HomomorphicEncryption/BasicsSnippet.swift

@@ -24,13 +24,11 @@ import HomomorphicEncryption
 let encryptParams =
     try EncryptionParameters<UInt64>(from: .insecure_n_8_logq_5x18_logt_5)
 // Perform pre-computation for HE computation with these parameters.
-let context = try Context<Bfv<UInt64>>(encryptionParameters: encryptParams)
+let context = try Bfv<UInt64>.Context(encryptionParameters: encryptParams)
 
 // We encode N values using coefficient encoding.
 let values: [UInt64] = [8, 5, 12, 12, 15, 0, 8, 5]
-let plaintext: Bfv<UInt64>.CoeffPlaintext = try context.encode(
-    values: values,
-    format: .coefficient)
+let plaintext = try context.encode(values: values, format: .coefficient)
 
 // We generate a secret key and use it to encrypt the plaintext.
 let secretKey = try context.generateSecretKey()

Snippets/HomomorphicEncryption/EncryptionParametersSnippet.swift

@@ -85,10 +85,8 @@ func summarize<Scheme: HeScheme>(
     parameters: EncryptionParameters<Scheme.Scalar>, _: Scheme.Type) throws
 {
     let values = (0..<8).map { Scheme.Scalar($0) }
-    let context = try Context<Scheme>(encryptionParameters: parameters)
-    let plaintext: Scheme.CoeffPlaintext = try context.encode(
-        values: values,
-        format: .coefficient)
+    let context = try Scheme.Context(encryptionParameters: parameters)
+    let plaintext = try context.encode(values: values, format: .coefficient)
     let secretKey = try context.generateSecretKey()
     let ciphertext = try plaintext.encrypt(using: secretKey)
     let noiseBudget = try ciphertext.noiseBudget(

Snippets/HomomorphicEncryption/EvaluationKeySnippet.swift

@@ -46,14 +46,12 @@ let encryptionParameters =
     try EncryptionParameters<UInt32>(from: .insecure_n_8_logq_5x18_logt_5)
 precondition(encryptionParameters.plaintextModulus == 17)
 // Perform pre-computation for HE computation with these parameters.
-let context = try Context<Bfv<UInt32>>(encryptionParameters: encryptionParameters)
+let context = try Bfv<UInt32>.Context(encryptionParameters: encryptionParameters)
 
 // We encode N values using SIMD encoding.
 // Each value is a Bfv<UInt32>.Scalar, which is UInt32.
 let values = (0..<8).map { Bfv<UInt32>.Scalar($0) }
-let plaintext: Bfv<UInt32>.CoeffPlaintext = try context.encode(
-    values: values,
-    format: .simd)
+let plaintext = try context.encode(values: values, format: .simd)
 
 // We generate a secret key and encrypt the plaintext.
 let secretKey = try context.generateSecretKey()

Snippets/HomomorphicEncryption/MultiplicationSnippet.swift

@@ -34,16 +34,14 @@ let encryptParams =
     try EncryptionParameters<UInt32>(from: .insecure_n_8_logq_5x18_logt_5)
 precondition(encryptParams.plaintextModulus == 17)
 // Perform pre-computation for HE computation with these parameters.
-let context = try Context<Bfv<UInt32>>(encryptionParameters: encryptParams)
+let context = try Bfv<UInt32>.Context(encryptionParameters: encryptParams)
 
 // We don't need to use all the slots in the encoding.
 // However, performing HE operations on ciphertexts with fewer slots doesn't give
 // any runtime savings.
 let valueCount = encryptParams.polyDegree / 2
 let values = (0..<valueCount).map { UInt32($0) }
-let plaintext: Bfv<UInt32>.CoeffPlaintext = try context.encode(
-    values: values,
-    format: .simd)
+let plaintext = try context.encode(values: values, format: .simd)
 
 // We generate a secret key and encrypt the plaintext.
 let secretKey = try context.generateSecretKey()

Snippets/HomomorphicEncryption/NoiseBudgetSnippet.swift

@@ -45,13 +45,11 @@ func checkDecryptsDecodes<Scheme: HeScheme>(
 let encryptionParameters =
     try EncryptionParameters<UInt32>(from: .insecure_n_8_logq_5x18_logt_5)
 // Perform pre-computation for HE computation with these parameters.
-let context = try Context<Bfv<UInt32>>(encryptionParameters: encryptionParameters)
+let context = try Bfv<UInt32>.Context(encryptionParameters: encryptionParameters)
 
 // We encode N values in coefficient format and SIMD encoding.
 let values = (0..<8).map { UInt32($0) }
-let plaintext: Bfv<UInt32>.CoeffPlaintext = try context.encode(
-    values: values,
-    format: .simd)
+let plaintext = try context.encode(values: values, format: .simd)
 
 // We generate a secret key and encrypt the plaintext.
 let secretKey = try context.generateSecretKey()

Snippets/HomomorphicEncryption/SerializationSnippet.swift

@@ -36,18 +36,15 @@ extension SerializedCiphertext {
 let encryptionParameters =
     try EncryptionParameters<UInt32>(from: .n_4096_logq_27_28_28_logt_5)
 // Perform pre-computation for HE computation with these parameters.
-let context = try Context<Bfv<UInt32>>(encryptionParameters: encryptionParameters)
+let context = try Bfv<UInt32>.Context(encryptionParameters: encryptionParameters)
 
 // We encode some values in coefficient format and coefficient encoding.
 let values = (0..<8).map { UInt32($0) }
-let plaintext: Bfv<UInt32>.CoeffPlaintext = try context.encode(
-    values: values,
-    format: .coefficient)
+let plaintext = try context.encode(values: values, format: .coefficient)
 
 // We generate a secret key and encrypt the plaintext.
 let secretKey = try context.generateSecretKey()
-var ciphertext: Bfv<UInt32>.CanonicalCiphertext = try plaintext
-    .encrypt(using: secretKey)
+var ciphertext = try plaintext.encrypt(using: secretKey)
 
 // A freshly-encrypted ciphertext has one of its 2 polynomials sampled uniformly
 // randomly from a seeded random number generator (RNG). We compress such

Sources/BenchmarkUtilities/PirBenchmarkUtilities.swift

@@ -83,7 +83,7 @@ extension PrivateInformationRetrieval.Response {
 
 struct ProcessBenchmarkContext<Server: IndexPirServer> {
     let database: [[UInt8]]
-    let context: Context<Server.Scheme>
+    let context: Server.Scheme.Context
     let parameter: IndexPirParameter
     init(server _: Server.Type, pirConfig: IndexPirConfig,
          parameterConfig: PirEncryptionParametersConfig) throws
@@ -93,7 +93,7 @@ struct ProcessBenchmarkContext<Server: IndexPirServer> {
         self.database = getDatabaseForTesting(
             numberOfEntries: pirConfig.entryCount,
             entrySizeInBytes: pirConfig.entrySizeInBytes)
-        self.context = try Context(encryptionParameters: encryptParameter)
+        self.context = try Server.Scheme.Context(encryptionParameters: encryptParameter)
         self.parameter = Server.generateParameter(config: pirConfig, with: context)
     }
 }
@@ -171,7 +171,7 @@ struct IndexPirBenchmarkContext<Server: IndexPirServer, Client: IndexPirClient>
     {
         let encryptParameter: EncryptionParameters<Server.Scheme.Scalar> =
             try EncryptionParameters(from: parameterConfig)
-        let context = try Context<Server.Scheme>(encryptionParameters: encryptParameter)
+        let context = try Server.Scheme.Context(encryptionParameters: encryptParameter)
         let indexPirParameters = Server.generateParameter(config: pirConfig, with: context)
         let database = getDatabaseForTesting(
             numberOfEntries: pirConfig.entryCount,
@@ -286,7 +286,7 @@ struct KeywordPirBenchmarkContext<IndexServer: IndexPirServer, IndexClient: Inde
     {
         let encryptParameter: EncryptionParameters<Server.Scheme.Scalar> =
             try EncryptionParameters(from: parameterConfig)
-        let context = try Context<Server.Scheme>(encryptionParameters: encryptParameter)
+        let context = try Server.Scheme.Context(encryptionParameters: encryptParameter)
         let rows = (0..<databaseCount).map { index in KeywordValuePair(
             keyword: [UInt8](String(index).utf8),
             value: (0..<payloadSize).map { _ in UInt8.random(in: 0..<UInt8.max) })

Sources/BenchmarkUtilities/PnnsBenchmarkUtilities.swift

@@ -188,7 +188,7 @@ extension PrivateNearestNeighborSearch.Response {
 
 struct PnnsProcessBenchmarkContext<Scheme: HeScheme> {
     let database: Database
-    let contexts: [Context<Scheme>]
+    let contexts: [Scheme.Context]
     let serverConfig: ServerConfig<Scheme>
 
     init(databaseConfig: DatabaseConfig,
@@ -238,7 +238,7 @@ struct PnnsProcessBenchmarkContext<Scheme: HeScheme> {
 
         self.database = getDatabaseForTesting(config: databaseConfig)
         self.contexts = try serverConfig.encryptionParameters.map { encryptionParameters in
-            try Context(encryptionParameters: encryptionParameters)
+            try Scheme.Context(encryptionParameters: encryptionParameters)
         }
     }
 }
@@ -305,7 +305,7 @@ struct PnnsBenchmarkContext<Scheme: HeScheme> {
 
         let database = getDatabaseForTesting(config: databaseConfig)
         let contexts = try clientConfig.encryptionParameters
-            .map { encryptionParameters in try Context<Scheme>(encryptionParameters: encryptionParameters) }
+            .map { encryptionParameters in try Scheme.Context(encryptionParameters: encryptionParameters) }
         self.processedDatabase = try database.process(config: serverConfig, contexts: contexts)
         self.client = try Client(config: clientConfig, contexts: contexts)
         self.server = try Server(database: processedDatabase)

Sources/HomomorphicEncryption/Bfv/Bfv+Decrypt.swift

@@ -26,8 +26,8 @@ extension Bfv {
 
     @inlinable
     // swiftlint:disable:next missing_docs attributes
-    public static func decryptEval(_ ciphertext: EvalCiphertext,
-                                   using secretKey: SecretKey<Bfv<T>>) throws -> CoeffPlaintext
+    public static func decryptEval(_ ciphertext: EvalCiphertext, using secretKey: SecretKey<Self>) throws ->
+        CoeffPlaintext
     {
         let t = ciphertext.context.plaintextModulus
         let dotProduct = try Self.dotProduct(ciphertext: ciphertext, with: secretKey)

Sources/HomomorphicEncryption/Bfv/Bfv+Encode.swift

@@ -24,23 +24,23 @@ extension Bfv {
 
     @inlinable
     // swiftlint:disable:next missing_docs attributes
-    public static func encode(context: Context<Bfv<T>>, values: some Collection<Scalar>,
+    public static func encode(context: Context, values: some Collection<Scalar>,
                               format: EncodeFormat) throws -> CoeffPlaintext
     {
         try context.encode(values: values, format: format)
     }
 
     @inlinable
     // swiftlint:disable:next missing_docs attributes
-    public static func encode(context: Context<Bfv<T>>, signedValues: some Collection<SignedScalar>,
+    public static func encode(context: Context, signedValues: some Collection<SignedScalar>,
                               format: EncodeFormat) throws -> CoeffPlaintext
     {
         try context.encode(signedValues: signedValues, format: format)
     }
 
     @inlinable
     // swiftlint:disable:next missing_docs attributes
-    public static func encode(context: Context<Bfv<T>>, values: some Collection<Scalar>, format: EncodeFormat,
+    public static func encode(context: Context, values: some Collection<Scalar>, format: EncodeFormat,
                               moduliCount: Int?) throws -> EvalPlaintext
     {
         let coeffPlaintext = try Self.encode(context: context, values: values, format: format)
@@ -50,7 +50,7 @@ extension Bfv {
     @inlinable
     // swiftlint:disable:next missing_docs attributes
     public static func encode(
-        context: Context<Bfv<T>>,
+        context: Context,
         signedValues: some Collection<SignedScalar>,
         format: EncodeFormat,
         moduliCount: Int?) throws -> EvalPlaintext

Sources/HomomorphicEncryption/Bfv/Bfv+Encrypt.swift

@@ -1,4 +1,4 @@
-// Copyright 2024 Apple Inc. and the Swift Homomorphic Encryption project authors
+// 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.
@@ -23,7 +23,7 @@ extension Bfv {
 
     @inlinable
     // swiftlint:disable:next missing_docs attributes
-    public static func zeroCiphertextCoeff(context: Context<Self>, moduliCount: Int?) throws -> CoeffCiphertext {
+    public static func zeroCiphertextCoeff(context: Context, moduliCount: Int?) throws -> CoeffCiphertext {
         let moduliCount = moduliCount ?? context.ciphertextContext.moduli.count
         let zeroPoly = try PolyRq<Scalar, Coeff>.zero(
             context: context.ciphertextContext
@@ -34,7 +34,7 @@ extension Bfv {
 
     @inlinable
     // swiftlint:disable:next missing_docs attributes
-    public static func zeroCiphertextEval(context: Context<Self>, moduliCount: Int?) throws -> EvalCiphertext {
+    public static func zeroCiphertextEval(context: Context, moduliCount: Int?) throws -> EvalCiphertext {
         let moduliCount = moduliCount ?? context.ciphertextContext.moduli.count
         let zeroPoly = try PolyRq<Scalar, Eval>.zero(
             context: context.ciphertextContext
@@ -143,15 +143,15 @@ extension Bfv {
     }
 
     @inlinable
-    static func encryptZero(for context: Context<Bfv<T>>,
+    static func encryptZero(for context: Context,
                             using secretKey: SecretKey<Bfv<T>>) throws -> CanonicalCiphertext
     {
         let ciphertextContext = context.ciphertextContext
         return try encryptZero(for: context, using: secretKey, with: ciphertextContext)
     }
 
     @inlinable
-    static func encryptZero(for context: Context<Bfv<T>>,
+    static func encryptZero(for context: Context,
                             using secretKey: SecretKey<Bfv<T>>,
                             with ciphertextContext: PolyContext<T>) throws -> CanonicalCiphertext
     {

Sources/HomomorphicEncryption/Bfv/Bfv+Keys.swift

@@ -1,4 +1,4 @@
-// Copyright 2024 Apple Inc. and the Swift Homomorphic Encryption project authors
+// 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.
@@ -17,7 +17,7 @@ import ModularArithmetic
 extension Bfv {
     @inlinable
     // swiftlint:disable:next missing_docs attributes
-    public static func generateSecretKey(context: Context<Self>) throws -> SecretKey<Bfv<T>> {
+    public static func generateSecretKey(context: Context) throws -> SecretKey<Bfv<T>> {
         var s = PolyRq<Scalar, Coeff>.zero(context: context.secretKeyContext)
         var rng = SystemRandomNumberGenerator()
         s.randomizeTernary(using: &rng)
@@ -28,7 +28,7 @@ extension Bfv {
     @inlinable
     // swiftlint:disable:next missing_docs attributes
     public static func generateEvaluationKey(
-        context: Context<Bfv<T>>,
+        context: Context,
         config: EvaluationKeyConfig,
         using secretKey: borrowing SecretKey<Bfv<T>>) throws -> EvaluationKey<Bfv<T>>
     {
@@ -55,7 +55,7 @@ extension Bfv {
     }
 
     @inlinable
-    static func generateRelinearizationKey(context: Context<Self>,
+    static func generateRelinearizationKey(context: Context,
                                            secretKey: borrowing SecretKey<Self>) throws
         -> RelinearizationKey<Self>
     {
@@ -65,7 +65,7 @@ extension Bfv {
     }
 
     @inlinable
-    static func generateKeySwitchKey(context: Context<Bfv<T>>,
+    static func generateKeySwitchKey(context: Context,
                                      currentKey: consuming PolyRq<T, Eval>,
                                      targetKey: borrowing SecretKey<Bfv<T>>) throws -> KeySwitchKey<Bfv<T>>
     {
@@ -120,7 +120,7 @@ extension Bfv {
     /// - seealso: ``Bfv/generateEvaluationKey(context:config:using:)``.
     @inlinable
     static func computeKeySwitchingUpdate(
-        context: Context<Bfv<T>>,
+        context: Context,
         target: PolyRq<Scalar, CanonicalCiphertextFormat>,
         keySwitchingKey: KeySwitchKey<Self>) throws -> [PolyRq<Scalar, CanonicalCiphertextFormat>]
     {

Sources/HomomorphicEncryption/Bfv/Bfv+Multiply.swift

@@ -1,4 +1,4 @@
-// Copyright 2024 Apple Inc. and the Swift Homomorphic Encryption project authors
+// 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.

Sources/HomomorphicEncryption/Bfv/Bfv.swift

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