diff --git a/utils/iterator/filter.go b/utils/iterator/filter.go index e8a11464457..f26b082aeab 100644 --- a/utils/iterator/filter.go +++ b/utils/iterator/filter.go @@ -3,6 +3,8 @@ package iterator +import "github.com/ava-labs/avalanchego/utils/set" + var _ Iterator[any] = (*filtered[any])(nil) type filtered[T any] struct { @@ -19,6 +21,19 @@ func Filter[T any](it Iterator[T], filter func(T) bool) Iterator[T] { } } +// Deduplicate returns an iterator that skips the elements that have already +// been returned from [it]. +func Deduplicate[T comparable](it Iterator[T]) Iterator[T] { + var seen set.Set[T] + return Filter(it, func(e T) bool { + if seen.Contains(e) { + return true + } + seen.Add(e) + return false + }) +} + func (i *filtered[_]) Next() bool { for i.it.Next() { element := i.it.Value() diff --git a/utils/iterator/filter_test.go b/utils/iterator/filter_test.go index bf523017fe9..56c47892095 100644 --- a/utils/iterator/filter_test.go +++ b/utils/iterator/filter_test.go @@ -10,8 +10,9 @@ import ( "github.com/stretchr/testify/require" "github.com/ava-labs/avalanchego/ids" - "github.com/ava-labs/avalanchego/utils/iterator" "github.com/ava-labs/avalanchego/vms/platformvm/state" + + . "github.com/ava-labs/avalanchego/utils/iterator" ) func TestFilter(t *testing.T) { @@ -40,8 +41,8 @@ func TestFilter(t *testing.T) { stakers[3].TxID: stakers[3], } - it := iterator.Filter( - iterator.FromSlice(stakers[:3]...), + it := Filter( + FromSlice(stakers[:3]...), func(staker *state.Staker) bool { _, ok := maskedStakers[staker.TxID] return ok @@ -55,3 +56,11 @@ func TestFilter(t *testing.T) { it.Release() require.False(it.Next()) } + +func TestDeduplicate(t *testing.T) { + require.Equal( + t, + []int{0, 1, 2, 3}, + ToSlice(Deduplicate(FromSlice(0, 1, 2, 1, 2, 0, 3))), + ) +}