From 55079aa893e91f4d15dab6342a4d02fb1aef0dd3 Mon Sep 17 00:00:00 2001
From: Dan Laine <daniel.laine@avalabs.org>
Date: Tue, 23 Jun 2020 15:01:55 -0400
Subject: [PATCH] add CappedList for ids.Set and use it in fetchANeededVtx

---
 ids/set.go                            | 23 +++++++++++++-
 ids/set_test.go                       | 43 +++++++++++++++++++++++++++
 snow/engine/avalanche/bootstrapper.go |  2 +-
 3 files changed, 66 insertions(+), 2 deletions(-)

diff --git a/ids/set.go b/ids/set.go
index c3aa024..d632949 100644
--- a/ids/set.go
+++ b/ids/set.go
@@ -78,7 +78,7 @@ func (ids *Set) Clear() { *ids = nil }
 
 // List converts this set into a list
 func (ids Set) List() []ID {
-	idList := make([]ID, ids.Len(), ids.Len())
+	idList := make([]ID, ids.Len())
 	i := 0
 	for id := range ids {
 		idList[i] = NewID(id)
@@ -87,6 +87,27 @@ func (ids Set) List() []ID {
 	return idList
 }
 
+// CappedList returns a list of length at most [size].
+// Size should be >= 0. If size < 0, returns nil.
+func (ids Set) CappedList(size int) []ID {
+	if size < 0 {
+		return nil
+	}
+	if l := ids.Len(); l < size {
+		size = l
+	}
+	i := 0
+	idList := make([]ID, size)
+	for id := range ids {
+		if i >= size {
+			break
+		}
+		idList[i] = NewID(id)
+		i++
+	}
+	return idList
+}
+
 // Equals returns true if the sets contain the same elements
 func (ids Set) Equals(oIDs Set) bool {
 	if ids.Len() != oIDs.Len() {
diff --git a/ids/set_test.go b/ids/set_test.go
index 3c7ab15..b4e05db 100644
--- a/ids/set_test.go
+++ b/ids/set_test.go
@@ -55,3 +55,46 @@ func TestSet(t *testing.T) {
 		t.Fatalf("Sets overlap")
 	}
 }
+
+func TestSetCappedList(t *testing.T) {
+	set := Set{}
+
+	id := Empty
+
+	if list := set.CappedList(0); len(list) != 0 {
+		t.Fatalf("List should have been empty but was %v", list)
+	}
+
+	set.Add(id)
+
+	if list := set.CappedList(0); len(list) != 0 {
+		t.Fatalf("List should have been empty but was %v", list)
+	} else if list := set.CappedList(1); len(list) != 1 {
+		t.Fatalf("List should have had length %d but had %d", 1, len(list))
+	} else if returnedID := list[0]; !id.Equals(returnedID) {
+		t.Fatalf("List should have been %s but was %s", id, returnedID)
+	} else if list := set.CappedList(2); len(list) != 1 {
+		t.Fatalf("List should have had length %d but had %d", 1, len(list))
+	} else if returnedID := list[0]; !id.Equals(returnedID) {
+		t.Fatalf("List should have been %s but was %s", id, returnedID)
+	}
+
+	id2 := NewID([32]byte{1})
+	set.Add(id2)
+
+	if list := set.CappedList(0); len(list) != 0 {
+		t.Fatalf("List should have been empty but was %v", list)
+	} else if list := set.CappedList(1); len(list) != 1 {
+		t.Fatalf("List should have had length %d but had %d", 1, len(list))
+	} else if returnedID := list[0]; !id.Equals(returnedID) && !id2.Equals(returnedID) {
+		t.Fatalf("List should have been %s but was %s", id, returnedID)
+	} else if list := set.CappedList(2); len(list) != 2 {
+		t.Fatalf("List should have had length %d but had %d", 2, len(list))
+	} else if list := set.CappedList(3); len(list) != 2 {
+		t.Fatalf("List should have had length %d but had %d", 2, len(list))
+	} else if returnedID := list[0]; !id.Equals(returnedID) && !id2.Equals(returnedID) {
+		t.Fatalf("list contains unexpected element %s", returnedID)
+	} else if returnedID := list[1]; !id.Equals(returnedID) && !id2.Equals(returnedID) {
+		t.Fatalf("list contains unexpected element %s", returnedID)
+	}
+}
diff --git a/snow/engine/avalanche/bootstrapper.go b/snow/engine/avalanche/bootstrapper.go
index 749694e..352f40b 100644
--- a/snow/engine/avalanche/bootstrapper.go
+++ b/snow/engine/avalanche/bootstrapper.go
@@ -111,7 +111,7 @@ func (b *bootstrapper) FilterAccepted(containerIDs ids.Set) ids.Set {
 // Calls fetch for a pending vertex if there are any
 func (b *bootstrapper) fetchANeededVtx() error {
 	if b.needToFetch.Len() > 0 {
-		return b.fetch(b.needToFetch.List()[0])
+		return b.fetch(b.needToFetch.CappedList(1)[0])
 	}
 	return nil
 }