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.

11 files changed, 242 insertions, 184 deletions

diff --git a/commands/msg/read.go b/commands/msg/read.go
index 1e264c27..08881894 100644
--- a/commands/msg/read.go
+++ b/commands/msg/read.go
@@ -1,9 +1,9 @@
 package msg
 
 import (
-	"errors"
 	"sync"
 	"time"
+	"fmt"
 
 	"git.sr.ht/~sircmpwn/getopt"
 
@@ -13,36 +13,106 @@ import (
 	"git.sr.ht/~sircmpwn/aerc/worker/types"
 )
 
-type Read struct{}
+type FlagMsg struct{}
 
 func init() {
-	register(Read{})
+	register(FlagMsg{})
 }
 
-func (Read) Aliases() []string {
-	return []string{"read", "unread"}
+func (FlagMsg) Aliases() []string {
+	return []string{"flag", "unflag", "read", "unread"}
 }
 
-func (Read) Complete(aerc *widgets.Aerc, args []string) []string {
+func (FlagMsg) Complete(aerc *widgets.Aerc, args []string) []string {
 	return nil
 }
 
-func (Read) Execute(aerc *widgets.Aerc, args []string) error {
-	opts, optind, err := getopt.Getopts(args, "t")
+// If this was called as 'flag' or 'unflag', without the toggle (-t)
+// option, then it will flag the corresponding messages with the given
+// flag.  If the toggle option was given, it will individually toggle
+// the given flag for the corresponding messages.
+//
+// If this was called as 'read' or 'unread', it has the same effect as
+// 'flag' or 'unflag', respectively, but the 'Seen' flag is affected.
+func (FlagMsg) Execute(aerc *widgets.Aerc, args []string) error {
+
+	// The flag to change
+	var flag models.Flag
+	// User-readable name of the flag to change
+	var flagName string
+	// Whether to toggle the flag (true) or to enable/disable it (false)
+	var toggle bool
+	// Whether to enable (true) or disable (false) the flag
+	enable := (args[0] == "read" || args[0] == "flag")
+	// User-readable name for the action being performed
+	var actionName string
+	// Getopt option string, varies by command name
+	var getoptString string
+	// Help message to provide on parsing failure
+	var helpMessage string
+	// Used during parsing to prevent choosing a flag muliple times
+	// A default flag will be used if this is false
+	flagChosen := false
+
+	if args[0] == "read" || args[0] == "unread" {
+		flag = models.SeenFlag
+		flagName = "read"
+		getoptString = "t"
+		helpMessage = "Usage: " + args[0] + " [-t]"
+	} else { // 'flag' / 'unflag'
+		flag = models.FlaggedFlag
+		flagName = "flagged"
+		getoptString = "tax:"
+		helpMessage = "Usage: " + args[0] + " [-t] [-a | -x <flag>]"
+	}
+
+	opts, optind, err := getopt.Getopts(args, getoptString)
 	if err != nil {
 		return err
 	}
-	if optind != len(args) {
-		return errors.New("Usage: " + args[0] + " [-t]")
-	}
-	var toggle bool
-
 	for _, opt := range opts {
 		switch opt.Option {
 		case 't':
 			toggle = true
+		case 'a':
+			if flagChosen {
+				return fmt.Errorf("Cannot choose a flag multiple times! " + helpMessage)
+			}
+			flag = models.AnsweredFlag
+			flagName = "answered"
+			flagChosen = true
+		case 'x':
+			if flagChosen {
+				return fmt.Errorf("Cannot choose a flag multiple times! " + helpMessage)
+			}
+			// TODO: Support all flags?
+			switch opt.Value {
+			case "Seen":
+				flag = models.SeenFlag
+				flagName = "seen"
+			case "Answered":
+				flag = models.AnsweredFlag
+				flagName = "answered"
+			case "Flagged":
+				flag = models.FlaggedFlag
+				flagName = "flagged"
+			default:
+				return fmt.Errorf("Unknown / Prohibited flag \"%v\"", opt.Value)
+			}
+			flagChosen = true
 		}
 	}
+	if toggle {
+		actionName = "Toggling"
+	} else if enable {
+		actionName = "Setting"
+	} else {
+		actionName = "Unsetting"
+	}
+	if optind != len(args) {
+		// Any non-option arguments: Error
+		return fmt.Errorf(helpMessage)
+	}
 
 	h := newHelper(aerc)
 	store, err := h.store()
@@ -50,57 +120,68 @@ func (Read) Execute(aerc *widgets.Aerc, args []string) error {
 		return err
 	}
 
-	if toggle {
-		// ignore command given, simply toggle all the read states
-		return submitToggle(aerc, store, h)
-	}
-	msgUids, err := h.markedOrSelectedUids()
-	if err != nil {
-		return err
-	}
-	switch args[0] {
-	case "read":
-		submitReadChange(aerc, store, msgUids, true)
-	case "unread":
-		submitReadChange(aerc, store, msgUids, false)
-
-	}
-	return nil
-}
+	// UIDs of messages to enable or disable the flag for.
+	var toEnable []uint32
+	var toDisable []uint32
 
-func splitMessages(msgs []*models.MessageInfo) (read []uint32, unread []uint32) {
-	for _, m := range msgs {
-		var seen bool
-		for _, flag := range m.Flags {
-			if flag == models.SeenFlag {
-				seen = true
-				break
+	if toggle {
+		// If toggling, split messages into those that need to
+		// be enabled / disabled.
+		msgs, err := h.messages()
+		if err != nil {
+			return err
+		}
+		for _, m := range msgs {
+			var enabled bool
+			for _, mFlag := range m.Flags {
+				if mFlag == flag {
+					enabled = true
+					break
+				}
 			}
+			if enabled {
+				toDisable = append(toDisable, m.Uid)
+			} else {
+				toEnable = append(toEnable, m.Uid)
+			}
+		}
+	} else {
+		msgUids, err := h.markedOrSelectedUids()
+		if err != nil {
+			return err
 		}
-		if seen {
-			read = append(read, m.Uid)
+		if enable {
+			toEnable = msgUids
 		} else {
-			unread = append(unread, m.Uid)
+			toDisable = msgUids
 		}
 	}
-	return read, unread
-}
 
-func submitReadChange(aerc *widgets.Aerc, store *lib.MessageStore,
-	uids []uint32, newState bool) {
-	store.Read(uids, newState, func(msg types.WorkerMessage) {
-		switch msg := msg.(type) {
-		case *types.Done:
-			aerc.PushStatus(msg_success, 10*time.Second)
-		case *types.Error:
-			aerc.PushError(" " + msg.Error.Error())
+	var wg sync.WaitGroup
+	success := true
+
+	if len(toEnable) != 0 {
+		submitFlagChange(aerc, store, toEnable, flag, true, &wg, &success)
+	}
+	if len(toDisable) != 0 {
+		submitFlagChange(aerc, store, toDisable, flag, false, &wg, &success)
+	}
+
+	// We need to do flagging in the background, else we block the main thread
+	go func() {
+		wg.Wait()
+		if success {
+			aerc.PushStatus(actionName + " flag '" + flagName + "' successful", 10*time.Second)
 		}
-	})
+	}()
+
+	return nil
 }
 
-func submitReadChangeWg(aerc *widgets.Aerc, store *lib.MessageStore,
-	uids []uint32, newState bool, wg *sync.WaitGroup, success *bool) {
-	store.Read(uids, newState, func(msg types.WorkerMessage) {
+func submitFlagChange(aerc *widgets.Aerc, store *lib.MessageStore,
+	uids []uint32, flag models.Flag, newState bool,
+	wg *sync.WaitGroup, success *bool) {
+	store.Flag(uids, flag, newState, func(msg types.WorkerMessage) {
 		wg.Add(1)
 		switch msg := msg.(type) {
 		case *types.Done:
@@ -112,35 +193,3 @@ func submitReadChangeWg(aerc *widgets.Aerc, store *lib.MessageStore,
 		}
 	})
 }
-
-func submitToggle(aerc *widgets.Aerc, store *lib.MessageStore, h *helper) error {
-	msgs, err := h.messages()
-	if err != nil {
-		return err
-	}
-	read, unread := splitMessages(msgs)
-
-	var wg sync.WaitGroup
-	success := true
-
-	if len(read) != 0 {
-		newState := false
-		submitReadChangeWg(aerc, store, read, newState, &wg, &success)
-	}
-
-	if len(unread) != 0 {
-		newState := true
-		submitReadChangeWg(aerc, store, unread, newState, &wg, &success)
-	}
-	// we need to do that in the background, else we block the main thread
-	go func() {
-		wg.Wait()
-		if success {
-			aerc.PushStatus(msg_success, 10*time.Second)
-		}
-	}()
-	return nil
-
-}
-
-const msg_success = "read state set successfully"
diff --git a/doc/aerc.1.scd b/doc/aerc.1.scd
index 61e35bbc..929aaa9c 100644
--- a/doc/aerc.1.scd
+++ b/doc/aerc.1.scd
@@ -136,16 +136,39 @@ message list, the message in the message viewer, etc).
 
 	*-T*: Use the specified template file for creating the initial message body
 
-*read*
+*read* [-t]
 	Marks the marked or selected messages as read.
 
 	*-t*: Toggle the messages between read and unread.
 
-*unread*
+*unread* [-t]
 	Marks the marked or selected messages as unread.
 
 	*-t*: Toggle the messages between read and unread.
 
+*flag* [-t] [-a | -x <flag>]
+	Sets (enables) a certain flag on the marked or selected messages.
+
+	*-t*: Toggle the flag instead of setting (enabling) it.
+
+	*-a*: Mark message as answered/unanswered.
+
+	*-x <flag>*: Mark message with specific flag.
+
+	The available flags are (adapted from RFC 3501, section 2.3.2):
+
+	Seen
+		Message has been read
+	
+	Answered
+		Message has been answered
+	
+	Flagged
+		Message is flagged for urgent/special attention
+
+*unflag* [-t] <flag>
+	Operates exactly like *flag*, defaulting to unsetting (disabling) flags.
+
 *modify-labels* <[+-]label>...
 	Modify message labels (e.g. notmuch tags). Labels prefixed with a '+' are
 	added, those prefixed with a '-' removed. As a convenience, labels without
diff --git a/lib/messageview.go b/lib/messageview.go
index 384a9472..08ea92fb 100644
--- a/lib/messageview.go
+++ b/lib/messageview.go
@@ -94,7 +94,7 @@ func NewMessageStoreView(messageInfo *models.MessageInfo,
 	} else {
 		cb(msv, nil)
 	}
-	store.Read([]uint32{messageInfo.Uid}, true, nil)
+	store.Flag([]uint32{messageInfo.Uid}, models.SeenFlag, true, nil)
 }
 
 func (msv *MessageStoreView) MessageInfo() *models.MessageInfo {
diff --git a/lib/msgstore.go b/lib/msgstore.go
index 86215a72..dda15d3b 100644
--- a/lib/msgstore.go
+++ b/lib/msgstore.go
@@ -333,11 +333,12 @@ func (store *MessageStore) Move(uids []uint32, dest string, createDest bool,
 	store.update()
 }
 
-func (store *MessageStore) Read(uids []uint32, read bool,
-	cb func(msg types.WorkerMessage)) {
+func (store *MessageStore) Flag(uids []uint32, flag models.Flag,
+	enable bool, cb func(msg types.WorkerMessage)) {
 
-	store.worker.PostAction(&types.ReadMessages{
-		Read: read,
+	store.worker.PostAction(&types.FlagMessages{
+		Enable: enable,
+		Flag: flag,
 		Uids: uids,
 	}, cb)
 }
diff --git a/worker/imap/flags.go b/worker/imap/flags.go
index 26250866..aef10191 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 a43ac49e..0be51d71 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 5c6c9307..55e2613b 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 f14672e5..467f2aa0 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 65295073..3d3ed0c3 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 3fdd11d2..f14b7ff0 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 475a7aa3..c0446671 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
 }