Make saving ProcessedDatabaseWithParameters easier

Sources/PIRProcessDatabase/ProcessDatabase.swift

@@ -301,7 +301,6 @@ struct ProcessDatabase: ParsableCommand {
         let keywordConfig = try KeywordPirConfig(dimensionCount: 2,
                                                  cuckooTableConfig: config.cuckooTableConfig,
                                                  unevenDimensions: true)
-        let keywordPirParams = keywordConfig.parameter.proto()
         let databaseConfig = KeywordDatabaseConfig(
             sharding: config.sharding,
             keywordPirConfig: keywordConfig)
@@ -335,7 +334,6 @@ struct ProcessDatabase: ParsableCommand {
                 ProcessDatabase.logger.info("ValidationResults \(description)")
             }
 
-            let shardConfig = processed.pirParameter.proto(shardID: shardID)
             let outputDatabaseFilename = config.outputDatabase.replacingOccurrences(
                 of: "SHARD_ID",
                 with: String(shardID))
@@ -345,16 +343,7 @@ struct ProcessDatabase: ParsableCommand {
             let shardEvaluationKeyConfig = processed.evaluationKeyConfiguration
             evaluationKeyConfig = [evaluationKeyConfig, shardEvaluationKeyConfig].union()
 
-            let shardPirParameters = try Apple_SwiftHomomorphicEncryption_Pir_V1_PirParameters.with { params in
-                params.encryptionParameters = try encryptionParameters.proto()
-                params.numEntries = shardConfig.numEntries
-                params.entrySize = shardConfig.entrySize
-                params.dimensions = shardConfig.dimensions
-                params.batchSize = UInt64(processed.pirParameter.batchSize)
-                params.keywordPirParams = keywordPirParams
-                params.evaluationKeyConfig = try shardEvaluationKeyConfig
-                    .proto(encryptionParameters: encryptionParameters)
-            }
+            let shardPirParameters = try processed.proto(context: context)
             let outputParametersFilename = config.outputPirParameters.replacingOccurrences(
                 of: "SHARD_ID",
                 with: String(shardID))

Sources/PrivateInformationRetrieval/IndexPirProtocol.swift

@@ -148,7 +148,7 @@ public struct ProcessedDatabase<Scheme: HeScheme>: Equatable, Sendable {
     ///   - buffer: Serialized plaintexts.
     ///   - context: Context for HE computation.
     /// - Throws: Error upon failure to deserialize.
-    init(from buffer: [UInt8], context: Context<Scheme>) throws {
+    public init(from buffer: [UInt8], context: Context<Scheme>) throws {
         var offset = buffer.startIndex
         let versionNumber = buffer[offset]
         offset += MemoryLayout<SerializationVersionType>.size
@@ -181,9 +181,9 @@ public struct ProcessedDatabase<Scheme: HeScheme>: Equatable, Sendable {
         self.init(plaintexts: plaintexts)
     }
 
-    /// Returns the serilization size in bytes of the database.
+    /// Returns the serialization size in bytes of the database.
     @inlinable
-    func serializationByteCount() throws -> Int {
+    public func serializationByteCount() throws -> Int {
         let nonNilPlaintexts = plaintexts.compactMap { $0 }
         guard let polyContext = nonNilPlaintexts.first?.polyContext() else {
             throw PirError.emptyDatabase
@@ -208,7 +208,7 @@ public struct ProcessedDatabase<Scheme: HeScheme>: Equatable, Sendable {
     /// - Returns: The serialized database.
     /// - Throws: Error upon failure to serialize the database.
     @inlinable
-    func serialize() throws -> [UInt8] {
+    public func serialize() throws -> [UInt8] {
         var buffer: [UInt8] = []
         try buffer.reserveCapacity(serializationByteCount())
         buffer.append(Self.serializationVersion)
@@ -230,6 +230,8 @@ public struct ProcessedDatabase<Scheme: HeScheme>: Equatable, Sendable {
 public struct ProcessedDatabaseWithParameters<Scheme: HeScheme>: Sendable {
     /// Processed database.
     public let database: ProcessedDatabase<Scheme>
+    /// The algorithm that this database was processed for.
+    public let algorithm: PirAlgorithm
     /// Evaluation key configuration.
     public let evaluationKeyConfiguration: EvaluationKeyConfiguration
     /// Parameters for Index PIR queries.
@@ -240,16 +242,19 @@ public struct ProcessedDatabaseWithParameters<Scheme: HeScheme>: Sendable {
     /// Initializes a ``ProcessedDatabaseWithParameters``.
     /// - Parameters:
     ///   - database: Processed database.
+    ///   - algorithm: The PIR algorithm used.
     ///   - evaluationKeyConfiguration: Evaluation key configuration.
     ///   - pirParameter: Index PIR parameters.
     ///   - keywordPirParameter: Optional keyword PIR parameters.
     public init(
         database: ProcessedDatabase<Scheme>,
+        algorithm: PirAlgorithm,
         evaluationKeyConfiguration: EvaluationKeyConfiguration,
         pirParameter: IndexPirParameter,
         keywordPirParameter: KeywordPirParameter? = nil)
     {
         self.database = database
+        self.algorithm = algorithm
         self.evaluationKeyConfiguration = evaluationKeyConfiguration
         self.pirParameter = pirParameter
         self.keywordPirParameter = keywordPirParameter
@@ -296,6 +301,9 @@ public protocol IndexPirProtocol {
     /// Encrypted server response type.
     typealias Response = PrivateInformationRetrieval.Response<Scheme>
 
+    /// The PIR algorithm.
+    static var algorithm: PirAlgorithm { get }
+
     /// Generates the PIR parameters for a database.
     /// - Parameters:
     ///   - config: Database configuration.

Sources/PrivateInformationRetrieval/KeywordDatabase.swift

@@ -44,15 +44,15 @@ extension KeywordValuePair.Keyword {
     ///  - shardCount: The shard count.
     /// - Returns: The shard identifier.
     @inlinable
-    func shardID(shardCount: Int) -> String {
+    public func shardID(shardCount: Int) -> String {
         String(shardIndex(shardCount: shardCount))
     }
 
     /// Returns the shard index for the given shard count.
     /// - Parameter shardCount: The shard count.
     /// - Returns: The shard index.
     @inlinable
-    func shardIndex(shardCount: Int) -> Int {
+    public func shardIndex(shardCount: Int) -> Int {
         let digest = SHA256.hash(data: self)
         let truncatedHash = digest.withUnsafeBytes { buffer in
             buffer.load(as: UInt64.self)
@@ -163,7 +163,7 @@ public struct KeywordDatabaseShard: Hashable, Codable {
     ///   - shardID: Identifier for the database shard.
     ///   - rows: Rows in the database.
     @inlinable
-    init(
+    public init(
         shardID: String,
         rows: some Collection<(KeywordValuePair.Keyword, KeywordValuePair.Value)>)
     {

Sources/PrivateInformationRetrieval/KeywordPirProtocol.swift

@@ -199,6 +199,7 @@ public final class KeywordPirServer<PirServer: IndexPirServer>: KeywordPirProtoc
 
         return ProcessedDatabaseWithParameters(
             database: processedDb,
+            algorithm: PirServer.IndexPir.algorithm,
             evaluationKeyConfiguration: evaluationKeyConfig,
             pirParameter: indexPirParameter,
             keywordPirParameter: config.parameter)

Sources/PrivateInformationRetrieval/MulPir.swift

@@ -27,6 +27,8 @@ public enum MulPir<Scheme: HeScheme>: IndexPirProtocol {
     public typealias Response = PrivateInformationRetrieval.Response<Scheme>
     @usableFromInline typealias CanonicalCiphertext = Scheme.CanonicalCiphertext
 
+    public static var algorithm: PirAlgorithm { .mulPir }
+
     public static func generateParameter(config: IndexPirConfig, with context: Context<Scheme>) -> IndexPirParameter {
         let entrySizeInBytes = config.entrySizeInBytes
         let perChunkPlaintextCount = if entrySizeInBytes <= context.bytesPerPlaintext {

Sources/PrivateInformationRetrievalProtobuf/ConversionPir.swift

@@ -45,6 +45,30 @@ extension Query {
     }
 }
 
+extension ProcessedDatabaseWithParameters {
+    /// Converts the native processed database with parameters into protobuf object that contains only the parameters.
+    /// - Parameter context: The context that was used to create processed database.
+    /// - Returns: The PIR parameters protobuf object.
+    /// - Throws: Error when the parameters cannot be represented as a protobuf object.
+    public func proto(context: Context<Scheme>) throws -> Apple_SwiftHomomorphicEncryption_Pir_V1_PirParameters {
+        let encryptionParameters = context.encryptionParameters
+        return try Apple_SwiftHomomorphicEncryption_Pir_V1_PirParameters.with { params in
+            params.encryptionParameters = try encryptionParameters.proto()
+            params.numEntries = UInt64(pirParameter.entryCount)
+            params.entrySize = UInt64(pirParameter.entrySizeInBytes)
+            params.dimensions = pirParameter.dimensions.map(UInt64.init)
+            if let keywordPirParameter {
+                params.keywordPirParams = keywordPirParameter.proto()
+            }
+            params.algorithm = algorithm.proto()
+            params.batchSize = UInt64(pirParameter.batchSize)
+            params.evaluationKeyConfig = try evaluationKeyConfiguration
+                .proto(encryptionParameters: encryptionParameters)
+            params.keyCompressionStrategy = .unspecified
+        }
+    }
+}
+
 extension Apple_SwiftHomomorphicEncryption_Pir_V1_KeywordPirParameters {
     /// Converts the protobuf object to a native type.
     /// - Returns: The converted native type.