Fix race conditions

- Wait for completions of goroutines when cancelling a search
- Remove shared access to rank field of Item

7 files changed, 49 insertions, 23 deletions

diff --git a/src/core.go b/src/core.go
index 2970038f..ab2a48fa 100644
--- a/src/core.go
+++ b/src/core.go
@@ -64,7 +64,10 @@ func Run(options *Options) {
 	var chunkList *ChunkList
 	if len(opts.WithNth) == 0 {
 		chunkList = NewChunkList(func(data *string, index int) *Item {
-			return &Item{text: data, rank: Rank{0, 0, uint32(index)}}
+			return &Item{
+				text:  data,
+				index: uint32(index),
+				rank:  Rank{0, 0, uint32(index)}}
 		})
 	} else {
 		chunkList = NewChunkList(func(data *string, index int) *Item {
@@ -72,6 +75,7 @@ func Run(options *Options) {
 			item := Item{
 				text:     Transform(tokens, opts.WithNth).whole,
 				origText: data,
+				index:    uint32(index),
 				rank:     Rank{0, 0, uint32(index)}}
 			return &item
 		})
diff --git a/src/item.go b/src/item.go
index 9f90b8d4..41aa34bd 100644
--- a/src/item.go
+++ b/src/item.go
@@ -1,9 +1,6 @@
 package fzf
 
-import (
-	"fmt"
-	"sort"
-)
+import "fmt"
 
 type Offset [2]int32
 
@@ -11,6 +8,7 @@ type Item struct {
 	text        *string
 	origText    *string
 	transformed *Transformed
+	index       uint32
 	offsets     []Offset
 	rank        Rank
 }
@@ -21,11 +19,10 @@ type Rank struct {
 	index    uint32
 }
 
-func (i *Item) Rank() Rank {
-	if i.rank.matchlen > 0 || i.rank.strlen > 0 {
+func (i *Item) Rank(cache bool) Rank {
+	if cache && (i.rank.matchlen > 0 || i.rank.strlen > 0) {
 		return i.rank
 	}
-	sort.Sort(ByOrder(i.offsets))
 	matchlen := 0
 	prevEnd := 0
 	for _, offset := range i.offsets {
@@ -41,8 +38,11 @@ func (i *Item) Rank() Rank {
 			matchlen += end - begin
 		}
 	}
-	i.rank = Rank{uint16(matchlen), uint16(len(*i.text)), i.rank.index}
-	return i.rank
+	rank := Rank{uint16(matchlen), uint16(len(*i.text)), i.index}
+	if cache {
+		i.rank = rank
+	}
+	return rank
 }
 
 func (i *Item) Print() {
@@ -80,8 +80,8 @@ func (a ByRelevance) Swap(i, j int) {
 }
 
 func (a ByRelevance) Less(i, j int) bool {
-	irank := a[i].Rank()
-	jrank := a[j].Rank()
+	irank := a[i].Rank(true)
+	jrank := a[j].Rank(true)
 
 	return compareRanks(irank, jrank)
 }
diff --git a/src/item_test.go b/src/item_test.go
index 87d8be46..0e83631a 100644
--- a/src/item_test.go
+++ b/src/item_test.go
@@ -31,13 +31,13 @@ func TestRankComparison(t *testing.T) {
 // Match length, string length, index
 func TestItemRank(t *testing.T) {
 	strs := []string{"foo", "foobar", "bar", "baz"}
-	item1 := Item{text: &strs[0], rank: Rank{0, 0, 1}, offsets: []Offset{}}
-	rank1 := item1.Rank()
+	item1 := Item{text: &strs[0], index: 1, offsets: []Offset{}}
+	rank1 := item1.Rank(true)
 	if rank1.matchlen != 0 || rank1.strlen != 3 || rank1.index != 1 {
-		t.Error(item1.Rank())
+		t.Error(item1.Rank(true))
 	}
 	// Only differ in index
-	item2 := Item{text: &strs[0], rank: Rank{0, 0, 0}, offsets: []Offset{}}
+	item2 := Item{text: &strs[0], index: 0, offsets: []Offset{}}
 
 	items := []*Item{&item1, &item2}
 	sort.Sort(ByRelevance(items))
diff --git a/src/matcher.go b/src/matcher.go
index 234e7033..713b4dd6 100644
--- a/src/matcher.go
+++ b/src/matcher.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"runtime"
 	"sort"
+	"sync"
 	"time"
 )
 
@@ -134,10 +135,13 @@ func (m *Matcher) scan(request MatchRequest, limit int) (*Merger, bool) {
 	slices := m.sliceChunks(request.chunks)
 	numSlices := len(slices)
 	resultChan := make(chan partialResult, numSlices)
-	countChan := make(chan int, numSlices)
+	countChan := make(chan int, numChunks)
+	waitGroup := sync.WaitGroup{}
 
 	for idx, chunks := range slices {
+		waitGroup.Add(1)
 		go func(idx int, chunks []*Chunk) {
+			defer func() { waitGroup.Done() }()
 			sliceMatches := []*Item{}
 			for _, chunk := range chunks {
 				var matches []*Item
@@ -159,6 +163,12 @@ func (m *Matcher) scan(request MatchRequest, limit int) (*Merger, bool) {
 		}(idx, chunks)
 	}
 
+	wait := func() bool {
+		cancelled.Set(true)
+		waitGroup.Wait()
+		return true
+	}
+
 	count := 0
 	matchCount := 0
 	for matchesInChunk := range countChan {
@@ -166,7 +176,7 @@ func (m *Matcher) scan(request MatchRequest, limit int) (*Merger, bool) {
 		matchCount += matchesInChunk
 
 		if limit > 0 && matchCount > limit {
-			return nil, true // For --select-1 and --exit-0
+			return nil, wait() // For --select-1 and --exit-0
 		}
 
 		if count == numChunks {
@@ -174,8 +184,7 @@ func (m *Matcher) scan(request MatchRequest, limit int) (*Merger, bool) {
 		}
 
 		if !empty && m.reqBox.Peak(REQ_RESET) {
-			cancelled.Set(true)
-			return nil, true
+			return nil, wait()
 		}
 
 		if time.Now().Sub(startedAt) > PROGRESS_MIN_DURATION {
diff --git a/src/merger.go b/src/merger.go
index 08a3d154..16afdafc 100644
--- a/src/merger.go
+++ b/src/merger.go
@@ -57,7 +57,7 @@ func (mg *Merger) mergedGet(idx int) *Item {
 				continue
 			}
 			if cursor >= 0 {
-				rank := list[cursor].Rank()
+				rank := list[cursor].Rank(false)
 				if minIdx < 0 || compareRanks(rank, minRank) {
 					minRank = rank
 					minIdx = listIdx
diff --git a/src/merger_test.go b/src/merger_test.go
index 32a1228a..f79da09a 100644
--- a/src/merger_test.go
+++ b/src/merger_test.go
@@ -1,6 +1,7 @@
 package fzf
 
 import (
+	"fmt"
 	"math/rand"
 	"sort"
 	"testing"
@@ -13,8 +14,17 @@ func assert(t *testing.T, cond bool, msg ...string) {
 }
 
 func randItem() *Item {
+	str := fmt.Sprintf("%d", rand.Uint32())
+	offsets := make([]Offset, rand.Int()%3)
+	for idx := range offsets {
+		sidx := int32(rand.Uint32() % 20)
+		eidx := sidx + int32(rand.Uint32()%20)
+		offsets[idx] = Offset{sidx, eidx}
+	}
 	return &Item{
-		rank: Rank{uint16(rand.Uint32()), uint16(rand.Uint32()), rand.Uint32()}}
+		text:    &str,
+		index:   rand.Uint32(),
+		offsets: offsets}
 }
 
 func TestEmptyMerger(t *testing.T) {
diff --git a/src/pattern.go b/src/pattern.go
index 31ba8137..2e7d6f91 100644
--- a/src/pattern.go
+++ b/src/pattern.go
@@ -2,6 +2,7 @@ package fzf
 
 import (
 	"regexp"
+	"sort"
 	"strings"
 )
 
@@ -225,12 +226,14 @@ Loop:
 }
 
 func dupItem(item *Item, offsets []Offset) *Item {
+	sort.Sort(ByOrder(offsets))
 	return &Item{
 		text:        item.text,
 		origText:    item.origText,
 		transformed: item.transformed,
+		index:       item.index,
 		offsets:     offsets,
-		rank:        Rank{0, 0, item.rank.index}}
+		rank:        Rank{0, 0, item.index}}
 }
 
 func (p *Pattern) fuzzyMatch(chunk *Chunk) []*Item {