From 15f606701a4f9a4fc4f01d67306aaaa45a656a51 Mon Sep 17 00:00:00 2001
From: Johannes Link <business@johanneslink.net>
Date: Thu, 16 Mar 2023 14:09:35 +0100
Subject: [PATCH] Add uniqueElements() methods to StreamableArbitrary interface
 See https://github.com/jqwik-team/jqwik/issues/466

---
 TODO.md                                       |  3 --
 .../jqwik/api/arbitraries/SetArbitrary.java   |  9 +++++
 .../api/arbitraries/StreamableArbitrary.java  | 29 +++++++++++++++
 .../arbitraries/DefaultArrayArbitrary.java    |  2 +-
 .../arbitraries/DefaultIteratorArbitrary.java |  2 +-
 .../arbitraries/DefaultListArbitrary.java     |  2 +-
 .../arbitraries/DefaultSetArbitrary.java      |  5 +++
 .../arbitraries/DefaultStreamArbitrary.java   |  2 +-
 .../arbitraries/MultivalueArbitraryBase.java  |  5 +++
 .../UniqueElementsConfigurator.java           | 35 +++++--------------
 10 files changed, 60 insertions(+), 34 deletions(-)

diff --git a/TODO.md b/TODO.md
index 41d3bf06b..e2fdbbd6a 100644
--- a/TODO.md
+++ b/TODO.md
@@ -1,8 +1,5 @@
 # 1.7.3
 
-    - Support UniqueElements annotation for extensions.
-      See: https://github.com/jqwik-team/jqwik/issues/466
-
     - Add Arbitrary.ignoreException(maxMisses, ...)
       See https://github.com/jqwik-team/jqwik/issues/462.
 
diff --git a/api/src/main/java/net/jqwik/api/arbitraries/SetArbitrary.java b/api/src/main/java/net/jqwik/api/arbitraries/SetArbitrary.java
index 191e0c63a..d8a66e983 100644
--- a/api/src/main/java/net/jqwik/api/arbitraries/SetArbitrary.java
+++ b/api/src/main/java/net/jqwik/api/arbitraries/SetArbitrary.java
@@ -72,6 +72,15 @@ default SetArbitrary<T> ofSize(int size) {
 	@API(status = MAINTAINED, since = "1.4.0")
 	<U> Arbitrary<Set<U>> flatMapEach(BiFunction<Set<T>, T, Arbitrary<U>> flatMapper);
 
+	/**
+	 * Do not use. Sets have unique elements anyway.
+	 * It only exists for purposes of symmetry.
+	 *
+	 * @return same instance of arbitrary
+	 */
+	@API(status = MAINTAINED, since = "1.7.3")
+	SetArbitrary<T> uniqueElements();
+
 	/**
 	 * Add the constraint that elements of the generated set must be unique
 	 * relating to an element's "feature" being extracted using the
diff --git a/api/src/main/java/net/jqwik/api/arbitraries/StreamableArbitrary.java b/api/src/main/java/net/jqwik/api/arbitraries/StreamableArbitrary.java
index 8e91bbc98..f16b2a213 100644
--- a/api/src/main/java/net/jqwik/api/arbitraries/StreamableArbitrary.java
+++ b/api/src/main/java/net/jqwik/api/arbitraries/StreamableArbitrary.java
@@ -50,4 +50,33 @@ default StreamableArbitrary<T, U> ofSize(int size) {
 	 */
 	@API(status = EXPERIMENTAL, since = "1.5.3")
 	StreamableArbitrary<T, U> withSizeDistribution(RandomDistribution distribution);
+
+	/**
+	 * Add the constraint that elements of the generated streamable must be unique,
+	 * i.e. no two elements must return true when being compared using {@linkplain Object#equals(Object)}.
+	 *
+	 * <p>
+	 *     The constraint can be combined with other {@linkplain #uniqueElements(Function)} constraints.
+	 * </p>
+	 *
+	 * @return new arbitrary instance
+	 */
+	@API(status = MAINTAINED, since = "1.7.3")
+	StreamableArbitrary<T, U> uniqueElements();
+
+	/**
+	 * Add the constraint that elements of the generated streamable must be unique
+	 * relating to an element's "feature" being extracted using the
+	 * {@code by} function.
+	 * The extracted features are being compared using {@linkplain Object#equals(Object)}.
+	 *
+	 * <p>
+	 *     The constraint can be combined with other {@linkplain #uniqueElements(Function)} constraints.
+	 * </p>
+	 *
+	 * @return new arbitrary instance
+	 */
+	@API(status = MAINTAINED, since = "1.7.3")
+	StreamableArbitrary<T, U> uniqueElements(Function<T, Object> by);
+
 }
