Implement PlaintextMatrix.diagonal encoding

Package.resolved

@@ -1,5 +1,5 @@
 {
-  "originHash" : "7ee1d955f789a2932c73eb859ad423a54265140cebf530f47f84bc523d26c86f",
+  "originHash" : "287ba43d965ab226c778e6332f9fbb9244920f2452fbe9bccef6f422c30d7e49",
   "pins" : [
     {
       "identity" : "hdrhistogram-swift",
@@ -28,6 +28,15 @@
         "version" : "1.0.0"
       }
     },
+    {
+      "identity" : "swift-algorithms",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/apple/swift-algorithms",
+      "state" : {
+        "revision" : "f6919dfc309e7f1b56224378b11e28bab5bccc42",
+        "version" : "1.2.0"
+      }
+    },
     {
       "identity" : "swift-argument-parser",
       "kind" : "remoteSourceControl",

Package.swift

@@ -51,6 +51,7 @@ let package = Package(
         .executable(name: "PIRShardDatabase", targets: ["PIRShardDatabase"]),
     ],
     dependencies: [
+        .package(url: "https://github.com/apple/swift-algorithms", from: "1.2.0"),
         .package(url: "https://github.com/apple/swift-argument-parser.git", from: "1.2.0"),
         .package(url: "https://github.com/apple/swift-crypto.git", from: "3.4.0"),
         .package(url: "https://github.com/apple/swift-log.git", from: "1.0.0"),
@@ -96,7 +97,10 @@ let package = Package(
             swiftSettings: librarySettings),
         .target(
             name: "PrivateNearestNeighborsSearch",
-            dependencies: ["HomomorphicEncryption"],
+            dependencies: [
+                .product(name: "Algorithms", package: "swift-algorithms"),
+                "HomomorphicEncryption",
+            ],
             swiftSettings: librarySettings),
         .target(
             name: "TestUtilities",

Sources/HomomorphicEncryption/Array2d.swift

@@ -15,11 +15,11 @@
 /// Stores values in a 2 dimensional array.
 public struct Array2d<T: Equatable & AdditiveArithmetic & Sendable>: Equatable, Sendable {
     @usableFromInline package var data: [T]
-    @usableFromInline var rowCount: Int
-    @usableFromInline var columnCount: Int
+    @usableFromInline package var rowCount: Int
+    @usableFromInline package var columnCount: Int
 
-    @usableFromInline var shape: (Int, Int) { (rowCount, columnCount) }
-    @usableFromInline var count: Int { rowCount * columnCount }
+    @usableFromInline package var shape: (Int, Int) { (rowCount, columnCount) }
+    @usableFromInline package var count: Int { rowCount * columnCount }
 
     @inlinable
     package init(data: [T], rowCount: Int, columnCount: Int) {
@@ -35,26 +35,34 @@ public struct Array2d<T: Equatable & AdditiveArithmetic & Sendable>: Equatable,
         self.rowCount = array.rowCount
         self.data = array.data.map { T($0) }
     }
+
+    @inlinable
+    package static func zero(rowCount: Int, columnCount: Int) -> Self {
+        self.init(
+            data: [T](Array(repeating: T.zero, count: rowCount * columnCount)),
+            rowCount: rowCount,
+            columnCount: columnCount)
+    }
 }
 
 extension Array2d {
     @inlinable
-    func index(row: Int, column: Int) -> Int {
+    package func index(row: Int, column: Int) -> Int {
         row &* columnCount &+ column
     }
 
     @inlinable
-    func rowIndices(row: Int) -> Range<Int> {
+    package func rowIndices(row: Int) -> Range<Int> {
         index(row: row, column: 0)..<index(row: row, column: columnCount)
     }
 
     @inlinable
-    func columnIndices(column: Int) -> StrideTo<Int> {
+    package func columnIndices(column: Int) -> StrideTo<Int> {
         stride(from: index(row: 0, column: column), to: index(row: rowCount, column: column), by: columnCount)
     }
 
     @inlinable
-    func row(row: Int) -> [T] {
+    package func row(row: Int) -> [T] {
         Array(data[rowIndices(row: row)])
     }
 
@@ -80,7 +88,7 @@ extension Array2d {
     }
 
     @inlinable
-    subscript(_ index: Int) -> T {
+    package subscript(_ index: Int) -> T {
         get {
             data[index]
         }
@@ -90,7 +98,7 @@ extension Array2d {
     }
 
     @inlinable
-    subscript(_ row: Int, _ column: Int) -> T {
+    package subscript(_ row: Int, _ column: Int) -> T {
         get {
             data[index(row: row, column: column)]
         }

Sources/HomomorphicEncryption/Encoding.swift

@@ -169,16 +169,11 @@ extension Context {
     func encodeSimd(values: [some ScalarType]) throws -> Plaintext<Scheme, Coeff> {
         guard !simdEncodingMatrix.isEmpty else { throw HeError.simdEncodingNotSupported(for: encryptionParameters) }
         let polyDegree = encryptionParameters.polyDegree
-        var array = Array2d<Scheme.Scalar>(
-            data: [Scheme.Scalar](repeating: 0,
-                                  count: polyDegree),
-            rowCount: 1,
-            columnCount: polyDegree)
+        var array = Array2d<Scheme.Scalar>.zero(rowCount: 1, columnCount: polyDegree)
         for index in 0..<values.count {
             array[0, simdEncodingMatrix[index]] = Scheme.Scalar(values[index])
         }
-        let poly = PolyRq<_, Eval>(context: plaintextContext,
-                                   data: array)
+        let poly = PolyRq<_, Eval>(context: plaintextContext, data: array)
         let coeffPoly = try poly.inverseNtt()
         return Plaintext<Scheme, Coeff>(context: self, poly: coeffPoly)
     }
@@ -189,10 +184,8 @@ extension Context {
             throw HeError.simdEncodingNotSupported(for: encryptionParameters)
         }
         let poly = try plaintext.poly.forwardNtt()
-        var values = [T](repeating: 0, count: encryptionParameters.polyDegree)
-        for index in 0..<encryptionParameters.polyDegree {
-            values[index] = T(poly.data[0, simdEncodingMatrix[index]])
+        return (0..<encryptionParameters.polyDegree).map { index in
+            T(poly.data[0, simdEncodingMatrix[index]])
         }
-        return values
     }
 }

Sources/HomomorphicEncryption/Scalar.swift

@@ -371,22 +371,40 @@ extension FixedWidthInteger {
         return self / divisor
     }
 
-    /// Computes the next value larger than or equal to this value that is a multiple of `factor`.
+    /// Computes the smallest value greater than or equal to this value that is a multiple of `rhs`.
     ///
     /// This value must be non-negative.
     /// - Parameters:
-    ///   - factor: Value of which the output is a product of.
+    ///   - rhs: Value of which the output is a multiple of.
     ///   - variableTime: Must be `true`, indicating this value and `other` are leaked through timing.
     /// - Returns: the next multiple of this value.
     /// - Warning: Leaks this value and `other` through timing.
     @inlinable
-    public func nextMultiple(of factor: Self, variableTime: Bool) -> Self {
+    public func nextMultiple(of rhs: Self, variableTime: Bool) -> Self {
         precondition(variableTime)
         precondition(self >= 0)
-        if factor == 0 {
+        if rhs == 0 {
             return 0
         }
-        return dividingCeil(factor, variableTime: true) * factor
+        return dividingCeil(rhs, variableTime: true) * rhs
+    }
+
+    /// Computes the largest value less than or equal to this value that is a multiple of `rhs`.
+    ///
+    /// This value must be non-negative.
+    /// - Parameters:
+    ///   - rhs: Value of which the output is a multiple of.
+    ///   - variableTime: Must be `true`, indicating this value and `other` are leaked through timing.
+    /// - Returns: the previous multiple of this value.
+    /// - Warning: Leaks this value and `other` through timing.
+    @inlinable
+    public func previousMultiple(of rhs: Self, variableTime: Bool) -> Self {
+        precondition(variableTime)
+        precondition(self >= 0)
+        if rhs == 0 {
+            return 0
+        }
+        return (self / rhs) * rhs
     }
 }
 

Sources/PrivateNearestNeighborsSearch/DotProduct.swift

@@ -0,0 +1,41 @@
+// Copyright 2024 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 Foundation
+
+/// Pre-computed values for matrix-vector multiplication using baby-step, giant-step algorithm.
+///
+/// - seealso: Section 6.3 of <https://eprint.iacr.org/2018/244.pdf>.
+struct BabyStepGiantStep: Codable, Equatable, Hashable, Sendable {
+    /// Dimension of the vector; "D" in the reference.
+    let vectorDimension: Int
+    /// Baby step; "g" in the reference.
+    let babyStep: Int
+    /// Giant step; "h" in the reference.
+    let giantStep: Int
+
+    init(vectorDimension: Int, babyStep: Int, giantStep: Int) {
+        self.vectorDimension = vectorDimension
+        self.babyStep = babyStep
+        self.giantStep = giantStep
+    }
+
+    init(vectorDimension: Int) {
+        let dimension = Int32(vectorDimension).nextPowerOfTwo
+        let babyStep = Int32(Double(dimension).squareRoot().rounded(.up))
+        let giantStep = dimension.dividingCeil(babyStep, variableTime: true)
+
+        self.init(vectorDimension: Int(dimension), babyStep: Int(babyStep), giantStep: Int(giantStep))
+    }
+}

Sources/PrivateNearestNeighborsSearch/PlaintextMatrix.swift

@@ -12,10 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+import Algorithms
 import HomomorphicEncryption
 
 /// Different algorithms for packing a matrix of scalar values into plaintexts.
-enum PlaintextMatrixPacking: Equatable, Sendable {
+enum PlaintextMatrixPacking: Codable, Equatable, Hashable, Sendable {
     /// As many columns of data are packed sequentially into each plaintext SIMD row as possible, such that no SIMD row
     /// contains data from multiple columns.
     case denseColumn
@@ -27,6 +28,11 @@ enum PlaintextMatrixPacking: Equatable, Sendable {
     /// with the constraint that either all or none of the entries stored within the last plaintext
     /// row are repeated.
     case denseRow
+    /// Packs the values using a generalized diagonal packing.
+    ///
+    /// Includes modifications for the baby-step, giant-step algorithm from Section 6.3 of
+    /// <https://eprint.iacr.org/2018/244.pdf>.
+    case diagonal(babyStepGiantStep: BabyStepGiantStep)
 }
 
 /// The dimensions of a matrix, a 2d array.
@@ -156,6 +162,13 @@ struct PlaintextMatrix<Scheme: HeScheme, Format: PolyFormat>: Equatable, Sendabl
                 dimensions: dimensions,
                 values: values)
             try self.init(dimensions: dimensions, packing: packing, plaintexts: plaintexts)
+        case .diagonal:
+            let plaintexts = try PlaintextMatrix.diagonalPlaintexts(
+                context: context,
+                dimensions: dimensions,
+                packing: packing,
+                values: values)
+            try self.init(dimensions: dimensions, packing: packing, plaintexts: plaintexts)
         }
     }
 
@@ -190,6 +203,11 @@ struct PlaintextMatrix<Scheme: HeScheme, Format: PolyFormat>: Equatable, Sendabl
             let rowsPerPlaintextCount = simdDimensions.rowCount * (
                 simdDimensions.columnCount / dimensions.columnCount.nextPowerOfTwo)
             return dimensions.rowCount.dividingCeil(rowsPerPlaintextCount, variableTime: true)
+        case .diagonal:
+            let plaintextsPerColumnCount = dimensions.rowCount.dividingCeil(
+                encryptionParameters.polyDegree,
+                variableTime: true)
+            return dimensions.columnCount.nextPowerOfTwo * plaintextsPerColumnCount
         }
     }
 
@@ -322,6 +340,83 @@ struct PlaintextMatrix<Scheme: HeScheme, Format: PolyFormat>: Equatable, Sendabl
         return plaintexts
     }
 
+    /// Computes the plaintexts for diagonal packing.
+    /// - Parameters:
+    ///   - context: Context for HE computation.
+    ///   - dimensions: Plaintext matrix dimensions.
+    ///   - packing: Plaintext packing; must be `.diagonal`.
+    ///   - values: The data values to store in the plaintext matrix; stored in row-major format.
+    /// - Returns: The plaintexts for diagonal packing.
+    /// - Throws: Error upon failure to compute the plaintexts.
+    @inlinable
+    static func diagonalPlaintexts<V: ScalarType>(
+        context: Context<Scheme>,
+        dimensions: Dimensions,
+        packing: PlaintextMatrixPacking,
+        values: [V]) throws -> [Scheme.CoeffPlaintext]
+    {
+        let encryptionParameters = context.encryptionParameters
+        guard let simdDimensions = context.simdDimensions else {
+            throw PNNSError.simdEncodingNotSupported(for: encryptionParameters)
+        }
+        let simdColumnCount = simdDimensions.columnCount
+        let simdRowCount = simdDimensions.rowCount
+        precondition(simdRowCount == 2, "simdRowCount must be 2")
+        guard dimensions.columnCount <= simdColumnCount else {
+            throw PNNSError.invalidMatrixDimensions(dimensions)
+        }
+        guard case let .diagonal(bsgs) = packing else {
+            let expectedBsgs = BabyStepGiantStep(vectorDimension: dimensions.columnCount)
+            throw PNNSError
+                .wrongPlaintextMatrixPacking(got: packing, expected: .diagonal(babyStepGiantStep: expectedBsgs))
+        }
+
+        let data = Array2d(data: values, rowCount: dimensions.rowCount, columnCount: dimensions.columnCount)
+        // Transposed from original shape, with extra zero columns.
+        // Encode diagonals
+        var packedValues = Array2d<V>.zero(
+            rowCount: dimensions.columnCount.nextPowerOfTwo,
+            columnCount: dimensions.rowCount)
+        for rowIndex in 0..<packedValues.rowCount {
+            for columnIndex in 0..<packedValues.columnCount {
+                let paddedColumnIndex = (columnIndex &+ rowIndex) % packedValues.rowCount
+                if paddedColumnIndex < dimensions.columnCount {
+                    packedValues[rowIndex, columnIndex] =
+                        data[columnIndex, paddedColumnIndex]
+                }
+            }
+        }
+
+        var plaintexts: [Scheme.CoeffPlaintext] = []
+        let expectedPlaintextCount = try PlaintextMatrix.plaintextCount(
+            encryptionParameters: encryptionParameters,
+            dimensions: dimensions,
+            packing: packing)
+        plaintexts.reserveCapacity(expectedPlaintextCount)
+        let plaintextsPerColumn = expectedPlaintextCount / packedValues.rowCount
+
+        // Perform baby-step giant-step rotations.
+        // See Section 6.3 of <https://eprint.iacr.org/2018/244.pdf>.
+        let n = context.degree
+        for rowIndex in 0..<packedValues.rowCount {
+            let row = packedValues.row(row: rowIndex)
+            for (chunkIndex, var chunk) in row.chunks(ofCount: n).enumerated() {
+                chunk += repeatElement(0, count: n - chunk.count)
+                let i = (plaintexts.count - chunkIndex) / plaintextsPerColumn
+                let rotationStep = i.previousMultiple(of: bsgs.babyStep, variableTime: true)
+                let middle = min(chunk.endIndex, chunk.startIndex + n / 2)
+                chunk[chunk.startIndex..<middle].rotate(toStartAt: chunk.startIndex + rotationStep)
+                chunk[middle...].rotate(toStartAt: middle + rotationStep)
+
+                let plaintext = try context.encode(values: Array(chunk), format: .simd)
+                plaintexts.append(plaintext)
+            }
+        }
+        precondition(plaintexts.count == expectedPlaintextCount)
+
+        return plaintexts
+    }
+
     /// Unpacks the plaintext matrix.
     /// - Returns: The stored data values in row-major format.
     /// - Throws: Error upon failure to unpack the matrix.
@@ -332,6 +427,9 @@ struct PlaintextMatrix<Scheme: HeScheme, Format: PolyFormat>: Equatable, Sendabl
             return try unpackDenseColumn()
         case .denseRow:
             return try unpackDenseRow()
+        case .diagonal:
+            // TODO: Implement
+            preconditionFailure("Unpacking diagonal plaintext matrix not supported")
         }
     }
 

Tests/HomomorphicEncryptionTests/Array2dTests.swift

@@ -16,7 +16,7 @@
 import XCTest
 
 class Array2dTests: XCTestCase {
-    func testZeroize() {
+    func testZeroAndZeroize() {
         func runTest<T: FixedWidthInteger & Sendable>(_: T.Type) {
             let data = [T](1...16)
             var array = Array2d(data: data, rowCount: 2, columnCount: 8)
@@ -27,6 +27,7 @@ class Array2dTests: XCTestCase {
                 rowCount: 2,
                 columnCount: 8)
             XCTAssertEqual(array, zero)
+            XCTAssertEqual(array, Array2d.zero(rowCount: 2, columnCount: 8))
         }
         runTest(Int.self)
         runTest(Int32.self)