Change Merger implementation on --no-sort

2 files changed, 29 insertions, 27 deletions

diff --git a/src/merger.go b/src/merger.go
index 7ff32434..08a3d154 100644
--- a/src/merger.go
+++ b/src/merger.go
@@ -1,12 +1,15 @@
 package fzf
 
+import "fmt"
+
 var EmptyMerger *Merger = NewMerger([][]*Item{}, false)
 
 type Merger struct {
 	lists   [][]*Item
 	merged  []*Item
 	cursors []int
-	done    bool
+	sorted  bool
+	count   int
 }
 
 func NewMerger(lists [][]*Item, sorted bool) *Merger {
@@ -14,41 +17,37 @@ func NewMerger(lists [][]*Item, sorted bool) *Merger {
 		lists:   lists,
 		merged:  []*Item{},
 		cursors: make([]int, len(lists)),
-		done:    false}
-	if !sorted {
-		for _, list := range lists {
-			mg.merged = append(mg.merged, list...)
-		}
-		mg.done = true
+		sorted:  sorted,
+		count:   0}
+
+	for _, list := range mg.lists {
+		mg.count += len(list)
 	}
 	return &mg
 }
 
 func (mg *Merger) Length() int {
-	cnt := 0
-	for _, list := range mg.lists {
-		cnt += len(list)
-	}
-	return cnt
+	return mg.count
 }
 
 func (mg *Merger) Get(idx int) *Item {
-	if mg.done {
-		return mg.merged[idx]
-	} else if len(mg.lists) == 1 {
+	if len(mg.lists) == 1 {
 		return mg.lists[0][idx]
+	} else if !mg.sorted {
+		for _, list := range mg.lists {
+			numItems := len(list)
+			if idx < numItems {
+				return list[idx]
+			}
+			idx -= numItems
+		}
+		panic(fmt.Sprintf("Index out of bounds (unsorted, %d/%d)", idx, mg.count))
 	}
-	mg.buildUpto(idx)
-	return mg.merged[idx]
+	return mg.mergedGet(idx)
 }
 
-func (mg *Merger) buildUpto(upto int) {
-	numBuilt := len(mg.merged)
-	if numBuilt > upto {
-		return
-	}
-
-	for i := numBuilt; i <= upto; i++ {
+func (mg *Merger) mergedGet(idx int) *Item {
+	for i := len(mg.merged); i <= idx; i++ {
 		minRank := Rank{0, 0, 0}
 		minIdx := -1
 		for listIdx, list := range mg.lists {
@@ -72,8 +71,8 @@ func (mg *Merger) buildUpto(upto int) {
 			mg.merged = append(mg.merged, chosen[mg.cursors[minIdx]])
 			mg.cursors[minIdx] += 1
 		} else {
-			mg.done = true
-			return
+			panic(fmt.Sprintf("Index out of bounds (sorted, %d/%d)", i, mg.count))
 		}
 	}
+	return mg.merged[idx]
 }
diff --git a/src/merger_test.go b/src/merger_test.go
index 19941b17..32a1228a 100644
--- a/src/merger_test.go
+++ b/src/merger_test.go
@@ -19,6 +19,9 @@ func randItem() *Item {
 
 func TestEmptyMerger(t *testing.T) {
 	assert(t, EmptyMerger.Length() == 0, "Not empty")
+	assert(t, EmptyMerger.count == 0, "Invalid count")
+	assert(t, len(EmptyMerger.lists) == 0, "Invalid lists")
+	assert(t, len(EmptyMerger.merged) == 0, "Invalid merged list")
 }
 
 func buildLists(partiallySorted bool) ([][]*Item, []*Item) {
@@ -72,7 +75,7 @@ func TestMergerSorted(t *testing.T) {
 
 	// Inverse order
 	mg2 := NewMerger(lists, true)
-	for i := cnt - 1; i >= cnt; i-- {
+	for i := cnt - 1; i >= 0; i-- {
 		if items[i] != mg2.Get(i) {
 			t.Error("Not sorted", items[i], mg2.Get(i))
 		}