Add wrapper types for short IDs

src/service/rooms/state_compressor.rs

@@ -1,4 +1,5 @@
 use std::{
+    array,
     collections::HashSet,
     mem::size_of,
     sync::{Arc, Mutex},
@@ -17,9 +18,11 @@ pub(crate) mod data;
 pub(crate) use data::Data;
 use data::StateDiff;
 
+use super::short::{ShortEventId, ShortStateHash, ShortStateKey};
+
 #[derive(Clone)]
 pub(crate) struct CompressedStateLayer {
-    pub(crate) shortstatehash: u64,
+    pub(crate) shortstatehash: ShortStateHash,
     pub(crate) full_state: Arc<HashSet<CompressedStateEvent>>,
     pub(crate) added: Arc<HashSet<CompressedStateEvent>>,
     pub(crate) removed: Arc<HashSet<CompressedStateEvent>>,
@@ -29,10 +32,45 @@ pub(crate) struct Service {
     pub(crate) db: &'static dyn Data,
 
     #[allow(clippy::type_complexity)]
-    pub(crate) stateinfo_cache: Mutex<LruCache<u64, Vec<CompressedStateLayer>>>,
+    pub(crate) stateinfo_cache:
+        Mutex<LruCache<ShortStateHash, Vec<CompressedStateLayer>>>,
 }
 
-pub(crate) type CompressedStateEvent = [u8; 2 * size_of::<u64>()];
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub(crate) struct CompressedStateEvent {
+    pub(crate) state: ShortStateKey,
+    pub(crate) event: ShortEventId,
+}
+
+impl CompressedStateEvent {
+    pub(crate) fn as_bytes(
+        &self,
+    ) -> [u8; size_of::<ShortStateKey>() + size_of::<ShortEventId>()] {
+        let mut bytes = self
+            .state
+            .get()
+            .to_be_bytes()
+            .into_iter()
+            .chain(self.event.get().to_be_bytes());
+        array::from_fn(|_| bytes.next().unwrap())
+    }
+
+    pub(crate) fn from_bytes(
+        bytes: [u8; size_of::<ShortStateKey>() + size_of::<ShortEventId>()],
+    ) -> Self {
+        let state = ShortStateKey::new(u64::from_be_bytes(
+            bytes[0..8].try_into().unwrap(),
+        ));
+        let event = ShortEventId::new(u64::from_be_bytes(
+            bytes[8..16].try_into().unwrap(),
+        ));
+
+        Self {
+            state,
+            event,
+        }
+    }
+}
 
 impl Service {
     /// Returns a stack with info on shortstatehash, full state, added diff and
@@ -41,7 +79,7 @@ impl Service {
     #[tracing::instrument(skip(self))]
     pub(crate) fn load_shortstatehash_info(
         &self,
-        shortstatehash: u64,
+        shortstatehash: ShortStateHash,
     ) -> Result<Vec<CompressedStateLayer>> {
         let lookup = Lookup::StateInfo;
 
@@ -96,18 +134,16 @@ impl Service {
     #[allow(clippy::unused_self)]
     pub(crate) fn compress_state_event(
         &self,
-        shortstatekey: u64,
+        shortstatekey: ShortStateKey,
         event_id: &EventId,
     ) -> Result<CompressedStateEvent> {
-        let mut v = shortstatekey.to_be_bytes().to_vec();
-        v.extend_from_slice(
-            &services()
+        Ok(CompressedStateEvent {
+            state: shortstatekey,
+            event: services()
                 .rooms
                 .short
-                .get_or_create_shorteventid(event_id)?
-                .to_be_bytes(),
-        );
-        Ok(v.try_into().expect("we checked the size above"))
+                .get_or_create_shorteventid(event_id)?,
+        })
     }
 
     /// Returns shortstatekey, event id
@@ -116,14 +152,13 @@ impl Service {
     pub(crate) fn parse_compressed_state_event(
         &self,
         compressed_event: &CompressedStateEvent,
-    ) -> Result<(u64, Arc<EventId>)> {
+    ) -> Result<(ShortStateKey, Arc<EventId>)> {
         Ok((
-            utils::u64_from_bytes(&compressed_event[0..size_of::<u64>()])
-                .expect("bytes have right length"),
-            services().rooms.short.get_eventid_from_short(
-                utils::u64_from_bytes(&compressed_event[size_of::<u64>()..])
-                    .expect("bytes have right length"),
-            )?,
+            compressed_event.state,
+            services()
+                .rooms
+                .short
+                .get_eventid_from_short(compressed_event.event)?,
         ))
     }
 
@@ -155,7 +190,7 @@ impl Service {
     ))]
     pub(crate) fn save_state_from_diff(
         &self,
-        shortstatehash: u64,
+        shortstatehash: ShortStateHash,
         statediffnew: Arc<HashSet<CompressedStateEvent>>,
         statediffremoved: Arc<HashSet<CompressedStateEvent>>,
         diff_to_sibling: usize,
@@ -275,7 +310,7 @@ impl Service {
         room_id: &RoomId,
         new_state_ids_compressed: Arc<HashSet<CompressedStateEvent>>,
     ) -> Result<(
-        u64,
+        ShortStateHash,
         Arc<HashSet<CompressedStateEvent>>,
         Arc<HashSet<CompressedStateEvent>>,
     )> {
@@ -283,7 +318,7 @@ impl Service {
             services().rooms.state.get_room_shortstatehash(room_id)?;
 
         let state_hash = utils::calculate_hash(
-            new_state_ids_compressed.iter().map(|bytes| &bytes[..]),
+            new_state_ids_compressed.iter().map(CompressedStateEvent::as_bytes),
         );
 
         let (new_shortstatehash, already_existed) =