Fix trait Source

Author: DanglingPointer

src/sync/channel.rs

@@ -1,6 +1,7 @@
 use super::shared_state::{SharedState, Source};
 use crate::sealed;
 use std::cell::Cell;
+use std::ops::ControlFlow;
 use std::pin::Pin;
 use std::rc::Rc;
 use std::task::{Context, Poll};
@@ -14,12 +15,14 @@ struct Data<T> {
 impl<T> Source for Data<T> {
     type Item = T;
 
-    fn closed(&self) -> bool {
-        self.sender_count.get() == 0
-    }
-
-    fn extract_item(&self) -> Option<Self::Item> {
-        self.queue.pop()
+    fn try_yield_one(&self) -> ControlFlow<Option<Self::Item>> {
+        if let Some(item) = self.queue.pop() {
+            ControlFlow::Break(Some(item))
+        } else if self.sender_count.get() == 0 {
+            ControlFlow::Break(None)
+        } else {
+            ControlFlow::Continue(())
+        }
     }
 }
 

src/sync/condvar.rs

@@ -2,6 +2,7 @@ use super::shared_state::{SharedState, Source};
 use futures::FutureExt;
 use std::cell::Cell;
 use std::future::{poll_fn, Future};
+use std::ops::ControlFlow;
 use std::rc::Rc;
 
 struct Data {
@@ -14,16 +15,13 @@ struct Data {
 impl Source for Data {
     type Item = ();
 
-    fn closed(&self) -> bool {
-        !self.has_sender.get()
-    }
-
-    fn extract_item(&self) -> Option<Self::Item> {
-        if !self.closed() && self.notified.get() {
-            self.notified.set(false);
-            Some(())
+    fn try_yield_one(&self) -> ControlFlow<Option<Self::Item>> {
+        if !self.has_sender.get() {
+            ControlFlow::Break(None)
+        } else if self.notified.replace(false) {
+            ControlFlow::Break(Some(()))
         } else {
-            None
+            ControlFlow::Continue(())
         }
     }
 }

src/sync/oneshot.rs

@@ -1,6 +1,7 @@
 use super::shared_state::{SharedState, Source};
 use std::cell::Cell;
 use std::future::Future;
+use std::ops::ControlFlow;
 use std::pin::Pin;
 use std::rc::Rc;
 use std::task::{Context, Poll};
@@ -14,12 +15,12 @@ struct Data<T> {
 impl<T> Source for Data<T> {
     type Item = T;
 
-    fn closed(&self) -> bool {
-        !self.has_sender.get()
-    }
-
-    fn extract_item(&self) -> Option<Self::Item> {
-        self.value.replace(None)
+    fn try_yield_one(&self) -> ControlFlow<Option<Self::Item>> {
+        match self.value.take() {
+            Some(value) => ControlFlow::Break(Some(value)),
+            None if !self.has_sender.get() => ControlFlow::Break(None),
+            None => ControlFlow::Continue(()),
+        }
     }
 }
 

src/sync/semaphore.rs

@@ -2,6 +2,7 @@ use super::shared_state::{SharedState, Source};
 use futures::FutureExt;
 use std::cell::Cell;
 use std::future::{poll_fn, Future};
+use std::ops::ControlFlow;
 use std::rc::Rc;
 
 struct Data {
@@ -14,17 +15,14 @@ struct Data {
 impl Source for Data {
     type Item = ();
 
-    fn closed(&self) -> bool {
-        !self.has_sender.get()
-    }
-
-    fn extract_item(&self) -> Option<Self::Item> {
-        let current_capacity = self.capacity.get();
-        if !self.closed() && current_capacity > 0 {
-            self.capacity.set(current_capacity - 1);
-            Some(())
+    fn try_yield_one(&self) -> ControlFlow<Option<Self::Item>> {
+        if !self.has_sender.get() {
+            ControlFlow::Break(None)
+        } else if self.capacity.get() > 0 {
+            self.capacity.update(|cap| cap - 1);
+            ControlFlow::Break(Some(()))
         } else {
-            None
+            ControlFlow::Continue(())
         }
     }
 }

src/sync/shared_state.rs

@@ -1,12 +1,11 @@
 use std::cell::Cell;
-use std::ops::Deref;
+use std::ops::{ControlFlow, Deref};
 use std::rc::Rc;
 use std::task::{Context, Poll, Waker};
 
 pub(super) trait Source {
     type Item;
-    fn closed(&self) -> bool;
-    fn extract_item(&self) -> Option<Self::Item>;
+    fn try_yield_one(&self) -> ControlFlow<Option<Self::Item>>;
 }
 
 pub(super) struct SharedState<T> {
@@ -36,10 +35,8 @@ impl<T: Source> SharedState<T> {
     // This should NEVER be called concurrently from different futures/tasks,
     // because we store only 1 waker
     pub(super) fn poll_wait(self: &mut Rc<Self>, cx: &mut Context<'_>) -> Poll<Option<T::Item>> {
-        if let Some(item) = self.inner.extract_item() {
-            Poll::Ready(Some(item))
-        } else if self.inner.closed() {
-            Poll::Ready(None)
+        if let ControlFlow::Break(output) = self.inner.try_yield_one() {
+            Poll::Ready(output)
         } else {
             let new_waker = match self.waker.replace(None) {
                 Some(waker) if waker.will_wake(cx.waker()) => waker,