do modulo on u128, update compile check for TAG_SIZE

ipa-core/Cargo.toml

@@ -82,6 +82,7 @@ ipa-step = { version = "*", path = "../ipa-step" }
 ipa-step-derive = { version = "*", path = "../ipa-step-derive" }
 
 aes = "0.8.3"
+assertions = "0.1.0"
 async-trait = "0.1.79"
 async-scoped = { version = "0.9.0", features = ["use-tokio"], optional = true }
 axum = { version = "0.7.5", optional = true, features = ["http2", "macros"] }

ipa-core/src/report/hybrid.rs

@@ -1,5 +1,6 @@
 use std::{collections::HashSet, convert::Infallible, ops::Add};
 
+use assertions::const_assert;
 use bytes::Bytes;
 use generic_array::{ArrayLength, GenericArray};
 use rand_core::{CryptoRng, RngCore};
@@ -180,6 +181,12 @@ impl UniqueBytes for EncryptedHybridReport {
 }
 
 impl UniqueTag {
+    fn _compile_check() {
+        // This will vaild at compile time if TAG_SIZE doesn't match U16
+        // the macro expansion needs to be wrapped in a function
+        const_assert!(TAG_SIZE == 16);
+    }
+
     // Function to attempt to create a UniqueTag from a UniqueBytes implementor
     pub fn from_unique_bytes<T: UniqueBytes>(item: &T) -> Self {
         UniqueTag {
@@ -188,26 +195,18 @@ impl UniqueTag {
     }
 
     /// Maps the tag into a consistent shard.
-    /// Note that `ShardIndex` is limited to u32, so we only use the first 4 bytes.
     ///
     /// ## Panics
-    /// if the `TAG_SIZE < 4`
+    /// if the `TAG_SIZE != 16`
     /// note: ~10 below this, we have a compile time check that `TAG_SIZE = 16`
     #[must_use]
     pub fn shard_picker(&self, shard_count: ShardIndex) -> ShardIndex {
-        let num = u32::from_le_bytes(
-            self.bytes[0..4]
-                .try_into()
-                .expect("This is larger than 4 bytes"),
-        );
-        let shard_count = u32::from(shard_count);
-        ShardIndex::from(num % shard_count)
+        let num = u128::from_le_bytes(self.bytes);
+        let shard_count = u128::from(shard_count);
+        ShardIndex::try_from(num % shard_count).expect("Modulo a u32 will fit in u32")
     }
 }
 
-// This will vaild at compile time if TAG_SIZE doesn't match U16
-const _: [(); 16] = [(); TAG_SIZE];
-
 impl Serializable for UniqueTag {
     type Size = U16; // This must match TAG_SIZE
     type DeserializationError = Infallible;

ipa-core/src/sharding.rs

@@ -112,6 +112,14 @@ impl TryFrom<usize> for ShardIndex {
     }
 }
 
+impl TryFrom<u128> for ShardIndex {
+    type Error = TryFromIntError;
+
+    fn try_from(value: u128) -> Result<Self, Self::Error> {
+        u32::try_from(value).map(Self)
+    }
+}
+
 #[cfg(all(test, unit_test))]
 mod tests {
     use std::iter::empty;