swarm/shed: add Index.CountFrom

swarm/shed/index.go

@@ -277,3 +277,23 @@ func (f Index) Count() (count int, err error) {
 	}
 	return count, it.Error()
 }
+
+// CountFrom returns the number of items in index keys
+// starting from the key encoded from the provided Item.
+func (f Index) CountFrom(start Item) (count int, err error) {
+	startKey, err := f.encodeKeyFunc(start)
+	if err != nil {
+		return 0, err
+	}
+	it := f.db.NewIterator()
+	defer it.Release()
+
+	for ok := it.Seek(startKey); ok; ok = it.Next() {
+		key := it.Key()
+		if key[0] != f.prefix[0] {
+			break
+		}
+		count++
+	}
+	return count, it.Error()
+}

swarm/shed/index_test.go

@@ -407,9 +407,9 @@ func TestIndex_iterate(t *testing.T) {
 	})
 }
 
-// TestIndex_Count tests if Index.Count returns the correct
+// TestIndex_count tests if Index.Count and Index.CountFrom
-// number of items.
+// returns the correct number of items.
-func TestIndex_Count(t *testing.T) {
+func TestIndex_count(t *testing.T) {
 	db, cleanupFunc := newTestDB(t)
 	defer cleanupFunc()
 
@@ -423,20 +423,20 @@ func TestIndex_Count(t *testing.T) {
 			Address: []byte("iterate-hash-01"),
 			Data:    []byte("data80"),
 		},
-		{
-			Address: []byte("iterate-hash-03"),
-			Data:    []byte("data22"),
-		},
-		{
-			Address: []byte("iterate-hash-05"),
-			Data:    []byte("data41"),
-		},
 		{
 			Address: []byte("iterate-hash-02"),
 			Data:    []byte("data84"),
 		},
 		{
-			Address: []byte("iterate-hash-06"),
+			Address: []byte("iterate-hash-03"),
+			Data:    []byte("data22"),
+		},
+		{
+			Address: []byte("iterate-hash-04"),
+			Data:    []byte("data41"),
+		},
+		{
+			Address: []byte("iterate-hash-05"),
 			Data:    []byte("data1"),
 		},
 	}
@@ -449,6 +449,7 @@ func TestIndex_Count(t *testing.T) {
 		t.Fatal(err)
 	}
 
+	t.Run("Count", func(t *testing.T) {
 		got, err := index.Count()
 		if err != nil {
 			t.Fatal(err)
@@ -458,11 +459,26 @@ func TestIndex_Count(t *testing.T) {
 		if got != want {
 			t.Errorf("got %v items count, want %v", got, want)
 		}
+	})
+
+	t.Run("CountFrom", func(t *testing.T) {
+		got, err := index.CountFrom(Item{
+			Address: items[1].Address,
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		want := len(items) - 1
+		if got != want {
+			t.Errorf("got %v items count, want %v", got, want)
+		}
+	})
 
 	// update the index with another item
-
+	t.Run("add item", func(t *testing.T) {
 		item04 := Item{
-		Address: []byte("iterate-hash-04"),
+			Address: []byte("iterate-hash-06"),
 			Data:    []byte("data0"),
 		}
 		err = index.Put(item04)
@@ -470,18 +486,37 @@ func TestIndex_Count(t *testing.T) {
 			t.Fatal(err)
 		}
 
-	got, err = index.Count()
+		count := len(items) + 1
+
+		t.Run("Count", func(t *testing.T) {
+			got, err := index.Count()
 			if err != nil {
 				t.Fatal(err)
 			}
 
-	want = len(items) + 1
+			want := count
 			if got != want {
 				t.Errorf("got %v items count, want %v", got, want)
 			}
+		})
+
+		t.Run("CountFrom", func(t *testing.T) {
+			got, err := index.CountFrom(Item{
+				Address: items[1].Address,
+			})
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			want := count - 1
+			if got != want {
+				t.Errorf("got %v items count, want %v", got, want)
+			}
+		})
+	})
 
 	// delete some items
-
+	t.Run("delete items", func(t *testing.T) {
 		deleteCount := 3
 
 		for _, item := range items[:deleteCount] {
@@ -491,15 +526,34 @@ func TestIndex_Count(t *testing.T) {
 			}
 		}
 
-	got, err = index.Count()
+		count := len(items) + 1 - deleteCount
+
+		t.Run("Count", func(t *testing.T) {
+			got, err := index.Count()
 			if err != nil {
 				t.Fatal(err)
 			}
 
-	want = len(items) + 1 - deleteCount
+			want := count
 			if got != want {
 				t.Errorf("got %v items count, want %v", got, want)
 			}
+		})
+
+		t.Run("CountFrom", func(t *testing.T) {
+			got, err := index.CountFrom(Item{
+				Address: items[deleteCount+1].Address,
+			})
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			want := count - 1
+			if got != want {
+				t.Errorf("got %v items count, want %v", got, want)
+			}
+		})
+	})
 }
 
 // checkItem is a test helper function that compares if two Index items are the same.