diff --git a/engine/src/main/java/net/jqwik/engine/properties/arbitraries/DefaultArrayArbitrary.java b/engine/src/main/java/net/jqwik/engine/properties/arbitraries/DefaultArrayArbitrary.java
index 0a06419f2..0f6e55912 100644
--- a/engine/src/main/java/net/jqwik/engine/properties/arbitraries/DefaultArrayArbitrary.java
+++ b/engine/src/main/java/net/jqwik/engine/properties/arbitraries/DefaultArrayArbitrary.java
@@ -91,7 +91,7 @@ protected Iterable<T> toIterable(A array) {
 
 	@Override
 	public ArrayArbitrary<T, A> uniqueElements() {
-		return (ArrayArbitrary<T, A>) uniqueElements(FeatureExtractor.identity());
+		return (ArrayArbitrary<T, A>) super.uniqueElements();
 	}
 
 	@Override
diff --git a/engine/src/main/java/net/jqwik/engine/properties/arbitraries/DefaultIteratorArbitrary.java b/engine/src/main/java/net/jqwik/engine/properties/arbitraries/DefaultIteratorArbitrary.java
index 49f967381..ffef597f4 100644
--- a/engine/src/main/java/net/jqwik/engine/properties/arbitraries/DefaultIteratorArbitrary.java
+++ b/engine/src/main/java/net/jqwik/engine/properties/arbitraries/DefaultIteratorArbitrary.java
@@ -67,7 +67,7 @@ public IteratorArbitrary<T> uniqueElements(Function<T, Object> by) {
 
 	@Override
 	public IteratorArbitrary<T> uniqueElements() {
-		return (IteratorArbitrary<T>) uniqueElements(FeatureExtractor.identity());
+		return (IteratorArbitrary<T>) super.uniqueElements();
 	}
 
 }
diff --git a/engine/src/main/java/net/jqwik/engine/properties/arbitraries/DefaultListArbitrary.java b/engine/src/main/java/net/jqwik/engine/properties/arbitraries/DefaultListArbitrary.java
index 6ccdb0980..194d1bfa4 100644
--- a/engine/src/main/java/net/jqwik/engine/properties/arbitraries/DefaultListArbitrary.java
+++ b/engine/src/main/java/net/jqwik/engine/properties/arbitraries/DefaultListArbitrary.java
@@ -84,6 +84,6 @@ public ListArbitrary<T> uniqueElements(Function<T, Object> by) {
 
 	@Override
 	public ListArbitrary<T> uniqueElements() {
-		return (ListArbitrary<T>) uniqueElements(FeatureExtractor.identity());
+		return (ListArbitrary<T>) super.uniqueElements();
 	}
 }
diff --git a/engine/src/main/java/net/jqwik/engine/properties/arbitraries/DefaultSetArbitrary.java b/engine/src/main/java/net/jqwik/engine/properties/arbitraries/DefaultSetArbitrary.java
index c5288213e..ad48dd029 100644
--- a/engine/src/main/java/net/jqwik/engine/properties/arbitraries/DefaultSetArbitrary.java
+++ b/engine/src/main/java/net/jqwik/engine/properties/arbitraries/DefaultSetArbitrary.java
@@ -87,6 +87,11 @@ public <U> Arbitrary<Set<U>> flatMapEach(BiFunction<Set<T>, T, Arbitrary<U>> fla
 		});
 	}
 
+	@Override
+	public SetArbitrary<T> uniqueElements() {
+		return this;
+	}
+
 	@Override
 	public SetArbitrary<T> uniqueElements(Function<T, Object> by) {
 		FeatureExtractor<T> featureExtractor = by::apply;
diff --git a/engine/src/main/java/net/jqwik/engine/properties/arbitraries/DefaultStreamArbitrary.java b/engine/src/main/java/net/jqwik/engine/properties/arbitraries/DefaultStreamArbitrary.java
index 3ac518b3a..fa289d647 100644
--- a/engine/src/main/java/net/jqwik/engine/properties/arbitraries/DefaultStreamArbitrary.java
+++ b/engine/src/main/java/net/jqwik/engine/properties/arbitraries/DefaultStreamArbitrary.java
@@ -69,7 +69,7 @@ public StreamArbitrary<T> uniqueElements(Function<T, Object> by) {
 
 	@Override
 	public StreamArbitrary<T> uniqueElements() {
-		return (StreamArbitrary<T>) uniqueElements(FeatureExtractor.identity());
+		return (StreamArbitrary<T>) super.uniqueElements();
 	}
 
 }
diff --git a/engine/src/main/java/net/jqwik/engine/properties/arbitraries/MultivalueArbitraryBase.java b/engine/src/main/java/net/jqwik/engine/properties/arbitraries/MultivalueArbitraryBase.java
index 0b9aaee91..0058b7752 100644
--- a/engine/src/main/java/net/jqwik/engine/properties/arbitraries/MultivalueArbitraryBase.java
+++ b/engine/src/main/java/net/jqwik/engine/properties/arbitraries/MultivalueArbitraryBase.java
@@ -81,6 +81,11 @@ public <R> Arbitrary<R> reduce(R initial, BiFunction<R, T, R> accumulator) {
 		});
 	}
 
