4 files changed, 45 insertions, 24 deletions
diff --git a/src/imap/command/selected.rs b/src/imap/command/selected.rs
index a6fa645..1978729 100644
--- a/src/imap/command/selected.rs
+++ b/src/imap/command/selected.rs
@@ -62,7 +62,12 @@ impl<'a> SelectedContext<'a> {
     }
 
     async fn expunge(self) -> Result<(Response, flow::Transition)> {
-        Ok((Response::bad("Not implemented")?, flow::Transition::None))
+        let data = self.mailbox.expunge().await?;
+
+        Ok((
+            Response::ok("EXPUNGE completed")?.with_body(data),
+            flow::Transition::None,
+        ))
     }
 
     async fn store(
diff --git a/src/imap/mailbox_view.rs b/src/imap/mailbox_view.rs
index 93d3b4d..f293bfa 100644
--- a/src/imap/mailbox_view.rs
+++ b/src/imap/mailbox_view.rs
@@ -169,29 +169,7 @@ impl MailboxView {
                     self.mailbox.del_flags(*uuid, &flags[..]).await?;
                 }
                 StoreType::Replace => {
-                    let old_flags = &self
-                        .known_state
-                        .table
-                        .get(uuid)
-                        .ok_or(anyhow!(
-                            "Missing message: {} (UID {}, UUID {})",
-                            i,
-                            uid,
-                            uuid
-                        ))?
-                        .1;
-                    let to_remove = old_flags
-                        .iter()
-                        .filter(|x| !flags.contains(&x))
-                        .cloned()
-                        .collect::<Vec<_>>();
-                    let to_add = flags
-                        .iter()
-                        .filter(|x| !old_flags.contains(&x))
-                        .cloned()
-                        .collect::<Vec<_>>();
-                    self.mailbox.add_flags(*uuid, &to_add[..]).await?;
-                    self.mailbox.del_flags(*uuid, &to_remove[..]).await?;
+                    self.mailbox.set_flags(*uuid, &flags[..]).await?;
                 }
             }
         }
@@ -199,6 +177,10 @@ impl MailboxView {
         self.update().await
     }
 
+    pub async fn expunge(&mut self) -> Result<Vec<Body>> {
+        unimplemented!()
+    }
+
     /// Looks up state changes in the mailbox and produces a set of IMAP
     /// responses describing the new state.
     pub async fn fetch(
diff --git a/src/mail/mailbox.rs b/src/mail/mailbox.rs
index 0e8af70..4b84cf2 100644
--- a/src/mail/mailbox.rs
+++ b/src/mail/mailbox.rs
@@ -95,6 +95,11 @@ impl Mailbox {
         self.mbox.write().await.del_flags(id, flags).await
     }
 
+    /// Define the new flags for this message
+    pub async fn set_flags<'a>(&self, id: UniqueIdent, flags: &[Flag]) -> Result<()> {
+        self.mbox.write().await.set_flags(id, flags).await
+    }
+
     /// Insert an email into the mailbox
     pub async fn append<'a>(
         &self,
@@ -265,6 +270,11 @@ impl MailboxInternal {
         self.uid_index.push(del_flag_op).await
     }
 
+    async fn set_flags(&mut self, ident: UniqueIdent, flags: &[Flag]) -> Result<()> {
+        let set_flag_op = self.uid_index.state().op_flag_set(ident, flags.to_vec());
+        self.uid_index.push(set_flag_op).await
+    }
+
     async fn append(
         &mut self,
         mail: IMF<'_>,
diff --git a/src/mail/uidindex.rs b/src/mail/uidindex.rs
index 6da08c1..3a3e252 100644
--- a/src/mail/uidindex.rs
+++ b/src/mail/uidindex.rs
@@ -36,6 +36,7 @@ pub enum UidIndexOp {
     MailDel(UniqueIdent),
     FlagAdd(UniqueIdent, Vec<Flag>),
     FlagDel(UniqueIdent, Vec<Flag>),
+    FlagSet(UniqueIdent, Vec<Flag>),
     BumpUidvalidity(u32),
 }
 
@@ -61,6 +62,11 @@ impl UidIndex {
     }
 
     #[must_use]
+    pub fn op_flag_set(&self, ident: UniqueIdent, flags: Vec<Flag>) -> UidIndexOp {
+        UidIndexOp::FlagSet(ident, flags)
+    }
+
+    #[must_use]
     pub fn op_bump_uidvalidity(&self, count: u32) -> UidIndexOp {
         UidIndexOp::BumpUidvalidity(count)
     }
@@ -162,6 +168,24 @@ impl BayouState for UidIndex {
                     new.idx_by_flag.remove(*uid, rm_flags);
                 }
             }
+            UidIndexOp::FlagSet(ident, new_flags) => {
+                if let Some((uid, existing_flags)) = new.table.get_mut(ident) {
+                    // Remove flags from the source of trust and the cache
+                    let (keep_flags, rm_flags): (Vec<String>, Vec<String>) = existing_flags
+                        .iter()
+                        .cloned()
+                        .partition(|x| new_flags.contains(x));
+                    *existing_flags = keep_flags;
+                    let mut to_add: Vec<Flag> = new_flags
+                        .iter()
+                        .filter(|f| !existing_flags.contains(f))
+                        .cloned()
+                        .collect();
+                    existing_flags.append(&mut to_add);
+                    new.idx_by_flag.remove(*uid, &rm_flags);
+                    new.idx_by_flag.insert(*uid, &to_add);
+                }
+            }
             UidIndexOp::BumpUidvalidity(count) => {
                 new.uidvalidity = ImapUidvalidity::new(new.uidvalidity.get() + *count)
                     .unwrap_or(ImapUidvalidity::new(u32::MAX).unwrap());