Fix sharding

There was a mistake in sharding. I added a test and made the sharding
compatible with the device code.

Sources/PrivateInformationRetrieval/KeywordDatabase.swift

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+import CryptoKit
 import Foundation
 import HomomorphicEncryption
 
@@ -52,11 +53,11 @@ extension KeywordValuePair.Keyword {
     /// - Returns: The shard index.
     @inlinable
     func shardIndex(shardCount: Int) -> Int {
-        HashKeyword
-            .indexFromHash(
-                keywordHash: HashKeyword.hash(keyword: self),
-                bucketCount: shardCount,
-                counter: 0)
+        let digest = SHA256.hash(data: self)
+        let truncatedHash = digest.withUnsafeBytes { buffer in
+            buffer.load(as: UInt64.self)
+        }
+        return Int(truncatedHash % UInt64(shardCount))
     }
 }
 

Tests/PrivateInformationRetrievalTests/KeywordDatabaseTests.swift

@@ -39,4 +39,15 @@ class KeywordDatabaseTests: XCTestCase {
             XCTAssert(database.shards.contains { shard in shard.value[row.keyword] == row.value })
         }
     }
+
+    func testShardingKnownAnswerTest() throws {
+        func checkKeywordShard(_ keyword: KeywordValuePair.Keyword, shardCount: Int, expectedShard: Int) {
+            XCTAssertEqual(keyword.shardIndex(shardCount: shardCount), expectedShard)
+        }
+
+        checkKeywordShard([0, 0, 0, 0], shardCount: 41, expectedShard: 2)
+        checkKeywordShard([0, 0, 0, 0], shardCount: 1001, expectedShard: 635)
+        checkKeywordShard([1, 2, 3], shardCount: 1001, expectedShard: 903)
+        checkKeywordShard([3, 2, 1], shardCount: 1001, expectedShard: 842)
+    }
 }