Cache several shrink resultsfor every chain step, so it does not re-access state model state during shrink phase

This might take slightly more memory and produce suboptimal shrinks,
however makes shrink possible for state-accessing transformations.

Fixes #428

Author: Vladimir Sitnikov

engine/src/main/java/net/jqwik/engine/properties/state/ShrinkableChainIteration.java

@@ -2,16 +2,63 @@
 
 import java.util.*;
 import java.util.function.*;
+import java.util.stream.*;
 
 import net.jqwik.api.*;
 import net.jqwik.api.state.*;
 
-import org.jetbrains.annotations.*;
-
 class ShrinkableChainIteration<T> {
+	// Larger values might improve shrink quality, however, they increase the shrink space, so it might increase shrink duration
+	private final static int NUM_SAMPLES_IN_EAGER_CHAIN_SHRINK = Integer.getInteger("jqwik.eagerChainShrinkSamples", 2);
+
+	static class ShrinkableWithEagerValue<T> implements Shrinkable<T> {
+		protected final Shrinkable<T> base;
+		private final ShrinkingDistance distance;
+		final T value;
+
+		ShrinkableWithEagerValue(Shrinkable<T> base) {
+			this.base = base;
+			this.distance = base.distance();
+			this.value = base.value();
+		}
+
+		@Override
+		public T value() {
+			return value;
+		}
+
+		@Override
+		public Stream<Shrinkable<T>> shrink() {
+			return base.shrink();
+		}
+
+		@Override
+		public ShrinkingDistance distance() {
+			return distance;
+		}
+	}
+
+	static class EagerShrinkable<T> extends ShrinkableWithEagerValue<T> {
+		private final List<Shrinkable<T>> shrinkResults;
+
+		EagerShrinkable(Shrinkable<T> base, int numSamples) {
+			super(base);
+			this.shrinkResults =
+				base.shrink()
+					.sorted(Comparator.comparing(Shrinkable::distance))
+					.limit(numSamples)
+					.map(ShrinkableWithEagerValue::new)
+					.collect(Collectors.toList());
+		}
+
+		@Override
+		public Stream<Shrinkable<T>> shrink() {
+			return shrinkResults.stream();
+		}
+	}
+
 	final Shrinkable<Transformer<T>> shrinkable;
 	private final Predicate<T> precondition;
-	private final @Nullable Transformer<T> transformer;
 	final boolean accessState;
 	final boolean changeState;
 
@@ -33,10 +80,17 @@ private ShrinkableChainIteration(
 		this.precondition = precondition;
 		this.accessState = accessState;
 		this.changeState = changeState;
-		this.shrinkable = shrinkable;
-		// Transformer method might access state, so we need to cache the value here
-		// otherwise it might be evaluated with wrong state (e.g. after chain executes)
-		this.transformer = accessState ? shrinkable.value() : null;
+		// When the shrinkable does not access state, we could just use it as is for ".value()", and ".shrink()"
+		// If we get LazyShrinkable here, it means we are in a shrinking phase, so we know ".shrink()" will be called only
+		// in case the subsequent execution fails. So we can just keep LazyShrinkable as is
+		// Otherwise, we need to eagerly evaluate the shrinkables to since the state might change by appyling subsequent transformers,
+		// so we won't be able to access the state anymore.
+		// See https://github.com/jlink/jqwik/issues/428
+		if (!accessState || shrinkable instanceof ShrinkableChainIteration.ShrinkableWithEagerValue) {
+			this.shrinkable = shrinkable;
+		} else {
+			this.shrinkable = new EagerShrinkable<>(shrinkable, NUM_SAMPLES_IN_EAGER_CHAIN_SHRINK);
+		}
 	}
 
 	@Override
@@ -71,6 +125,6 @@ String transformation() {
 	}
 
 	Transformer<T> transformer() {
-		return transformer == null ? shrinkable.value() : transformer;
+		return shrinkable.value();
 	}
 }

engine/src/test/java/net/jqwik/engine/properties/state/ActionChainArbitraryTests.java

@@ -121,6 +121,7 @@ ActionChainArbitrary<String> xOrFailing() {
 
 	static class SetMutatingChainState {
 		final List<String> actualOps = new ArrayList<>();
+		boolean hasPrints;
 		final Set<Integer> set = new HashSet<>();
 
 		@Override
@@ -131,6 +132,14 @@ public String toString() {
 
 	@Property
 	void chainActionsAreProperlyDescribedEvenAfterChainExecution(@ForAll("setMutatingChain") ActionChain<SetMutatingChainState> chain) {
+		chain = chain.withInvariant(
+			state -> {
+				if (state.hasPrints) {
+					assertThat(state.actualOps).hasSizeLessThan(5);
+				}
+			}
+		);
+
 		SetMutatingChainState finalState = chain.run();
 
 		assertThat(chain.transformations())
@@ -168,7 +177,7 @@ public ActionChainArbitrary<SetMutatingChainState> setMutatingChain() {
 								)
 				)
 				.addAction(
-					2,
+					4,
 					(Action.Dependent<SetMutatingChainState>)
 						state ->
 							Arbitraries
@@ -177,16 +186,33 @@ public ActionChainArbitrary<SetMutatingChainState> setMutatingChain() {
 								.map(i -> {
 										 if (state.set.contains(i)) {
 											 return Transformer.noop();
-										 } else {
-											 return Transformer.mutate("add " + i + " to " + state.set, newState -> {
-												 newState.actualOps.add("add " + i + " to " + newState.set);
-												 newState.set.add(i);
-											 });
 										 }
+										 return Transformer.mutate("add " + i + " to " + state.set, newState -> {
+											 newState.actualOps.add("add " + i + " to " + newState.set);
+											 newState.set.add(i);
+										 });
+									 }
+								)
+				)
+				.addAction(
+					2,
+					(Action.Dependent<SetMutatingChainState>)
+						state ->
+							state.set.isEmpty() ? Arbitraries.just(Transformer.noop()) :
+							Arbitraries
+								.of(state.set)
+								.map(i -> {
+										 if (!state.set.contains(i)) {
+											 throw new IllegalStateException("The set does not contain " + i + ", current state is " + state);
+										 }
+										 return Transformer.mutate("print " + i + " from " + state.set, newState -> {
+											 newState.actualOps.add("print " + i + " from " + newState.set);
+											 newState.hasPrints = true;
+										 });
 									 }
 								)
 				)
-				.withMaxTransformations(5);
+				.withMaxTransformations(7);
 	}
 
 	@Property