Add additional flagging functionality

More mail flags can now be set, unset, and toggled, not just the
read/seen flag.

This functionality is implemented with a new `:flag` and `:unflag`
command, which are extensions to the matching `:read` and `:unread`
commands, adding support for different flags.  In fact, the
`read`/`unread` commands are now recognized aliases to `flag`/`unflag`.
The new commands are also well documented in aerc(1).

The change mostly extends the previous read/unread setting functionality
by adding a selection for the flag to change.

7 files changed, 76 insertions, 91 deletions

diff --git a/worker/imap/flags.go b/worker/imap/flags.go
index 2625086..aef1019 100644
--- a/worker/imap/flags.go
+++ b/worker/imap/flags.go
@@ -76,12 +76,11 @@ func (imapw *IMAPWorker) handleAnsweredMessages(msg *types.AnsweredMessages) {
 	})
 }
 
-func (imapw *IMAPWorker) handleReadMessages(msg *types.ReadMessages) {
+func (imapw *IMAPWorker) handleFlagMessages(msg *types.FlagMessages) {
+	flags := []interface{}{flagToImap[msg.Flag]}
 	item := imap.FormatFlagsOp(imap.AddFlags, true)
-	flags := []interface{}{imap.SeenFlag}
-	if !msg.Read {
+	if !msg.Enable {
 		item = imap.FormatFlagsOp(imap.RemoveFlags, true)
-		flags = []interface{}{imap.SeenFlag}
 	}
 	uids := toSeqSet(msg.Uids)
 	emitErr := func(err error) {
@@ -99,7 +98,7 @@ func (imapw *IMAPWorker) handleReadMessages(msg *types.ReadMessages) {
 	}, func(_msg types.WorkerMessage) {
 		switch m := _msg.(type) {
 		case *types.Error:
-			err := fmt.Errorf("handleReadMessages: %v", m.Error)
+			err := fmt.Errorf("handleFlagMessages: %v", m.Error)
 			imapw.worker.Logger.Printf("could not fetch headers: %s", err)
 			emitErr(err)
 		case *types.Done:
diff --git a/worker/imap/worker.go b/worker/imap/worker.go
index a43ac49..0be51d7 100644
--- a/worker/imap/worker.go
+++ b/worker/imap/worker.go
@@ -173,8 +173,8 @@ func (w *IMAPWorker) handleMessage(msg types.WorkerMessage) error {
 		w.handleFetchFullMessages(msg)
 	case *types.DeleteMessages:
 		w.handleDeleteMessages(msg)
-	case *types.ReadMessages:
-		w.handleReadMessages(msg)
+	case *types.FlagMessages:
+		w.handleFlagMessages(msg)
 	case *types.AnsweredMessages:
 		w.handleAnsweredMessages(msg)
 	case *types.CopyMessages:
diff --git a/worker/maildir/message.go b/worker/maildir/message.go
index 5c6c930..55e2613 100644
--- a/worker/maildir/message.go
+++ b/worker/maildir/message.go
@@ -53,43 +53,29 @@ func (m Message) SetFlags(flags []maildir.Flag) error {
 	return m.dir.SetFlags(m.key, flags)
 }
 
-// MarkReplied either adds or removes the maildir.FlagReplied flag from the
-// message.
-func (m Message) MarkReplied(answered bool) error {
+// SetOneFlag enables or disables a single message flag on the message.
+func (m Message) SetOneFlag(flag maildir.Flag, enable bool) error {
 	flags, err := m.Flags()
 	if err != nil {
 		return fmt.Errorf("could not read previous flags: %v", err)
 	}
-	if answered {
-		flags = append(flags, maildir.FlagReplied)
+	if enable {
+		flags = append(flags, flag)
 		return m.SetFlags(flags)
 	}
 	var newFlags []maildir.Flag
-	for _, flag := range flags {
-		if flag != maildir.FlagReplied {
-			newFlags = append(newFlags, flag)
+	for _, oldFlag := range flags {
+		if oldFlag != flag {
+			newFlags = append(newFlags, oldFlag)
 		}
 	}
 	return m.SetFlags(newFlags)
 }
 
-// MarkRead either adds or removes the maildir.FlagSeen flag from the message.
-func (m Message) MarkRead(seen bool) error {
-	flags, err := m.Flags()
-	if err != nil {
-		return fmt.Errorf("could not read previous flags: %v", err)
-	}
-	if seen {
-		flags = append(flags, maildir.FlagSeen)
-		return m.SetFlags(flags)
-	}
-	var newFlags []maildir.Flag
-	for _, flag := range flags {
-		if flag != maildir.FlagSeen {
-			newFlags = append(newFlags, flag)
-		}
-	}
-	return m.SetFlags(newFlags)
+// MarkReplied either adds or removes the maildir.FlagReplied flag from the
+// message.
+func (m Message) MarkReplied(answered bool) error {
+	return m.SetOneFlag(maildir.FlagReplied, answered)
 }
 
 // Remove deletes the email immediately.
diff --git a/worker/maildir/worker.go b/worker/maildir/worker.go
index f14672e..467f2aa 100644
--- a/worker/maildir/worker.go
+++ b/worker/maildir/worker.go
@@ -193,8 +193,8 @@ func (w *Worker) handleMessage(msg types.WorkerMessage) error {
 		return w.handleFetchFullMessages(msg)
 	case *types.DeleteMessages:
 		return w.handleDeleteMessages(msg)
-	case *types.ReadMessages:
-		return w.handleReadMessages(msg)
+	case *types.FlagMessages:
+		return w.handleFlagMessages(msg)
 	case *types.AnsweredMessages:
 		return w.handleAnsweredMessages(msg)
 	case *types.CopyMessages:
@@ -473,7 +473,7 @@ func (w *Worker) handleAnsweredMessages(msg *types.AnsweredMessages) error {
 	return nil
 }
 
-func (w *Worker) handleReadMessages(msg *types.ReadMessages) error {
+func (w *Worker) handleFlagMessages(msg *types.FlagMessages) error {
 	for _, uid := range msg.Uids {
 		m, err := w.c.Message(*w.selected, uid)
 		if err != nil {
@@ -481,8 +481,9 @@ func (w *Worker) handleReadMessages(msg *types.ReadMessages) error {
 			w.err(msg, err)
 			continue
 		}
-		if err := m.MarkRead(msg.Read); err != nil {
-			w.worker.Logger.Printf("could not mark message as read: %v", err)
+		flag := flagToMaildir[msg.Flag]
+		if err := m.SetOneFlag(flag, msg.Enable); err != nil {
+			w.worker.Logger.Printf("could change flag %v to %v on message: %v", flag, msg.Enable, err)
 			w.err(msg, err)
 			continue
 		}
diff --git a/worker/notmuch/message.go b/worker/notmuch/message.go
index 6529507..3d3ed0c 100644
--- a/worker/notmuch/message.go
+++ b/worker/notmuch/message.go
@@ -64,70 +64,68 @@ func (m *Message) NewBodyPartReader(requestedParts []int) (io.Reader, error) {
 	return lib.FetchEntityPartReader(msg, requestedParts)
 }
 
-// MarkAnswered either adds or removes the "replied" tag from the message.
-func (m *Message) MarkAnswered(answered bool) error {
-	haveReplied := false
+// SetOneFlag adds or removes a single flag from the message.
+// Notmuch doesn't support all the flags, and for those this errors.
+func (m *Message) SetOneFlag(flag models.Flag, enable bool) error {
+	// Translate the flag into a notmuch tag, ignoring no-op flags.
+	var tag string
+	switch flag {
+	case models.SeenFlag:
+		// Note: Inverted properly later
+		tag = "unread"
+	case models.AnsweredFlag:
+		tag = "replied"
+	case models.FlaggedFlag:
+		tag = "flagged"
+	default:
+		return fmt.Errorf("Notmuch doesn't support flag %v", flag)
+	}
+
+	// Get the current state of the flag.
+	// Note that notmuch handles models.SeenFlag in an inverted sense.
+	oldState := false
 	tags, err := m.Tags()
 	if err != nil {
 		return err
 	}
 	for _, t := range tags {
-		if t == "replied" {
-			haveReplied = true
+		if t == tag {
+			oldState = true
 			break
 		}
 	}
-	if haveReplied == answered {
-		// we already have the desired state
-		return nil
+	if flag == models.SeenFlag {
+		oldState = !oldState
 	}
 
-	if haveReplied {
-		err := m.RemoveTag("replied")
-		if err != nil {
-			return err
-		}
+	// Skip if flag already in correct state.
+	if oldState == enable {
 		return nil
 	}
 
-	err = m.AddTag("replied")
-	if err != nil {
-		return err
+	if !enable {
+		if flag == models.SeenFlag {
+			return m.AddTag("unread")
+		} else {
+			return m.RemoveTag(tag)
+		}
+	} else {
+		if flag == models.SeenFlag {
+			return m.RemoveTag("unread")
+		} else {
+			return m.AddTag(tag)
+		}
 	}
-	return nil
+}
+
+// MarkAnswered either adds or removes the "replied" tag from the message.
+func (m *Message) MarkAnswered(answered bool) error {
+	return m.SetOneFlag(models.AnsweredFlag, answered)
 }
 
 // MarkRead either adds or removes the maildir.FlagSeen flag from the message.
 func (m *Message) MarkRead(seen bool) error {
-	haveUnread := false
-	tags, err := m.Tags()
-	if err != nil {
-		return err
-	}
-	for _, t := range tags {
-		if t == "unread" {
-			haveUnread = true
-			break
-		}
-	}
-	if (haveUnread && !seen) || (!haveUnread && seen) {
-		// we already have the desired state
-		return nil
-	}
-
-	if haveUnread {
-		err := m.RemoveTag("unread")
-		if err != nil {
-			return err
-		}
-		return nil
-	}
-
-	err = m.AddTag("unread")
-	if err != nil {
-		return err
-	}
-	return nil
+	return m.SetOneFlag(models.SeenFlag, seen)
 }
 
 // tags returns the notmuch tags of a message
diff --git a/worker/notmuch/worker.go b/worker/notmuch/worker.go
index 3fdd11d..f14b7ff 100644
--- a/worker/notmuch/worker.go
+++ b/worker/notmuch/worker.go
@@ -113,8 +113,8 @@ func (w *worker) handleMessage(msg types.WorkerMessage) error {
 		return w.handleFetchMessageBodyPart(msg)
 	case *types.FetchFullMessages:
 		return w.handleFetchFullMessages(msg)
-	case *types.ReadMessages:
-		return w.handleReadMessages(msg)
+	case *types.FlagMessages:
+		return w.handleFlagMessages(msg)
 	case *types.AnsweredMessages:
 		return w.handleAnsweredMessages(msg)
 	case *types.SearchDirectory:
@@ -392,7 +392,7 @@ func (w *worker) handleAnsweredMessages(msg *types.AnsweredMessages) error {
 	return nil
 }
 
-func (w *worker) handleReadMessages(msg *types.ReadMessages) error {
+func (w *worker) handleFlagMessages(msg *types.FlagMessages) error {
 	for _, uid := range msg.Uids {
 		m, err := w.msgFromUid(uid)
 		if err != nil {
@@ -400,8 +400,8 @@ func (w *worker) handleReadMessages(msg *types.ReadMessages) error {
 			w.err(msg, err)
 			continue
 		}
-		if err := m.MarkRead(msg.Read); err != nil {
-			w.w.Logger.Printf("could not mark message as read: %v", err)
+		if err := m.SetOneFlag(msg.Flag, msg.Enable); err != nil {
+			w.w.Logger.Printf("could not set flag %v as %v for message: %v", msg.Flag, msg.Enable, err)
 			w.err(msg, err)
 			continue
 		}
diff --git a/worker/types/messages.go b/worker/types/messages.go
index 475a7aa..c044667 100644
--- a/worker/types/messages.go
+++ b/worker/types/messages.go
@@ -113,10 +113,11 @@ type DeleteMessages struct {
 	Uids []uint32
 }
 
-// Marks messages as read or unread
-type ReadMessages struct {
+// Flag messages with different mail types
+type FlagMessages struct {
 	Message
-	Read bool
+	Enable bool
+	Flag models.Flag
 	Uids []uint32
 }