diff --git a/swarm/shed/index.go b/swarm/shed/index.go index fe799f7e60..c592dc570e 100644 --- a/swarm/shed/index.go +++ b/swarm/shed/index.go @@ -17,6 +17,8 @@ package shed import ( + "bytes" + "github.com/syndtr/goleveldb/leveldb" ) @@ -263,6 +265,76 @@ func (f Index) IterateFrom(start Item, fn IndexIterFunc) (err error) { return it.Error() } +// IterateWithPrefix iterates over all keys of the Index that have +// a common prefix. +func (f Index) IterateWithPrefix(prefix []byte, fn IndexIterFunc) (err error) { + it := f.db.NewIterator() + defer it.Release() + + // construct complete prefix with index prefix + p := append(f.prefix, prefix...) + + for ok := it.Seek(p); ok; ok = it.Next() { + key := it.Key() + if !bytes.HasPrefix(key, p) { + break + } + keyItem, err := f.decodeKeyFunc(key) + if err != nil { + return err + } + valueItem, err := f.decodeValueFunc(keyItem, it.Value()) + if err != nil { + return err + } + stop, err := fn(keyItem.Merge(valueItem)) + if err != nil { + return err + } + if stop { + break + } + } + return it.Error() +} + +// IterateWithPrefixFrom iterates over Index keys that have a common prefix, +// starting from the key encoded from the provided start Item. +func (f Index) IterateWithPrefixFrom(prefix []byte, start Item, fn IndexIterFunc) (err error) { + startKey, err := f.encodeKeyFunc(start) + if err != nil { + return err + } + it := f.db.NewIterator() + defer it.Release() + + // construct complete prefix with index prefix + p := append(f.prefix, prefix...) + + for ok := it.Seek(startKey); ok; ok = it.Next() { + key := it.Key() + if !bytes.HasPrefix(key, p) { + break + } + keyItem, err := f.decodeKeyFunc(key) + if err != nil { + return err + } + valueItem, err := f.decodeValueFunc(keyItem, it.Value()) + if err != nil { + return err + } + stop, err := fn(keyItem.Merge(valueItem)) + if err != nil { + return err + } + if stop { + break + } + } + return it.Error() +} + // Count returns the number of items in index. func (f Index) Count() (count int, err error) { it := f.db.NewIterator() diff --git a/swarm/shed/index_test.go b/swarm/shed/index_test.go index 39005129fb..87f4e4e999 100644 --- a/swarm/shed/index_test.go +++ b/swarm/shed/index_test.go @@ -250,7 +250,8 @@ func TestIndex(t *testing.T) { }) } -// TestIndex_iterate validates index iterator functions for correctness. +// TestIndex_iterate validates index IterateAll and IterateFrom +// functions for correctness. func TestIndex_iterate(t *testing.T) { db, cleanupFunc := newTestDB(t) defer cleanupFunc() @@ -407,6 +408,153 @@ func TestIndex_iterate(t *testing.T) { }) } +// TestIndex_iterateWithPrefix validates index IterateWithPrefix +// and IterateWithPrefixFrom functions for correctness. +func TestIndex_iterateWithPrefix(t *testing.T) { + db, cleanupFunc := newTestDB(t) + defer cleanupFunc() + + index, err := db.NewIndex("retrieval", retrievalIndexFuncs) + if err != nil { + t.Fatal(err) + } + + allItems := []Item{ + {Address: []byte("want-hash-00"), Data: []byte("data80")}, + {Address: []byte("skip-hash-01"), Data: []byte("data81")}, + {Address: []byte("skip-hash-02"), Data: []byte("data82")}, + {Address: []byte("skip-hash-03"), Data: []byte("data83")}, + {Address: []byte("want-hash-04"), Data: []byte("data84")}, + {Address: []byte("want-hash-05"), Data: []byte("data85")}, + {Address: []byte("want-hash-06"), Data: []byte("data86")}, + {Address: []byte("want-hash-07"), Data: []byte("data87")}, + {Address: []byte("want-hash-08"), Data: []byte("data88")}, + {Address: []byte("want-hash-09"), Data: []byte("data89")}, + {Address: []byte("skip-hash-10"), Data: []byte("data90")}, + } + batch := new(leveldb.Batch) + for _, i := range allItems { + index.PutInBatch(batch, i) + } + err = db.WriteBatch(batch) + if err != nil { + t.Fatal(err) + } + + prefix := []byte("want") + + items := make([]Item, 0) + for _, item := range allItems { + if bytes.HasPrefix(item.Address, prefix) { + items = append(items, item) + } + } + sort.SliceStable(items, func(i, j int) bool { + return bytes.Compare(items[i].Address, items[j].Address) < 0 + }) + + t.Run("with prefix", func(t *testing.T) { + var i int + err := index.IterateWithPrefix(prefix, func(item Item) (stop bool, err error) { + if i > len(items)-1 { + return true, fmt.Errorf("got unexpected index item: %#v", item) + } + want := items[i] + checkItem(t, item, want) + i++ + return false, nil + }) + if err != nil { + t.Fatal(err) + } + if i != len(items) { + t.Errorf("got %v items, want %v", i, len(items)) + } + }) + + t.Run("with prefix from", func(t *testing.T) { + startIndex := 2 + var count int + i := startIndex + err := index.IterateWithPrefixFrom(prefix, items[startIndex], func(item Item) (stop bool, err error) { + if i > len(items)-1 { + return true, fmt.Errorf("got unexpected index item: %#v", item) + } + want := items[i] + checkItem(t, item, want) + i++ + count++ + return false, nil + }) + if err != nil { + t.Fatal(err) + } + wantCount := len(items) - startIndex + if count != wantCount { + t.Errorf("got %v items, want %v", count, wantCount) + } + }) + + t.Run("stop", func(t *testing.T) { + var i int + stopIndex := 3 + var count int + err := index.IterateWithPrefix(prefix, func(item Item) (stop bool, err error) { + if i > len(items)-1 { + return true, fmt.Errorf("got unexpected index item: %#v", item) + } + want := items[i] + checkItem(t, item, want) + count++ + if i == stopIndex { + return true, nil + } + i++ + return false, nil + }) + if err != nil { + t.Fatal(err) + } + wantItemsCount := stopIndex + 1 + if count != wantItemsCount { + t.Errorf("got %v items, expected %v", count, wantItemsCount) + } + }) + + t.Run("no overflow", func(t *testing.T) { + secondIndex, err := db.NewIndex("second-index", retrievalIndexFuncs) + if err != nil { + t.Fatal(err) + } + + secondItem := Item{ + Address: []byte("iterate-hash-10"), + Data: []byte("data-second"), + } + err = secondIndex.Put(secondItem) + if err != nil { + t.Fatal(err) + } + + var i int + err = index.IterateWithPrefix(prefix, func(item Item) (stop bool, err error) { + if i > len(items)-1 { + return true, fmt.Errorf("got unexpected index item: %#v", item) + } + want := items[i] + checkItem(t, item, want) + i++ + return false, nil + }) + if err != nil { + t.Fatal(err) + } + if i != len(items) { + t.Errorf("got %v items, want %v", i, len(items)) + } + }) +} + // TestIndex_count tests if Index.Count and Index.CountFrom // returns the correct number of items. func TestIndex_count(t *testing.T) {