Rework threading and add REFERENCES

* Implement a simplified version of the REFERENCES algorithm
* Remove FormatThreads function
* Instead of acting on all threads, handle each thread independently

Signed-off-by: Kevin Kuehler <keur@xcf.berkeley.edu>

5 files changed, 278 insertions, 77 deletions

diff --git a/lib/msgstore.go b/lib/msgstore.go
index e2e968b..ccde2c2 100644
--- a/lib/msgstore.go
+++ b/lib/msgstore.go
@@ -1,12 +1,15 @@
 package lib
 
 import (
+	"fmt"
 	"io"
+	"sort"
 	"time"
 
-	"git.sr.ht/~sircmpwn/aerc/lib/sort"
+	aercSort "git.sr.ht/~sircmpwn/aerc/lib/sort"
 	"git.sr.ht/~sircmpwn/aerc/models"
 	"git.sr.ht/~sircmpwn/aerc/worker/types"
+	"github.com/emersion/go-message"
 )
 
 // Accesses to fields must be guarded by MessageStore.Lock/Unlock
@@ -15,15 +18,17 @@ type MessageStore struct {
 	DirInfo  models.DirectoryInfo
 	Messages map[uint32]*models.MessageInfo
 	// Ordered list of known UIDs
-	uids    []uint32
-	Threads []*types.Thread
+	uids []uint32
 
 	selected        int
 	bodyCallbacks   map[uint32][]func(io.Reader)
 	headerCallbacks map[uint32][]func(*types.MessageInfo)
 
 	// If set, messages in the mailbox will be threaded
-	thread bool
+	ThreadRoot  *types.Thread
+	FlatThreads []*types.Thread
+	thread      bool
+	threadRefs  map[string]*types.Thread
 
 	// Search/filter results
 	results     []uint32
@@ -57,8 +62,9 @@ func NewMessageStore(worker *types.Worker,
 		selected:        0,
 		bodyCallbacks:   make(map[uint32][]func(io.Reader)),
 		headerCallbacks: make(map[uint32][]func(*types.MessageInfo)),
-
-		thread: thread,
+		ThreadRoot:      &types.Thread{Uid: 0, Dummy: true},
+		thread:          thread,
+		threadRefs:      make(map[string]*types.Thread),
 
 		defaultSortCriteria: defaultSortCriteria,
 
@@ -175,36 +181,47 @@ func (store *MessageStore) Update(msg types.WorkerMessage) {
 		}
 		update = true
 	case *types.DirectoryThreaded:
-		var uids []uint32
+		var (
+			uids      []uint32
+			flattened []*types.Thread
+		)
 		newMap := make(map[uint32]*models.MessageInfo)
 
-		for i := len(msg.Threads) - 1; i >= 0; i-- {
-			msg.Threads[i].FormatThread(func(t *types.Thread, x []rune) bool {
-				uid := t.Uid
-				uids = append([]uint32{uid}, uids...)
-				if msg, ok := store.Messages[uid]; ok {
-					newMap[uid] = msg
-				} else {
-					newMap[uid] = nil
-					directoryChange = true
-				}
-				return false
-			})
-		}
+		msg.ThreadRoot.Traverse(false, func(t *types.Thread) bool {
+			uid := t.Uid
+			uids = append([]uint32{uid}, uids...)
+			flattened = append([]*types.Thread{t}, flattened...)
+			if msg, ok := store.Messages[uid]; ok {
+				newMap[uid] = msg
+			} else {
+				newMap[uid] = nil
+				directoryChange = true
+			}
+			return false
+		})
 		store.Messages = newMap
 		store.uids = uids
-		store.Threads = msg.Threads
+		store.FlatThreads = flattened
+		store.ThreadRoot = msg.ThreadRoot
 		update = true
 	case *types.DirectoryContents:
+		var needsHeaders []uint32
 		newMap := make(map[uint32]*models.MessageInfo)
 		for _, uid := range msg.Uids {
 			if msg, ok := store.Messages[uid]; ok {
 				newMap[uid] = msg
 			} else {
 				newMap[uid] = nil
+				needsHeaders = append(needsHeaders, uid)
 				directoryChange = true
 			}
 		}
+		if store.thread {
+			// We need the headers to perform references. Grab them all for
+			// now. We can probably be smarter here, but let's get something
+			// working first.
+			store.FetchHeaders(needsHeaders, nil)
+		}
 		store.Messages = newMap
 		store.uids = msg.Uids
 		update = true
@@ -234,6 +251,9 @@ func (store *MessageStore) Update(msg types.WorkerMessage) {
 				}
 			}
 		}
