threadbuilder: store uid/thread association

Store the association between uids and threads in a map instead of just
having the threads in a slice. This simplifies the lookup of a thread
when we have an uid and we can avoid computationally expensive tree
walks.

The threadbuilder will rebuild the uids from the given thread structure.
Hence there is no need now to keep a threads slice around.

Signed-off-by: Koni Marti <koni.marti@gmail.com>
Acked-by: Robin Jarry <robin@jarry.cc>
Tested-by: inwit <inwit@sindominio.net>

1 files changed, 16 insertions, 0 deletions

diff --git a/lib/threadbuilder.go b/lib/threadbuilder.go
index c2fee228..793034df 100644
--- a/lib/threadbuilder.go
+++ b/lib/threadbuilder.go
@@ -1,6 +1,7 @@
 package lib
 
 import (
+	"fmt"
 	"sync"
 	"time"
 
@@ -15,6 +16,7 @@ type ThreadBuilder struct {
 	sync.Mutex
 	threadBlocks map[uint32]jwz.Threadable
 	threadedUids []uint32
+	threadMap    map[uint32]*types.Thread
 	iterFactory  iterator.Factory
 }
 
@@ -22,10 +24,22 @@ func NewThreadBuilder(i iterator.Factory) *ThreadBuilder {
 	tb := &ThreadBuilder{
 		threadBlocks: make(map[uint32]jwz.Threadable),
 		iterFactory:  i,
+		threadMap:    make(map[uint32]*types.Thread),
 	}
 	return tb
 }
 
+func (builder *ThreadBuilder) ThreadForUid(uid uint32) (*types.Thread, error) {
+	builder.Lock()
+	defer builder.Unlock()
+	t, ok := builder.threadMap[uid]
+	var err error
+	if !ok {
+		err = fmt.Errorf("no thread found for uid '%d'", uid)
+	}
+	return t, err
+}
+
 // Uids returns the uids in threading order
 func (builder *ThreadBuilder) Uids() []uint32 {
 	builder.Lock()
@@ -188,6 +202,7 @@ func (builder *ThreadBuilder) RebuildUids(threads []*types.Thread, inverse bool)
 					return nil
 				}
 				threaduids = append(threaduids, t.Uid)
+				builder.threadMap[t.Uid] = t
 				return nil
 			})
 		if inverse {
@@ -198,6 +213,7 @@ func (builder *ThreadBuilder) RebuildUids(threads []*types.Thread, inverse bool)
 			uids = append(uids, threaduids...)
 		}
 	}
+
 	result := make([]uint32, 0, len(uids))
 	iterU := builder.iterFactory.NewIterator(uids)
 	for iterU.Next() {