Add PrincipalToUserIdMap backed by stable memory (#7017)

Cargo.toml

@@ -136,6 +136,7 @@ members = [
     "backend/libraries/ledger_utils",
     "backend/libraries/msgpack",
     "backend/libraries/p256_key_pair",
+    "backend/libraries/principal_to_user_id_map",
     "backend/libraries/proof_of_unique_personhood",
     "backend/libraries/storage_bucket_client",
     "backend/libraries/search",

backend/canisters/community/impl/src/model/events/stable_memory.rs

@@ -4,19 +4,11 @@ use serde::Serialize;
 use stable_memory_map::{with_map_mut, CommunityEventKeyPrefix, KeyPrefix};
 use types::EventWrapperInternal;
 
-#[derive(Serialize, Deserialize)]
+#[derive(Serialize, Deserialize, Default)]
 pub struct EventsStableStorage {
     prefix: CommunityEventKeyPrefix,
 }
 
-impl Default for EventsStableStorage {
-    fn default() -> Self {
-        EventsStableStorage {
-            prefix: CommunityEventKeyPrefix::new(),
-        }
-    }
-}
-
 impl EventsStableStorage {
     pub fn insert(&mut self, event: EventWrapperInternal<CommunityEventInternal>) {
         with_map_mut(|m| m.insert(self.prefix.create_key(&event.index), event_to_bytes(event)));

backend/libraries/principal_to_user_id_map/Cargo.toml

@@ -0,0 +1,10 @@
+[package]
+name = "principal_to_user_id_map"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+ic_principal = { workspace = true }
+stable_memory_map = { path = "../stable_memory_map" }
+types = { path = "../types" }
+serde = { workspace = true }

backend/libraries/principal_to_user_id_map/src/lib.rs

@@ -0,0 +1,40 @@
+use ic_principal::Principal;
+use serde::{Deserialize, Serialize};
+use stable_memory_map::{with_map, with_map_mut, KeyPrefix, PrincipalToUserIdKeyPrefix};
+use types::UserId;
+
+#[derive(Serialize, Deserialize, Default)]
+pub struct PrincipalToUserIdMap {
+    prefix: PrincipalToUserIdKeyPrefix,
+    count: u32,
+}
+
+impl PrincipalToUserIdMap {
+    pub fn get(&self, principal: &Principal) -> Option<UserId> {
+        with_map(|m| m.get(self.prefix.create_key(principal)).map(bytes_to_user_id))
+    }
+
+    pub fn insert(&mut self, principal: Principal, user_id: UserId) {
+        if with_map_mut(|m| m.insert(self.prefix.create_key(&principal), user_id.as_slice().to_vec())).is_none() {
+            self.count += 1;
+        }
+    }
+
+    pub fn remove(&mut self, principal: &Principal) -> Option<UserId> {
+        let bytes = with_map_mut(|m| m.remove(self.prefix.create_key(principal)))?;
+        self.count = self.count.saturating_sub(1);
+        Some(bytes_to_user_id(bytes))
+    }
+
+    pub fn len(&self) -> u32 {
+        self.count
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.count == 0
+    }
+}
+
+fn bytes_to_user_id(bytes: Vec<u8>) -> UserId {
+    UserId::from(Principal::from_slice(&bytes))
+}

backend/libraries/stable_memory_map/src/keys.rs

@@ -7,17 +7,19 @@ mod chat_event;
 mod community_event;
 mod macros;
 mod member;
+mod principal_to_user_id;
 
 pub use chat_event::*;
 pub use community_event::*;
 pub use member::*;
+pub use principal_to_user_id::*;
 
 #[derive(Serialize, Deserialize, Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
 #[serde(transparent)]
 pub struct BaseKey(#[serde(with = "serde_bytes")] Vec<u8>);
 
 impl BaseKey {
-    pub fn starts_with(&self, prefix: &BaseKeyPrefix) -> bool {
+    pub fn matches_prefix(&self, prefix: &BaseKeyPrefix) -> bool {
         self.0.starts_with(prefix.0.as_slice())
     }
 
@@ -88,6 +90,7 @@ pub enum KeyType {
     ChannelMember = 8,
     CommunityMember = 9,
     CommunityEvent = 10,
+    PrincipalToUserId = 11,
 }
 
 fn extract_key_type(bytes: &[u8]) -> Option<KeyType> {
@@ -109,6 +112,7 @@ impl TryFrom<u8> for KeyType {
             8 => Ok(KeyType::ChannelMember),
             9 => Ok(KeyType::CommunityMember),
             10 => Ok(KeyType::CommunityEvent),
+            11 => Ok(KeyType::PrincipalToUserId),
             _ => Err(()),
         }
     }

backend/libraries/stable_memory_map/src/keys/principal_to_user_id.rs

@@ -0,0 +1,65 @@
+use crate::keys::macros::key;
+use crate::{BaseKey, KeyPrefix, KeyType};
+use ic_principal::Principal;
+
+key!(PrincipalToUserIdKey, PrincipalToUserIdKeyPrefix, KeyType::PrincipalToUserId);
+
+impl PrincipalToUserIdKeyPrefix {
+    pub fn new() -> Self {
+        // KeyType::PrincipalToUserId   1 byte
+        PrincipalToUserIdKeyPrefix(vec![KeyType::PrincipalToUserId as u8])
+    }
+}
+
+impl KeyPrefix for PrincipalToUserIdKeyPrefix {
+    type Key = PrincipalToUserIdKey;
+    type Suffix = Principal;
+
+    fn create_key(&self, principal: &Principal) -> PrincipalToUserIdKey {
+        let principal_bytes = principal.as_slice();
+        let mut bytes = Vec::with_capacity(principal_bytes.len() + 1);
+        bytes.push(KeyType::PrincipalToUserId as u8);
+        bytes.extend_from_slice(principal_bytes);
+        PrincipalToUserIdKey(bytes)
+    }
+}
+
+impl Default for PrincipalToUserIdKeyPrefix {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl PrincipalToUserIdKey {
+    pub fn principal(&self) -> Principal {
+        Principal::from_slice(&self.0[1..])
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::Key;
+    use rand::{thread_rng, RngCore};
+
+    #[test]
+    fn principal_to_user_id_key_e2e() {
+        for _ in 0..100 {
+            let prefix = PrincipalToUserIdKeyPrefix::new();
+            let principal = Principal::from_slice(&thread_rng().next_u32().to_be_bytes());
+            let key = BaseKey::from(prefix.create_key(&principal));
+            let principal_to_user_id_key = PrincipalToUserIdKey::try_from(key.clone()).unwrap();
+
+            assert_eq!(*principal_to_user_id_key.0.first().unwrap(), KeyType::PrincipalToUserId as u8);
+            assert_eq!(principal_to_user_id_key.0.len(), principal.as_slice().len() + 1);
+            assert!(principal_to_user_id_key.matches_prefix(&prefix));
+            assert_eq!(principal_to_user_id_key.principal(), principal);
+
+            let serialized = msgpack::serialize_then_unwrap(&principal_to_user_id_key);
+            assert_eq!(serialized.len(), principal_to_user_id_key.0.len() + 2);
+            let deserialized: PrincipalToUserIdKey = msgpack::deserialize_then_unwrap(&serialized);
+            assert_eq!(deserialized, principal_to_user_id_key);
+            assert_eq!(deserialized.0, key.0);
+        }
+    }
+}

backend/libraries/stable_memory_map/src/lib.rs

@@ -41,8 +41,8 @@ impl StableMemoryMap {
         self.map.get(&key.into())
     }
 
-    pub fn insert<K: Key>(&mut self, key: K, value: Vec<u8>) {
-        self.map.insert(key.into(), value);
+    pub fn insert<K: Key>(&mut self, key: K, value: Vec<u8>) -> Option<Vec<u8>> {
+        self.map.insert(key.into(), value)
     }
 
     pub fn remove<K: Key>(&mut self, key: K) -> Option<Vec<u8>> {
@@ -68,7 +68,7 @@ pub fn garbage_collect(prefix: BaseKeyPrefix) -> Result<u32, u32> {
             let keys: Vec<_> = m
                 .map
                 .range(BaseKey::from(prefix.clone())..)
-                .take_while(|(k, _)| k.starts_with(&prefix))
+                .take_while(|(k, _)| k.matches_prefix(&prefix))
                 .map(|(k, _)| k)
                 .take(100)
                 .collect();