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::Session;
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<dyn Session>,
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(&self, peer: PeerId) -> Option<PeerRef> {
39        match self.map.get(&peer) {
40            Some(peer) => peer.to_ref().await,
41            None => self.session.peer_ref(peer).await,
42        }
43    }
44
45    /// Take the full `Peer` object given its `PeerId`.
46    ///
47    /// The peer will be removed from the map if there are no other strong references to it.
48    pub fn take(&mut self, peer: PeerId) -> Option<Peer> {
49        match Arc::get_mut(&mut self.map) {
50            Some(map) => map.remove(&peer),
51            None => self.get(peer).cloned(),
52        }
53    }
54
55    pub(crate) fn take_user(&mut self, user_id: i64) -> Option<User> {
56        self.take(PeerId::user_unchecked(user_id))
57            .map(|peer| match peer {
58                Peer::User(user) => user,
59                _ => unreachable!(),
60            })
61    }
62
63    /// Iterate over the peers and peers in the map.
64    pub fn iter(&self) -> impl Iterator<Item = (PeerId, &Peer)> {
65        self.map.iter().map(|(k, v)| (*k, v))
66    }
67
68    /// Iterate over the peers in the map.
69    pub fn iter_peers(&self) -> impl Iterator<Item = &Peer> {
70        self.map.values()
71    }
72
73    /// Return a new strong reference to the map and session contained within.
74    pub fn handle(&self) -> Self {
75        Self {
76            map: Arc::clone(&self.map),
77            session: Arc::clone(&self.session),
78        }
79    }
80}
grammers_client/peer/peer_map.rs

grammers_client/peer/
peer_map.rs
