threading: honor user-defined sort criteria

Apply the user-defined sort criteria to the message with the highest
uid in a threaded discussion. Restore the default sort order when
leaving threading mode.

Commit 7811620eb809 ("threading: implement on-the-fly message
threading") introduced message threading with the threaded messages
being only sorted by their message uids irrespective of the defined sorting
criteria. It did not restore the default sort order either.

Reported-by: Sebastien Binet <s@sbinet.org>
Signed-off-by: Koni Marti <koni.marti@gmail.com>
Acked-by: Robin Jarry <robin@jarry.cc>

3 files changed, 70 insertions, 50 deletions

diff --git a/lib/msgstore.go b/lib/msgstore.go
index 369f4b4..bd9e935 100644
--- a/lib/msgstore.go
+++ b/lib/msgstore.go
@@ -2,7 +2,6 @@ package lib
 
 import (
 	"io"
-	gosort "sort"
 	"time"
 
 	"git.sr.ht/~rjarry/aerc/lib/sort"
@@ -339,8 +338,6 @@ func (store *MessageStore) SetBuildThreads(buildThreads bool) {
 	store.buildThreads = buildThreads
 	if store.BuildThreads() {
 		store.runThreadBuilder()
-	} else {
-		store.rebuildUids()
 	}
 }
 
@@ -354,46 +351,18 @@ func (store *MessageStore) BuildThreads() bool {
 
 func (store *MessageStore) runThreadBuilder() {
 	if store.builder == nil {
-		store.builder = NewThreadBuilder(store, store.worker.Logger)
+		store.builder = NewThreadBuilder(store.worker.Logger)
 		for _, msg := range store.Messages {
 			store.builder.Update(msg)
 		}
 	}
-	store.Threads = store.builder.Threads()
-	store.rebuildUids()
-}
-
-func (store *MessageStore) rebuildUids() {
-	start := time.Now()
-
-	uids := make([]uint32, 0, len(store.Uids()))
-
-	if store.BuildThreads() {
-		gosort.Sort(types.ByUID(store.Threads))
-		for i := len(store.Threads) - 1; i >= 0; i-- {
-			store.Threads[i].Walk(func(t *types.Thread, level int, currentErr error) error {
-				uids = append(uids, t.Uid)
-				return nil
-			})
-		}
-		uidsReversed := make([]uint32, len(uids))
-		for i := 0; i < len(uids); i++ {
-			uidsReversed[i] = uids[len(uids)-1-i]
-		}
-		uids = uidsReversed
-	} else {
-		uids = store.Uids()
-		gosort.SliceStable(uids, func(i, j int) bool { return uids[i] < uids[j] })
-	}
-
+	var uids []uint32
 	if store.filter {
-		store.results = uids
+		uids = store.results
 	} else {
-		store.uids = uids
+		uids = store.uids
 	}
-
-	elapsed := time.Since(start)
-	store.worker.Logger.Println("Store: Rebuilding UIDs took", elapsed)
+	store.Threads = store.builder.Threads(uids)
 }
 
 func (store *MessageStore) Delete(uids []uint32,
@@ -472,6 +441,13 @@ func (store *MessageStore) Answered(uids []uint32, answered bool,
 }
 
 func (store *MessageStore) Uids() []uint32 {
+
+	if store.BuildThreads() && store.builder != nil {
+		if uids := store.builder.Uids(); len(uids) > 0 {
+			return uids
+		}
+	}
+
 	if store.filter {
 		return store.results
 	}
diff --git a/lib/threadbuilder.go b/lib/threadbuilder.go
index c87d0bf..9607f1c 100644
--- a/lib/threadbuilder.go
+++ b/lib/threadbuilder.go
@@ -9,29 +9,33 @@ import (
 	"github.com/gatherstars-com/jwz"
 )
 
-type UidStorer interface {
-	Uids() []uint32
-}
-
 type ThreadBuilder struct {
 	threadBlocks   map[uint32]jwz.Threadable
 	messageidToUid map[string]uint32
 	seen           map[uint32]bool
-	store          UidStorer
+	threadedUids   []uint32
 	logger         *log.Logger
 }
 
-func NewThreadBuilder(store UidStorer, logger *log.Logger) *ThreadBuilder {
+func NewThreadBuilder(logger *log.Logger) *ThreadBuilder {
 	tb := &ThreadBuilder{
 		threadBlocks:   make(map[uint32]jwz.Threadable),
 		messageidToUid: make(map[string]uint32),
 		seen:           make(map[uint32]bool),
-		store:          store,
 		logger:         logger,
 	}
 	return tb
 }
 
+// Uids returns the uids in threading order
+func (builder *ThreadBuilder) Uids() []uint32 {
+	if builder.threadedUids == nil {
+		return []uint32{}
+	}
+	return builder.threadedUids
+}
+
+// Update updates the thread builder with a new message header
 func (builder *ThreadBuilder) Update(msg *models.MessageInfo) {
 	if msg != nil {
 		if threadable := newThreadable(msg); threadable != nil {
@@ -41,10 +45,17 @@ func (builder *ThreadBuilder) Update(msg *models.MessageInfo) {
 	}
 }
 
-func (builder *ThreadBuilder) Threads() []*types.Thread {
+// Threads returns a slice of threads for the given list of uids
+func (builder *ThreadBuilder) Threads(uids []uint32) []*types.Thread {
 	start := time.Now()
 
-	threads := builder.buildAercThreads(builder.generateStructure())
+	threads := builder.buildAercThreads(builder.generateStructure(uids), uids)
+
+	// sort threads according to uid ordering
+	builder.sortThreads(threads, uids)
+
+	// rebuild uids from threads
+	builder.RebuildUids(threads)
 
 	elapsed := time.Since(start)
 	builder.logger.Println("ThreadBuilder:", len(threads), "threads created in", elapsed)
@@ -52,9 +63,9 @@ func (builder *ThreadBuilder) Threads() []*types.Thread {
 	return threads
 }
 
-func (builder *ThreadBuilder) generateStructure() jwz.Threadable {
+func (builder *ThreadBuilder) generateStructure(uids []uint32) jwz.Threadable {
 	jwzThreads := make([]jwz.Threadable, 0, len(builder.threadBlocks))
-	for _, uid := range builder.store.Uids() {
+	for _, uid := range uids {
 		if thr, ok := builder.threadBlocks[uid]; ok {
 			jwzThreads = append(jwzThreads, thr)
 		}
@@ -68,15 +79,15 @@ func (builder *ThreadBuilder) generateStructure() jwz.Threadable {
 	return threadStructure
 }
 
-func (builder *ThreadBuilder) buildAercThreads(structure jwz.Threadable) []*types.Thread {
+func (builder *ThreadBuilder) buildAercThreads(structure jwz.Threadable, uids []uint32) []*types.Thread {
 	threads := make([]*types.Thread, 0, len(builder.threadBlocks))
 	if structure == nil {
-		for _, uid := range builder.store.Uids() {
+		for _, uid := range uids {
 			threads = append(threads, &types.Thread{Uid: uid})
 		}
 	} else {
 		// fill threads with nil messages
-		for _, uid := range builder.store.Uids() {
+		for _, uid := range uids {
 			if _, ok := builder.threadBlocks[uid]; !ok {
 				threads = append(threads, &types.Thread{Uid: uid})
 			}
@@ -127,6 +138,26 @@ func (builder *ThreadBuilder) buildTree(treeNode jwz.Threadable, target *types.T
 	}
 }
 
+func (builder *ThreadBuilder) sortThreads(threads []*types.Thread, orderedUids []uint32) {
+	types.SortThreadsBy(threads, orderedUids)
+}
+
+// RebuildUids rebuilds the uids from the given slice of threads
+func (builder *ThreadBuilder) RebuildUids(threads []*types.Thread) {
+	uids := make([]uint32, 0, len(threads))
+	for i := len(threads) - 1; i >= 0; i-- {
+		threads[i].Walk(func(t *types.Thread, level int, currentErr error) error {
+			uids = append(uids, t.Uid)
+			return nil
+		})
+	}
+	// copy in reverse as msgList displays backwards
+	for i, j := 0, len(uids)-1; i < j; i, j = i+1, j-1 {
+		uids[i], uids[j] = uids[j], uids[i]
+	}
+	builder.threadedUids = uids
+}
+
 // threadable implements the jwz.threadable interface which is required for the
 // jwz threading algorithm
 type threadable struct {
diff --git a/worker/types/thread.go b/worker/types/thread.go
index 48e4a00..7c0cc5b 100644
--- a/worker/types/thread.go
+++ b/worker/types/thread.go
@@ -3,6 +3,7 @@ package types
 import (
 	"errors"
 	"fmt"
+	"sort"
 )
 
 type Thread struct {
@@ -120,3 +121,15 @@ func (s ByUID) Less(i, j int) bool {
 	maxUID_j := getMaxUID(s[j])
 	return maxUID_i < maxUID_j
 }
+
+func SortThreadsBy(toSort []*Thread, sortBy []uint32) {
+	// build a map from sortBy
+	uidMap := make(map[uint32]int)
+	for i, uid := range sortBy {
+		uidMap[uid] = i
+	}
+	// sortslice of toSort with less function of indexing the map sortBy
+	sort.Slice(toSort, func(i, j int) bool {
+		return uidMap[getMaxUID(toSort[i])] < uidMap[getMaxUID(toSort[j])]
+	})
+}