Add PolyRq/removeLastRows

Package.swift

@@ -192,6 +192,11 @@ let package = Package(
                 "HomomorphicEncryption", "_TestUtilities",
                 .product(name: "Numerics", package: "swift-numerics"),
             ], swiftSettings: executableSettings),
+        .testTarget(
+            name: "HomomorphicEncryptionExtrasTests",
+            dependencies: [
+                "_HomomorphicEncryptionExtras",
+            ], swiftSettings: executableSettings),
         .testTarget(
             name: "HomomorphicEncryptionProtobufTests",
             dependencies: [

Sources/HomomorphicEncryption/Array2d.swift

@@ -189,7 +189,7 @@ extension Array2d {
     /// - Parameter k: The number of rows to remove. `k` Must be greater than or equal to zero and must not exceed
     /// `rowCount`.
     @inlinable
-    mutating func removeLastRows(_ k: Int) {
+    package mutating func removeLastRows(_ k: Int) {
         precondition(k >= 0 && k <= rowCount)
         rowCount -= k
         data.removeLast(columnCount * k)

Sources/HomomorphicEncryption/Error.swift

@@ -129,7 +129,7 @@ extension HeError {
     }
 
     @inlinable
-    static func invalidPolyContext(_ context: PolyContext<some Any>) -> Self {
+    package static func invalidPolyContext(_ context: PolyContext<some Any>) -> Self {
         .invalidPolyContext(context.description)
     }
 

Sources/HomomorphicEncryption/PolyRq/PolyContext.swift

@@ -24,7 +24,7 @@ public final class PolyContext<T: ScalarType>: Sendable {
     /// The modulus `Q = product_{i=0}^{L-1} q_i`, if representable by a `Width32<T>`
     @usableFromInline let modulus: Width32<T>?
     /// Next context, typically formed by dropping `q_{L-1}`.
-    @usableFromInline let next: PolyContext<T>?
+    @usableFromInline package let next: PolyContext<T>?
     /// Operations mod `q_0` up to `q_{L-1}`.
     @usableFromInline let reduceModuli: [Modulus<T>]
     /// Operations mod `UInt64(q_0), ..., UInt64(q_{L-1})`.

Sources/_HomomorphicEncryptionExtras/PolyRq.swift

@@ -0,0 +1,36 @@
+// Copyright 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 PolyRq {
+    /// Removes the last k rows of coefficients and drop the last k moduli from the context.
+    ///
+    ///  - Parameter k: The number of moduli to drop. It must be greater than or equal to zero and less than the number
+    /// of moduli in the context.
+    /// - Throws: Error upon failure to drop context.
+    @inlinable
+    public mutating func removeLastModuli(_ k: Int) throws {
+        precondition(k >= 0 && k < moduli.count)
+        var context = context
+        for _ in 0..<k {
+            guard let newContext = context.next else {
+                throw HeError.invalidPolyContext(context)
+            }
+            context = newContext
+        }
+        self.context = context
+        data.removeLastRows(k)
+    }
+}

Tests/HomomorphicEncryptionExtrasTests/PolyRqTests/PolyRqTests.swift

@@ -0,0 +1,36 @@
+// Copyright 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 _HomomorphicEncryptionExtras
+import HomomorphicEncryption
+import Testing
+
+@Suite
+struct PolyRqTests {
+    @Test
+    func removeLastModuli() async throws {
+        let context: PolyContext<UInt32> = try PolyContext(degree: 4, moduli: [2, 3, 5])
+        let data: [UInt32] = [0, 1, 0, 1,
+                              0, 1, 2, 0,
+                              0, 1, 2, 3]
+        var poly = PolyRq<_, Coeff>(
+            context: context,
+            data: Array2d(data: data, rowCount: 3, columnCount: 4))
+        #expect(poly.moduli == [2, 3, 5])
+        try poly.removeLastModuli(1)
+        #expect(poly.moduli == [2, 3])
+        try poly.removeLastModuli(1)
+        #expect(poly.moduli == [2])
+    }
+}