grammers_client/peer/
peer_map.rs

1// Copyright 2020 - developers of the `grammers` project.
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
6// option. This file may not be copied, modified, or distributed
7// except according to those terms.
8
9use std::collections::HashMap;
10use std::sync::Arc;
11
12use grammers_session::ErasedSession;
13use grammers_session::types::{PeerId, PeerRef};
14
15use crate::peer::{Peer, User};
16
17/// Helper structure to efficiently retrieve peers via their peer.
18///
19/// A lot of responses include the peers related to them in the form of a list of users
20/// and peers, making it annoying to extract a specific peer. This structure lets you
21/// save those separate vectors in a single place and query them by using a `Peer`.
22///
23/// While this type derives `Clone` for convenience, it is recommended to use
24/// [`PeerMap::handle`] instead to signal that it is a cheap clone.
25#[derive(Clone)]
26pub struct PeerMap {
27    pub(crate) map: Arc<HashMap<PeerId, Peer>>,
28    pub(crate) session: Arc<ErasedSession>,
29}
30
31impl PeerMap {
32    /// Retrieve the full `Peer` object given its `PeerId`.
33    pub fn get(&self, peer: PeerId) -> Option<&Peer> {
34        self.map.get(&peer)
35    }
36
37    /// Retrieve a non-min `PeerRef` from either the in-memory cache or the session.
38    pub async fn get_ref(
39        &self,
40        peer: PeerId,
41    ) -> Result<Option<PeerRef>, Box<dyn std::error::Error + Send + Sync>> {
42        match self.map.get(&peer) {
43            Some(peer) => peer.to_ref().await,
44            None => self.session.peer_ref(peer).await,
45        }
46    }
47
48    /// Take the full `Peer` object given its `PeerId`.
49    ///
50    /// The peer will be removed from the map if there are no other strong references to it.
51    pub fn take(&mut self, peer: PeerId) -> Option<Peer> {
52        match Arc::get_mut(&mut self.map) {
53            Some(map) => map.remove(&peer),
54            None => self.get(peer).cloned(),
55        }
56    }
57
58    pub(crate) fn take_user(&mut self, user_id: i64) -> Option<User> {
59        self.take(PeerId::user_unchecked(user_id))
60            .map(|peer| match peer {
61                Peer::User(user) => user,
62                _ => unreachable!(),
63            })
64    }
65
66    /// Iterate over the peers and peers in the map.
67    pub fn iter(&self) -> impl Iterator<Item = (PeerId, &Peer)> {
68        self.map.iter().map(|(k, v)| (*k, v))
69    }
70
71    /// Iterate over the peers in the map.
72    pub fn iter_peers(&self) -> impl Iterator<Item = &Peer> {
73        self.map.values()
74    }
75
76    /// Return a new strong reference to the map and session contained within.
77    pub fn handle(&self) -> Self {
78        Self {
79            map: Arc::clone(&self.map),
80            session: Arc::clone(&self.session),
81        }
82    }
83}
grammers_client/peer/peer_map.rs

grammers_client/peer/
peer_map.rs
