Remove public get_mut from peer addresses

libp2p · Dec 14, 2023 · 88833e9 · 88833e9
1 parent 5dd9846
commit 88833e9
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 17 deletions.
diff --git a/protocols/identify/src/behaviour.rs b/protocols/identify/src/behaviour.rs
@@ -199,9 +199,9 @@ impl Behaviour {
             .or_default()
             .insert(conn, addr);
 
-        if let Some(entry) = self.discovered_peers.get_mut(&peer_id) {
+        if let Some(cache) = self.discovered_peers.0.as_mut() {
             for addr in failed_addresses {
-                entry.remove(addr);
+                cache.remove(&peer_id, addr);
             }
         }
     }
@@ -386,11 +386,11 @@ impl NetworkBehaviour for Behaviour {
                 self.our_observed_addresses.remove(&connection_id);
             }
             FromSwarm::DialFailure(DialFailure { peer_id, error, .. }) => {
-                if let Some(entry) = peer_id.and_then(|id| self.discovered_peers.get_mut(&id)) {
-                    if let DialError::Transport(errors) = error {
-                        for (addr, _error) in errors {
-                            entry.remove(addr);
-                        }
+                if let (Some(peer_id), Some(cache), DialError::Transport(errors)) =
+                    (peer_id, self.discovered_peers.0.as_mut(), error)
+                {
+                    for (addr, _error) in errors {
+                        cache.remove(&peer_id, addr);
                     }
                 }
             }
@@ -455,10 +455,6 @@ impl PeerCache {
         Self(Some(PeerAddresses::new(size)))
     }
 
-    fn get_mut(&mut self, peer: &PeerId) -> Option<&mut HashSet<Multiaddr>> {
-        self.0.as_mut()?.get_mut(peer)
-    }
-
     fn put(&mut self, peer: PeerId, addresses: impl Iterator<Item = Multiaddr>) {
         if let Some(cache) = self.0.as_mut() {
             cache.put(peer, addresses);

diff --git a/swarm/src/behaviour/peer_addresses.rs b/swarm/src/behaviour/peer_addresses.rs
@@ -20,7 +20,7 @@ impl PeerAddresses {
     pub fn on_swarm_event(&mut self, event: &FromSwarm) -> bool {
         match event {
             FromSwarm::NewExternalAddrOfPeer(NewExternalAddrOfPeer { peer_id, addr }) => {
-                if let Some(peer_addrs) = self.get_mut(peer_id) {
+                if let Some(peer_addrs) = self.0.get_mut(peer_id) {
                     let addr = prepare_addr(peer_id, addr);
                     peer_addrs.insert(addr)
                 } else {
@@ -42,10 +42,6 @@ impl PeerAddresses {
         Self(LruCache::new(cache_size))
     }
 
-    pub fn get_mut(&mut self, peer: &PeerId) -> Option<&mut HashSet<Multiaddr>> {
-        self.0.get_mut(peer)
-    }
-
     pub fn put(&mut self, peer: PeerId, addresses: impl Iterator<Item = Multiaddr>) -> bool {
         let addresses = addresses.filter_map(|a| a.with_p2p(peer).ok());
         self.0.put(peer, HashSet::from_iter(addresses));
@@ -65,7 +61,7 @@ impl PeerAddresses {
     /// Removes address from peer addresses cache.
     /// Returns true if the address was removed.
     pub fn remove(&mut self, peer: &PeerId, address: &Multiaddr) -> bool {
-        self.get_mut(peer).map_or_else(
+        self.0.get_mut(peer).map_or_else(
             || false,
             |addrs| {
                 let address = prepare_addr(peer, address);