+		if store.thread {
+			store.threadMessage(msg.Info)
+		}
 		update = true
 	case *types.FullMessage:
 		if _, ok := store.pendingBodies[msg.Content.Uid]; ok {
@@ -248,19 +268,49 @@ func (store *MessageStore) Update(msg types.WorkerMessage) {
 	case *types.MessagesDeleted:
 		toDelete := make(map[uint32]interface{})
 		for _, uid := range msg.Uids {
+			if store.thread {
+				if needle := store.ThreadRoot.Find(uid); needle != nil {
+					_msg := store.Messages[uid]
+					delete(store.threadRefs, _msg.Envelope.MessageId)
+					needle.Dummy = true
+				}
+			}
 			toDelete[uid] = nil
 			delete(store.Messages, uid)
 			delete(store.Deleted, uid)
 		}
 		uids := make([]uint32, len(store.uids)-len(msg.Uids))
-		j := 0
-		for _, uid := range store.uids {
-			if _, deleted := toDelete[uid]; !deleted && j < len(uids) {
-				uids[j] = uid
-				j += 1
+		if store.thread {
+			flattened := make([]*types.Thread, len(store.FlatThreads)-len(msg.Uids))
+			j := 0
+			for _, uid := range store.uids {
+				if _, deleted := toDelete[uid]; !deleted && j < len(uids) {
+					uids[j] = uid
+					j += 1
+				}
+			}
+			j = 0
+			for _, t := range store.FlatThreads {
+				uid := t.Uid
+				if _, deleted := toDelete[uid]; !deleted && j < len(flattened) {
+					flattened[j] = t
+					j += 1
+				}
 			}
+			fmt.Printf("DELETE UID: prev: %d, new: %d\n", len(store.uids), len(uids))
+			fmt.Printf("DELETE FLAT: prev: %d, new: %d\n", len(store.FlatThreads), len(flattened))
+			store.uids = uids
+			store.FlatThreads = flattened
+		} else {
+			j := 0
+			for _, uid := range store.uids {
+				if _, deleted := toDelete[uid]; !deleted && j < len(uids) {
+					uids[j] = uid
+					j += 1
+				}
+			}
+			store.uids = uids
 		}
-		store.uids = uids
 		update = true
 	}
 
@@ -273,6 +323,96 @@ func (store *MessageStore) Update(msg types.WorkerMessage) {
 	}
 }
 
+func (store *MessageStore) threadMessage(msg *models.MessageInfo) {
+	var (
+		fields message.HeaderFields
+
+		childThread *types.Thread
+		irt         *types.Thread
+		roots       []*types.Thread
+	)
+	if msg.Envelope == nil {
+		return
+	}
+
+	newRefs := make(map[string]*types.Thread)
+
+	if thread, ok := store.threadRefs[msg.Envelope.MessageId]; ok {
+		// Are we in the references table as someone else's parent?
+		thread.Dummy = false
+		thread.Uid = msg.Uid
+		childThread = thread
+	} else {
+		// Then we create a new thread
+		childThread = &types.Thread{Uid: msg.Uid}
+	}
+
+	newRefs[msg.Envelope.MessageId] = childThread
+
+	fields = msg.RFC822Headers.FieldsByKey("In-Reply-To")
+	if fields.Next() {
+		inReplyHeader, err := fields.Text()
+		if err != nil {
+			return
+		}
+		if p, ok := store.threadRefs[inReplyHeader]; ok {
+			irt = p
+		} else {
+			irt = &types.Thread{Uid: 0, Dummy: true}
+		}
+		childThread.Parent = irt
+		irt.Children = append(irt.Children, childThread)
+		newRefs[inReplyHeader] = irt
+	}
+
+	for r, t := range store.threadRefs {
+		if _, ok := newRefs[r]; !ok {
+			newRefs[r] = t
+		}
+	}
+
+	for _, t := range newRefs {
+		if t.Parent == nil || t.Parent == store.ThreadRoot {
+			roots = append(roots, t)
+			t.Parent = store.ThreadRoot
+		}
+	}
+	store.ThreadRoot.Children = roots
+
+	var (
+		uids      []uint32
+		flattened []*types.Thread
+	)
+
+	if len(store.pendingHeaders) == 0 {
+		// Sort the root of the tree
+		children := store.ThreadRoot.Children
+		sort.Slice(children, func(i, j int) bool {
+			ci, cj := children[i], children[j]
+			if ci.Dummy {
+				ci = ci.Children[0]
+			}
+			if cj.Dummy {
+				cj = cj.Children[0]
+			}
+			mi, mj := store.Messages[ci.Uid], store.Messages[cj.Uid]
+			return mi.InternalDate.After(mj.InternalDate)
+		})
+
+		// Linearize tree
+		store.ThreadRoot.Traverse(false, func(t *types.Thread) bool {
+			uid := t.Uid
+			uids = append([]uint32{uid}, uids...)
+			flattened = append(flattened, t)
+			return false
+		})
+	}
+
+	store.FlatThreads = flattened
+	store.threadRefs = newRefs
+	store.uids = uids
+}
+
 func (store *MessageStore) OnUpdate(fn func(store *MessageStore)) {
 	store.onUpdate = fn
 }
@@ -418,7 +558,7 @@ func (store *MessageStore) Search(args []string, cb func([]uint32)) {
 	}, func(msg types.WorkerMessage) {
 		switch msg := msg.(type) {
 		case *types.SearchResults:
-			sort.SortBy(msg.Uids, store.uids)
+			aercSort.SortBy(msg.Uids, store.uids)
 			cb(msg.Uids)
 		}
 	})
diff --git a/widgets/msglist.go b/widgets/msglist.go
index 228477d..df0366d 100644
--- a/widgets/msglist.go
+++ b/widgets/msglist.go
@@ -12,7 +12,7 @@ import (
 	"git.sr.ht/~sircmpwn/aerc/lib/format"
 	"git.sr.ht/~sircmpwn/aerc/lib/ui"
 	"git.sr.ht/~sircmpwn/aerc/models"
-	"git.sr.ht/~sircmpwn/aerc/worker/types"
+	//"git.sr.ht/~sircmpwn/aerc/worker/types"
 )
 
 type MessageList struct {
@@ -71,16 +71,16 @@ func (ml *MessageList) Draw(ctx *ui.Context) {
 	uids := store.Uids()
 
 	if ml.conf.Ui.ThreadingEnabled {
-		threads := store.Threads
+		if store.ThreadRoot == nil {
+			return
+		}
 
-		for i := len(threads) - 1; i >= 0; i-- {
-			threads[i].FormatThread(func(thread *types.Thread, threadFmt []rune) bool {
-				if ml.drawRow(ctx, thread.Uid, row, row, &needsHeaders, string(threadFmt)) {
-					return true
-				}
-				row += 1
-				return false
-			})
+		for i := ml.scroll; i < len(store.FlatThreads); i++ {
+			thread := store.FlatThreads[i]
+			if ml.drawRow(ctx, thread.Uid, row, row, &needsHeaders, string(thread.DrawThread())) {
+				break
+			}
+			row += 1
 		}
 	} else {
 		for i := len(uids) - 1 - ml.scroll; i >= 0; i-- {
diff --git a/worker/imap/open.go b/worker/imap/open.go
index bdb794b..1152887 100644
--- a/worker/imap/open.go
+++ b/worker/imap/open.go
@@ -64,18 +64,23 @@ func (imapw *IMAPWorker) handleDirectoryThreaded(
 			Error:   err,
 		}, nil)
 	} else {
-		aercThreads, count := convertThreads(threads)
+		root := &types.Thread{
+			Uid:   0,
+			Dummy: true,
+		}
+		aercThreads, count := convertThreads(threads, root)
+		root.Children = aercThreads
 		imapw.seqMap = make([]uint32, count)
 		imapw.worker.PostMessage(&types.DirectoryThreaded{
-			Message: types.RespondTo(msg),
-			Threads: aercThreads,
+			Message:    types.RespondTo(msg),
+			ThreadRoot: root,
 		}, nil)
 		imapw.worker.PostMessage(&types.Done{types.RespondTo(msg)}, nil)
 	}
 }
 
 // This sucks... TODO: find a better way to do this.
-func convertThreads(threads []*sortthread.Thread) ([]*types.Thread, int) {
+func convertThreads(threads []*sortthread.Thread, parent *types.Thread) ([]*types.Thread, int) {
 	if threads == nil {
 		return nil, 0
 	}
@@ -84,11 +89,13 @@ func convertThreads(threads []*sortthread.Thread) ([]*types.Thread, int) {
 
 	for i := 0; i < len(threads); i++ {
 		t := threads[i]
-		children, childCount := convertThreads(t.Children)
 		conv[i] = &types.Thread{
-			Uid:      t.Id,
-			Children: children,
+			Uid:   t.Id,
+			Dummy: false,
 		}
+		conv[i].Parent = parent
+		children, childCount := convertThreads(t.Children, conv[i])
+		conv[i].Children = children
 		count += childCount + 1
 	}
 	return conv, count
diff --git a/worker/types/messages.go b/worker/types/messages.go
index 0b9dc6e..d5f2484 100644
--- a/worker/types/messages.go
+++ b/worker/types/messages.go
@@ -159,7 +159,7 @@ type DirectoryContents struct {
 
 type DirectoryThreaded struct {
 	Message
-	Threads []*Thread
+	ThreadRoot *Thread
 }
 
 type SearchResults struct {
diff --git a/worker/types/thread.go b/worker/types/thread.go
index a51a5c4..51cd599 100644
--- a/worker/types/thread.go
+++ b/worker/types/thread.go
@@ -3,44 +3,98 @@ package types
 type Thread struct {
 	Uid      uint32
 	Children []*Thread
+	Parent   *Thread
+	Dummy    bool
 }
 
-func (t *Thread) FormatThread(cb func(*Thread, []rune) bool) {
-	cb(t, []rune{})
-	walkThreads(t.Children, []rune{}, cb)
+func (t *Thread) Find(uid uint32) *Thread {
+	if t.Uid == uid && !t.Dummy {
+		return t
+	}
+
+	for _, child := range t.Children {
+		if needle := child.Find(uid); needle != nil {
+			return needle
+		}
+	}
+
+	return nil
 }
 
-func walkThreads(threads []*Thread, threadFmt []rune,
-	cb func(*Thread, []rune) bool) {
-	if threads == nil {
-		return
-	}
-
-	var (
-		indent          []rune = []rune{'\u0020', '\u0020'}
-		indentConnected []rune = []rune{'\u2502', '\u0020'}
-		arrow           []rune = []rune{'\u251c', '\u2500', '\u003e'}
-		arrowConnected  []rune = []rune{'\u2514', '\u2500', '\u003e'}
-
-		threadPrefix          []rune = append(threadFmt, arrow...)
-		threadPrefixConnected []rune = append(threadFmt, arrowConnected...)
-		nextThread            []rune = append(threadFmt, indent...)
-		nextThreadConnected   []rune = append(threadFmt, indentConnected...)
-	)
-
-	for i := len(threads) - 1; i >= 0; i-- {
-		t := threads[i]
-		if i > 0 && len(threads) > 1 {
-			if cb(t, threadPrefix) {
-				return
-			}
-			walkThreads(t.Children, nextThreadConnected, cb)
+func (t *Thread) Traverse(dummies bool, cb func(*Thread) bool) {
+	if dummies || !t.Dummy {
+		if cb(t) {
+			return
+		}
+	}
+
+	threads := t.Children
+	for i := 0; i < len(threads); i++ {
+		child := threads[i]
+		child.Traverse(dummies, cb)
+	}
+}
+
+var (
+	BINDENT []rune = []rune{'\u0020', '\u0020'}
+	CINDENT []rune = []rune{'\u2502', '\u0020'}
+	LARROW  []rune = []rune{'\u2514', '\u2500', '\u003e'}
+	TARROW  []rune = []rune{'\u252c', '\u2500', '\u003e'}
+	IARROW  []rune = []rune{'\u251c', '\u2500', '\u003e'}
+)
+
+func (t *Thread) isRoot() bool {
+	return t.Dummy && t.Parent == nil
+}
+
+func (t *Thread) isDummyRoot() bool {
+	p := t
+	for p.Dummy {
+		if p.isRoot() {
+			return true
+		}
+		p = p.Parent
+	}
+	return false
+}
+
+func (t *Thread) DrawThread() []rune {
+	var line []rune
+	if t.Parent.isRoot() {
+		return line
+	}
+
+	children := t.Parent.Children
+	if t.Parent.Dummy {
+		children := t.Parent.Children
+		if len(children) > 1 && t == children[0] {
+			line = TARROW
+		} else if len(children) > 1 && t != children[len(children)-1] {
+			line = IARROW
+		} else if len(children) > 1 {
+			line = LARROW
+		}
+	} else {
+		if len(children) > 1 && t != children[len(children)-1] {
+			line = IARROW
 		} else {
-			if cb(t, threadPrefixConnected) {
-				return
+			line = LARROW
+		}
+	}
+
+	p := t.Parent
+
+	for p != nil && !p.Parent.isDummyRoot() {
+		if !p.Dummy {
+			children := p.Parent.Children
+			if len(children) > 1 && p != children[len(children)-1] {
+				line = append(CINDENT, line...)
+			} else {
+				line = append(BINDENT, line...)
 			}
-			walkThreads(t.Children, nextThread, cb)
 		}
+		p = p.Parent
 	}
 
+	return line
 }