From b42273490467fc5dd7a9a2a546f7d71eccecfde2 Mon Sep 17 00:00:00 2001
From: andeya <andeyalee@outlook.com>
Date: Tue, 15 Nov 2022 14:57:15 +0800
Subject: [PATCH] feat(iter): add ToUniqueBy and ToDeUniqueBy

Change-Id: Icb9e63d3ee87fee190ea2642de49e1bee67ba990
---
 iter/export.go      | 57 +++++++++++++++++++++++++++++++++++++++++++++
 iter/export_test.go | 15 ++++++++++++
 2 files changed, 72 insertions(+)

diff --git a/iter/export.go b/iter/export.go
index ec29f36..20fd9cd 100644
--- a/iter/export.go
+++ b/iter/export.go
@@ -185,6 +185,63 @@ func ToDeUnique[T comparable](iter DeIterator[T]) DeIterator[T] {
 	})
 }
 
+// ToUniqueBy return an iterator adaptor that filters out elements that have
+// already been produced once during the iteration.
+//
+// Duplicates are detected by comparing the key they map to
+// with the keying function `f` by hash and equality.
+// The keys are stored in a hash set in the iterator.
+//
+// The iterator is stable, returning the non-duplicate items in the order
+// in which they occur in the adapted iterator. In a set of duplicate
+// items, the first item encountered is the item retained.
+//
+// ```
+// var data = FromElements("a", "bb", "aa", "c", "ccc");
+// ToUniqueBy(data, func(s string)int {return len(s)}).Collect() // "a", "bb", "ccc"
+// ```
+func ToUniqueBy[T any, K comparable](iter Iterator[T], f func(T) K) Iterator[T] {
+	min, _ := iter.SizeHint()
+	var set = make(map[K]struct{}, min)
+	return iter.ToFilter(func(x T) bool {
+		k := f(x)
+		if _, ok := set[k]; ok {
+			return false
+		} else {
+			set[k] = struct{}{}
+			return true
+		}
+	})
+}
+
+// ToDeUniqueBy return an iterator adaptor that filters out elements that have
+// already been produced once during the iteration.
+//
+// Duplicates are detected by comparing the key they map to
+// with the keying function `f` by hash and equality.
+// The keys are stored in a hash set in the iterator.
+//
+// The iterator is stable, returning the non-duplicate items in the order
+// in which they occur in the adapted iterator. In a set of duplicate
+// items, the first item encountered is the item retained.
+//
+// ```
+// var data = FromElements("a", "bb", "aa", "c", "ccc");
+// ToDeUniqueBy(data, func(s string)int {return len(s)}).Collect() // "a", "bb", "ccc"
+// ```
+func ToDeUniqueBy[T any, K comparable](iter DeIterator[T], f func(T) K) DeIterator[T] {
+	var set = make(map[K]struct{}, iter.Remaining())
+	return iter.ToDeFilter(func(x T) bool {
+		k := f(x)
+		if _, ok := set[k]; ok {
+			return false
+		} else {
+			set[k] = struct{}{}
+			return true
+		}
+	})
+}
+
 // ToMap takes a closure and creates an iterator which calls that closure on each
 // element.
 //
diff --git a/iter/export_test.go b/iter/export_test.go
index d42c09d..6a2f417 100644
--- a/iter/export_test.go
+++ b/iter/export_test.go
@@ -46,3 +46,18 @@ func TestToDeUnique(t *testing.T) {
 	var data3 = FromElements(10, 20, 30, 20, 40, 10, 50)
 	assert.Equal(t, []int{50, 10, 40, 20, 30}, ToDeUnique(data3.ToRev()).Collect())
 }
+
+func TestToUniqueBy(t *testing.T) {
+	var data = FromElements("a", "bb", "aa", "c", "ccc")
+	assert.Equal(t, []string{"a", "bb", "ccc"}, ToUniqueBy[string, int](data, func(s string) int { return len(s) }).Collect())
+}
+
+func TestToDeUniqueBy(t *testing.T) {
+	var f = func(s string) int { return len(s) }
+	var data = FromElements("a", "bb", "aa", "c", "ccc")
+	assert.Equal(t, []string{"a", "bb", "ccc"}, ToDeUniqueBy[string, int](data, f).Collect())
+	var data2 = FromElements("a", "bb", "aa", "c", "ccc")
+	assert.Equal(t, []string{"ccc", "c", "aa"}, ToDeUniqueBy[string, int](data2, f).ToRev().Collect())
+	var data3 = FromElements("a", "bb", "aa", "c", "ccc")
+	assert.Equal(t, []string{"ccc", "c", "aa"}, ToDeUniqueBy[string, int](data3.ToRev(), f).Collect())
+}