Adds caching logic

Author: assemblaj

src/stream/stream/partial_cmp.rs

@@ -1,30 +1,36 @@
 use std::cmp::Ordering;
 use std::pin::Pin;
 
+use super::fuse::Fuse;
+use crate::prelude::*;
 use crate::future::Future;
 use crate::stream::Stream;
 use crate::task::{Context, Poll};
 
 // Lexicographically compares the elements of this `Stream` with those
 // of another.
 #[derive(Debug)]
-pub struct PartialCmpFuture<L: Stream, R> {
-    l: L,
-    r: R,
+pub struct PartialCmpFuture<L: Stream, R: Stream> {
+    l: Fuse<L>,
+    r: Fuse<R>,
+    l_cache: Option<L::Item>,
+    r_cache: Option<R::Item>, 
 }
 
-impl<L: Stream + Unpin, R: Unpin> Unpin for PartialCmpFuture<L, R> {}
+impl<L: Stream + Unpin, R: Stream + Unpin> Unpin for PartialCmpFuture<L, R> {}
 
-impl<L: Stream, R> PartialCmpFuture<L, R> {
-    pub(crate) fn new(l: L, r: R) -> Self {
+impl<L: Stream, R: Stream> PartialCmpFuture<L, R> {
+    pin_utils::unsafe_pinned!(l: Fuse<L>);
+    pin_utils::unsafe_pinned!(r: Fuse<R>);
+
+    pub(super) fn new(l: L, r: R) -> Self {
         PartialCmpFuture {
-            l,
-            r,
+            l: l.fuse(),
+            r: r.fuse(),
+            l_cache: None, 
+            r_cache: None, 
         }
     }
-
-    pin_utils::unsafe_pinned!(l: L);
-    pin_utils::unsafe_pinned!(r: R);
 }
 
 impl<L: Stream, R: Stream> Future for PartialCmpFuture<L, R> 
@@ -35,22 +41,53 @@ where
 {
     type Output = Option<Ordering>;
 
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        let l_next  = futures_core::ready!(self.as_mut().l().poll_next(cx));
-        let r_next  = futures_core::ready!(self.as_mut().r().poll_next(cx));
-        
-        match (l_next, r_next) {
-            (None, None) => Poll::Ready(Some(Ordering::Equal)),
-            (None, _   ) => Poll::Ready(Some(Ordering::Less)),
-            (_,    None) => Poll::Ready(Some(Ordering::Greater)),
-            (Some(x), Some(y)) =>  match x.partial_cmp(&y) {
-                    Some(Ordering::Equal) => { 
-                        // wakes task to pull the next item from stream
-                        cx.waker().wake_by_ref();
-                        Poll::Pending
-                    },
-                    non_eq      => Poll::Ready(non_eq),
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {        
+        // Short circuit logic 
+        // Stream that completes earliest can be considered Less, etc
+        let l_complete = self.l.done && self.as_mut().l_cache.is_none();
+        let r_complete = self.r.done && self.as_mut().r_cache.is_none();
+
+        if l_complete && r_complete {
+            return Poll::Ready(Some(Ordering::Equal))
+        } else if l_complete {
+            return Poll::Ready(Some(Ordering::Less))
+        } else if r_complete {
+            return Poll::Ready(Some(Ordering::Greater))
+        }
+
+        // Get next value if possible and necesary 
+        if !self.l.done && self.as_mut().l_cache.is_none() {
+            let l_next = futures_core::ready!(self.as_mut().l().poll_next(cx));
+            if let Some(item) = l_next {
+                self.as_mut().l_cache = Some(item);
+            }
+        }
+
+        if !self.r.done && self.as_mut().r_cache.is_none() {
+            let r_next = futures_core::ready!(self.as_mut().r().poll_next(cx));
+            if let Some(item) = r_next {
+                self.as_mut().r_cache = Some(item);
             }
         }
+
+        // Compare if both values are available.
+        if self.as_mut().l_cache.is_some() && self.as_mut().r_cache.is_some() {
+            let l_value = self.as_mut().l_cache.take().unwrap();
+            let r_value = self.as_mut().r_cache.take().unwrap(); 
+            let result = l_value.partial_cmp(&r_value);
+
+            if let Some(Ordering::Equal) = result {
+                // Reset cache to prepare for next comparison 
+                self.as_mut().l_cache = None;  
+                self.as_mut().r_cache = None;
+            } else {
+                // Return non equal value 
+                return Poll::Ready(result);
+            }            
+        }
+
+        // wakes task to pull the next item from stream
+        cx.waker().wake_by_ref();
+        Poll::Pending
     }
 }