Add ptr_eq and ptr_hash to Shared (#2673)

Author: khuey

futures-util/src/future/future/shared.rs

@@ -4,7 +4,9 @@ use futures_core::task::{Context, Poll, Waker};
 use slab::Slab;
 use std::cell::UnsafeCell;
 use std::fmt;
+use std::hash::Hasher;
 use std::pin::Pin;
+use std::ptr;
 use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering::{Acquire, SeqCst};
 use std::sync::{Arc, Mutex, Weak};
@@ -156,6 +158,35 @@ where
     pub fn weak_count(&self) -> Option<usize> {
         self.inner.as_ref().map(Arc::weak_count)
     }
+
+    /// Hashes the internal state of this `Shared` in a way that's compatible with `ptr_eq`.
+    pub fn ptr_hash<H: Hasher>(&self, state: &mut H) {
+        match self.inner.as_ref() {
+            Some(arc) => {
+                state.write_u8(1);
+                ptr::hash(Arc::as_ptr(arc), state);
+            }
+            None => {
+                state.write_u8(0);
+            }
+        }
+    }
+
+    /// Returns `true` if the two `Shared`s point to the same future (in a vein similar to
+    /// `Arc::ptr_eq`).
+    ///
+    /// Returns `false` if either `Shared` has terminated.
+    pub fn ptr_eq(&self, rhs: &Self) -> bool {
+        let lhs = match self.inner.as_ref() {
+            Some(lhs) => lhs,
+            None => return false,
+        };
+        let rhs = match rhs.inner.as_ref() {
+            Some(rhs) => rhs,
+            None => return false,
+        };
+        Arc::ptr_eq(lhs, rhs)
+    }
 }
 
 impl<Fut> Inner<Fut>

futures/tests/future_shared.rs

@@ -151,6 +151,52 @@ fn downgrade() {
     assert!(shared2.downgrade().is_none());
 }
 
+#[test]
+fn ptr_eq() {
+    use future::FusedFuture;
+    use std::collections::hash_map::DefaultHasher;
+    use std::hash::Hasher;
+
+    let (tx, rx) = oneshot::channel::<i32>();
+    let shared = rx.shared();
+    let mut shared2 = shared.clone();
+    let mut hasher = DefaultHasher::new();
+    let mut hasher2 = DefaultHasher::new();
+
+    // Because these two futures share the same underlying future,
+    // `ptr_eq` should return true.
+    assert!(shared.ptr_eq(&shared2));
+    // Equivalence relations are symmetric
+    assert!(shared2.ptr_eq(&shared));
+
+    // If `ptr_eq` returns true, they should hash to the same value.
+    shared.ptr_hash(&mut hasher);
+    shared2.ptr_hash(&mut hasher2);
+    assert_eq!(hasher.finish(), hasher2.finish());
+
+    tx.send(42).unwrap();
+    assert_eq!(block_on(&mut shared2).unwrap(), 42);
+
+    // Now that `shared2` has completed, `ptr_eq` should return false.
+    assert!(shared2.is_terminated());
+    assert!(!shared.ptr_eq(&shared2));
+
+    // `ptr_eq` should continue to work for the other `Shared`.
+    let shared3 = shared.clone();
+    let mut hasher3 = DefaultHasher::new();
+    assert!(shared.ptr_eq(&shared3));
+
+    shared3.ptr_hash(&mut hasher3);
+    assert_eq!(hasher.finish(), hasher3.finish());
+
+    let (_tx, rx) = oneshot::channel::<i32>();
+    let shared4 = rx.shared();
+
+    // And `ptr_eq` should return false for two futures that don't share
+    // the underlying future.
+    assert!(!shared.ptr_eq(&shared4));
+}
+
 #[test]
 fn dont_clone_in_single_owner_shared_future() {
     let counter = CountClone(Rc::new(Cell::new(0)));