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 7811620eb8 ("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>

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
 	}

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 {

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])]
+	})
+}