Compile streams

4 files changed, 98 insertions, 118 deletions

diff --git a/src/imap/mailbox_view.rs b/src/imap/mailbox_view.rs
index d57e9a3..5055c40 100644
--- a/src/imap/mailbox_view.rs
+++ b/src/imap/mailbox_view.rs
@@ -4,7 +4,7 @@ use std::sync::Arc;
 
 use anyhow::{anyhow, Error, Result};
 
-use futures::stream::{FuturesOrdered, StreamExt};
+use futures::stream::{StreamExt, TryStreamExt};
 
 use imap_codec::imap_types::core::Charset;
 use imap_codec::imap_types::fetch::MessageDataItem;
@@ -362,46 +362,33 @@ impl MailboxView {
             .iter()
             .map(|midx| midx.uuid)
             .collect::<Vec<_>>();
-        let query_result = self.internal.query(&uuids, query_scope).fetch().await?;
 
-        // [3/6] Derive an IMAP-specific view from the results, apply the filters
-        let views = query_result
-            .iter()
-            .zip(mail_idx_list.into_iter())
-            .map(|(qr, midx)| MailView::new(qr, midx))
-            .collect::<Result<Vec<_>, _>>()?;
-
-        // [4/6] Apply the IMAP transformation, bubble up any error
-        // We get 2 results:
-        //   - The one we send to the client
-        //   - The \Seen flags we must set internally
-        let (flag_mgmt, imap_ret): (Vec<_>, Vec<_>) = views
-            .iter()
-            .map(|mv| mv.filter(&ap).map(|(body, seen)| ((mv, seen), body)))
-            .collect::<Result<Vec<_>, _>>()?
-            .into_iter()
-            .unzip();
+        let query = self.internal.query(&uuids, query_scope);
+        //let query_result = self.internal.query(&uuids, query_scope).fetch().await?;
 
-        // [5/6] Register the \Seen flags
-        flag_mgmt
-            .iter()
-            .filter(|(_mv, seen)| matches!(seen, SeenFlag::MustAdd))
-            .map(|(mv, _seen)| async move {
-                let seen_flag = Flag::Seen.to_string();
-                self.internal
-                    .mailbox
-                    .add_flags(*mv.query_result.uuid(), &[seen_flag])
-                    .await?;
-                Ok::<_, anyhow::Error>(())
+        let query_stream = query
+            .fetch()
+            .zip(futures::stream::iter(mail_idx_list))
+            // [3/6] Derive an IMAP-specific view from the results, apply the filters
+            .map(|(maybe_qr, midx)|  match maybe_qr {
+                Ok(qr) => Ok((MailView::new(&qr, midx)?.filter(&ap)?, midx)),
+                Err(e) => Err(e),
             })
-            .collect::<FuturesOrdered<_>>()
-            .collect::<Vec<_>>()
-            .await
-            .into_iter()
-            .collect::<Result<_, _>>()?;
+            // [4/6] Apply the IMAP transformation
+            .then(|maybe_ret| async move {
+                let ((body, seen), midx) = maybe_ret?;
+
+                // [5/6] Register the \Seen flags
+                if matches!(seen, SeenFlag::MustAdd) {
+                    let seen_flag = Flag::Seen.to_string();
+                    self.internal.mailbox.add_flags(midx.uuid, &[seen_flag]).await?;
+                }
+
+                Ok::<_, anyhow::Error>(body)
+            });
 
         // [6/6] Build the final result that will be sent to the client.
-        Ok(imap_ret)
+        query_stream.try_collect().await
     }
 
     /// A naive search implementation...
@@ -423,39 +410,55 @@ impl MailboxView {
         // 3. Filter the selection based on the ID / UID / Flags
         let (kept_idx, to_fetch) = crit.filter_on_idx(&selection);
 
-        // 4. Fetch additional info about the emails
+        // 4.a Fetch additional info about the emails
         let query_scope = crit.query_scope();
         let uuids = to_fetch.iter().map(|midx| midx.uuid).collect::<Vec<_>>();
-        let query_result = self.internal.query(&uuids, query_scope).fetch().await?;
+        let query = self.internal.query(&uuids, query_scope);
+
+        // 4.b We don't want to keep all data in memory, so we do the computing in a stream
+        let query_stream = query
+            .fetch()
+            .zip(futures::stream::iter(&to_fetch))
+            // 5.a Build a mailview with the body, might fail with an error
+            // 5.b If needed, filter the selection based on the body, but keep the errors
+            // 6. Drop the query+mailbox, keep only the mail index
+            // Here we release a lot of memory, this is the most important part ^^
+            .filter_map(|(maybe_qr, midx)| {
+                let r = match maybe_qr {
+                    Ok(qr) => match MailView::new(&qr, midx).map(|mv| crit.is_keep_on_query(&mv)) {
+                        Ok(true) => Some(Ok(*midx)),
+                        Ok(_) => None,
+                        Err(e) => Some(Err(e)),
+                    }
+                    Err(e) => Some(Err(e)),
+                };
+                futures::future::ready(r)
+            });
 
-        // 5. If needed, filter the selection based on the body
-        let kept_query = crit.filter_on_query(&to_fetch, &query_result)?;
 
-        // 6. Format the result according to the client's taste:
-        // either return UID or ID.
-        let final_selection = kept_idx.iter().chain(kept_query.iter());
-        let selection_fmt = match uid {
-            true => final_selection.map(|in_idx| in_idx.uid).collect(),
-            _ => final_selection.map(|in_idx| in_idx.i).collect(),
-        };
+        // 7. Chain both streams (part resolved from index, part resolved from metadata+body)
+        let main_stream = futures::stream::iter(kept_idx)
+            .map(Ok)
+            .chain(query_stream)
+            .map_ok(|idx| match uid {
+                true => (idx.uid, idx.modseq),
+                _ => (idx.i, idx.modseq),
+            });
 
-        // 7. Add the modseq entry if needed
-        let is_modseq = crit.is_modseq();
-        let maybe_modseq = match is_modseq {
-            true => {
-                let final_selection = kept_idx.iter().chain(kept_query.iter());
-                final_selection
-                    .map(|in_idx| in_idx.modseq)
-                    .max()
-                    .map(|r| NonZeroU64::try_from(r))
-                    .transpose()?
-            }
+        // 8. Do the actual computation
+        let internal_result: Vec<_> = main_stream.try_collect().await?;
+        let (selection, modseqs): (Vec<_>, Vec<_>) = internal_result.into_iter().unzip();
+
+        // 9. Aggregate the maximum modseq value
+        let maybe_modseq = match crit.is_modseq() {
+            true => modseqs.into_iter().max(),
             _ => None,
         };
 
+        // 10. Return the final result
         Ok((
-            vec![Body::Data(Data::Search(selection_fmt, maybe_modseq))],
-            is_modseq,
+            vec![Body::Data(Data::Search(selection, maybe_modseq))],
+            maybe_modseq.is_some(),
         ))
     }
 
@@ -678,7 +681,7 @@ mod tests {
             content: rfc822.to_vec(),
         };
 
-        let mv = MailView::new(&qr, &mail_in_idx)?;
+        let mv = MailView::new(std::borrow::Cow::Borrowed(&qr), &mail_in_idx)?;
         let (res_body, _seen) = mv.filter(&ap)?;
 
         let fattr = match res_body {
diff --git a/src/imap/search.rs b/src/imap/search.rs
index d06c3bd..4ff70ee 100644
--- a/src/imap/search.rs
+++ b/src/imap/search.rs
@@ -1,13 +1,12 @@
 use std::num::{NonZeroU32, NonZeroU64};
 
-use anyhow::Result;
 use imap_codec::imap_types::core::NonEmptyVec;
 use imap_codec::imap_types::search::{MetadataItemSearch, SearchKey};
 use imap_codec::imap_types::sequence::{SeqOrUid, Sequence, SequenceSet};
 
 use crate::imap::index::MailIndex;
 use crate::imap::mail_view::MailView;
-use crate::mail::query::{QueryResult, QueryScope};
+use crate::mail::query::QueryScope;
 
 pub enum SeqType {
     Undefined,
@@ -145,22 +144,6 @@ impl<'a> Criteria<'a> {
         (to_keep, to_fetch)
     }
 
-    pub fn filter_on_query<'b>(
-        &self,
-        midx_list: &[&'b MailIndex<'b>],
-        query_result: &'b Vec<QueryResult>,
-    ) -> Result<Vec<&'b MailIndex<'b>>> {
-        Ok(midx_list
-            .iter()
-            .zip(query_result.iter())
-            .map(|(midx, qr)| MailView::new(qr, midx))
-            .collect::<Result<Vec<_>, _>>()?
-            .into_iter()
-            .filter(|mail_view| self.is_keep_on_query(mail_view))
-            .map(|mail_view| mail_view.in_idx)
-            .collect())
-    }
-
     // ----
 
     /// Here we are doing a partial filtering: we do not have access
@@ -213,7 +196,7 @@ impl<'a> Criteria<'a> {
     /// the email, as body(x) might be false. So we need to check it. But as seqid(x) is true,
     /// we could simplify the request to just body(x) and truncate the first OR. Today, we are
     /// not doing that, and thus we reevaluate everything.
-    fn is_keep_on_query(&self, mail_view: &MailView) -> bool {
+    pub fn is_keep_on_query(&self, mail_view: &MailView) -> bool {
         use SearchKey::*;
         match self.0 {
             // Combinator logic
diff --git a/src/mail/mailbox.rs b/src/mail/mailbox.rs
index c20d815..9190883 100644
--- a/src/mail/mailbox.rs
+++ b/src/mail/mailbox.rs
@@ -498,7 +498,7 @@ fn dump(uid_index: &Bayou<UidIndex>) {
 
 /// The metadata of a message that is stored in K2V
 /// at pk = mail/<mailbox uuid>, sk = <message uuid>
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct MailMeta {
     /// INTERNALDATE field (milliseconds since epoch)
     pub internaldate: u64,
diff --git a/src/mail/query.rs b/src/mail/query.rs
index a183c5a..cbc5fe6 100644
--- a/src/mail/query.rs
+++ b/src/mail/query.rs
@@ -2,7 +2,8 @@ use super::mailbox::MailMeta;
 use super::snapshot::FrozenMailbox;
 use super::unique_ident::UniqueIdent;
 use anyhow::Result;
-use futures::stream::{FuturesOrdered, StreamExt};
+use futures::stream::{Stream, StreamExt, BoxStream};
+use futures::future::FutureExt;
 
 /// Query is in charge of fetching efficiently
 /// requested data for a list of emails
@@ -28,64 +29,57 @@ impl QueryScope {
     }
 }
 
+//type QueryResultStream = Box<dyn Stream<Item = Result<QueryResult>>>;
+
 impl<'a, 'b> Query<'a, 'b> {
-    pub async fn fetch(&self) -> Result<Vec<QueryResult>> {
+    pub fn fetch(&self) -> BoxStream<Result<QueryResult>> {
         match self.scope {
-            QueryScope::Index => Ok(self
-                .emails
-                .iter()
-                .map(|&uuid| QueryResult::IndexResult { uuid })
-                .collect()),
-            QueryScope::Partial => self.partial().await,
-            QueryScope::Full => self.full().await,
+            QueryScope::Index => Box::pin(futures::stream::iter(self.emails).map(|&uuid| Ok(QueryResult::IndexResult { uuid }))),
+            QueryScope::Partial => Box::pin(self.partial()),
+            QueryScope::Full => Box::pin(self.full()),
         }
     }
 
     // --- functions below are private *for reasons*
+    fn partial<'d>(&'d self) -> impl Stream<Item = Result<QueryResult>> + 'd + Send {
+        async move { 
+            let maybe_meta_list: Result<Vec<MailMeta>> = self.frozen.mailbox.fetch_meta(self.emails).await;
+            let list_res = maybe_meta_list
+                .map(|meta_list| meta_list
+                     .into_iter()
+                     .zip(self.emails)
+                     .map(|(metadata, &uuid)| Ok(QueryResult::PartialResult { uuid, metadata }))
+                     .collect()
+                    )
+                .unwrap_or_else(|e| vec![Err(e)]);
 
-    async fn partial(&self) -> Result<Vec<QueryResult>> {
-        let meta = self.frozen.mailbox.fetch_meta(self.emails).await?;
-        let result = meta
-            .into_iter()
-            .zip(self.emails.iter())
-            .map(|(metadata, &uuid)| QueryResult::PartialResult { uuid, metadata })
-            .collect::<Vec<_>>();
-
-        Ok(result)
+            futures::stream::iter(list_res)
+        }.flatten_stream()
     }
 
-    /// @FIXME WARNING: THIS CAN ALLOCATE A LOT OF MEMORY
-    /// AND GENERATE SO MUCH NETWORK TRAFFIC.
-    /// THIS FUNCTION SHOULD BE REWRITTEN, FOR EXAMPLE WITH
-    /// SOMETHING LIKE AN ITERATOR
-    async fn full(&self) -> Result<Vec<QueryResult>> {
-        let meta_list = self.partial().await?;
-        meta_list
-            .into_iter()
-            .map(|meta| async move {
+    fn full<'d>(&'d self) -> impl Stream<Item = Result<QueryResult>> + 'd + Send {
+        self.partial()
+            .then(move |maybe_meta| async move {
+                let meta = maybe_meta?;
+
                 let content = self
                     .frozen
                     .mailbox
                     .fetch_full(
                         *meta.uuid(),
                         &meta
-                            .metadata()
-                            .expect("meta to be PartialResult")
-                            .message_key,
-                    )
+                        .metadata()
+                        .expect("meta to be PartialResult")
+                        .message_key,
+                        )
                     .await?;
 
                 Ok(meta.into_full(content).expect("meta to be PartialResult"))
             })
-            .collect::<FuturesOrdered<_>>()
-            .collect::<Vec<_>>()
-            .await
-            .into_iter()
-            .collect::<Result<Vec<_>, _>>()
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub enum QueryResult {
     IndexResult {
         uuid: UniqueIdent,