+	@Override
+	public StreamableArbitrary<T, U> uniqueElements() {
+		return uniqueElements(FeatureExtractor.identity());
+	}
+
 	protected abstract Iterable<T> toIterable(U streamable);
 
 	protected StreamableArbitrary<T, U> uniqueElements(FeatureExtractor<T> by) {
diff --git a/engine/src/main/java/net/jqwik/engine/properties/configurators/UniqueElementsConfigurator.java b/engine/src/main/java/net/jqwik/engine/properties/configurators/UniqueElementsConfigurator.java
index 9557795ea..9e11ebc6f 100644
--- a/engine/src/main/java/net/jqwik/engine/properties/configurators/UniqueElementsConfigurator.java
+++ b/engine/src/main/java/net/jqwik/engine/properties/configurators/UniqueElementsConfigurator.java
@@ -19,20 +19,12 @@ public class UniqueElementsConfigurator implements ArbitraryConfigurator {
 	@Override
 	public <T> Arbitrary<T> configure(Arbitrary<T> arbitrary, TypeUsage targetType) {
 		return targetType.findAnnotation(UniqueElements.class).map(uniqueness -> {
-			if (arbitrary instanceof ListArbitrary) {
-				return (Arbitrary<T>) configureListArbitrary((ListArbitrary<?>) arbitrary, uniqueness);
-			}
 			if (arbitrary instanceof SetArbitrary) {
+				// Handle SetArbitrary explicitly for optimization
 				return (Arbitrary<T>) configureSetArbitrary((SetArbitrary<?>) arbitrary, uniqueness);
 			}
-			if (arbitrary instanceof ArrayArbitrary) {
-				return (Arbitrary<T>) configureArrayArbitrary((ArrayArbitrary<?, ?>) arbitrary, uniqueness);
-			}
-			if (arbitrary instanceof StreamArbitrary) {
-				return (Arbitrary<T>) configureStreamArbitrary((StreamArbitrary<?>) arbitrary, uniqueness);
-			}
-			if (arbitrary instanceof IteratorArbitrary) {
-				return (Arbitrary<T>) configureIteratorArbitrary((IteratorArbitrary<?>) arbitrary, uniqueness);
+			if (arbitrary instanceof StreamableArbitrary) {
+				return (Arbitrary<T>) configureStreamableArbitrary((StreamableArbitrary<?, ?>) arbitrary, uniqueness);
 			}
 			if (targetType.isAssignableFrom(List.class)) {
 				Arbitrary<List<?>> listArbitrary = (Arbitrary<List<?>>) arbitrary;
@@ -76,27 +68,16 @@ private boolean isUnique(Collection<?> list, Function<Object, Object> extractor)
 		return set.size() == list.size();
 	}
 
-	private <T> Arbitrary<?> configureListArbitrary(ListArbitrary<T> arbitrary, UniqueElements uniqueness) {
+	private <T> Arbitrary<?> configureStreamableArbitrary(StreamableArbitrary<T, ?> arbitrary, UniqueElements uniqueness) {
 		Function<T, Object> extractor = (Function<T, Object>) extractor(uniqueness);
 		return arbitrary.uniqueElements(extractor);
 	}
 
 	private <T> Arbitrary<?> configureSetArbitrary(SetArbitrary<T> arbitrary, UniqueElements uniqueness) {
-		Function<T, Object> extractor = (Function<T, Object>) extractor(uniqueness);
-		return arbitrary.uniqueElements(extractor);
-	}
-
-	private <T> Arbitrary<?> configureArrayArbitrary(ArrayArbitrary<T, ?> arbitrary, UniqueElements uniqueness) {
-		Function<T, Object> extractor = (Function<T, Object>) extractor(uniqueness);
-		return arbitrary.uniqueElements(extractor);
-	}
-
-	private <T> Arbitrary<?> configureStreamArbitrary(StreamArbitrary<T> arbitrary, UniqueElements uniqueness) {
-		Function<T, Object> extractor = (Function<T, Object>) extractor(uniqueness);
-		return arbitrary.uniqueElements(extractor);
-	}
-
-	private <T> Arbitrary<?> configureIteratorArbitrary(IteratorArbitrary<T> arbitrary, UniqueElements uniqueness) {
+		Class<? extends Function<?, Object>> extractorClass = uniqueness.by();
+		if (extractorClass.equals(UniqueElements.NOT_SET.class)) {
+			return arbitrary;
+		}
 		Function<T, Object> extractor = (Function<T, Object>) extractor(uniqueness);
 		return arbitrary.uniqueElements(extractor);
 	}