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

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

@@ -2,11 +2,61 @@
 
 import java.util.*;
 import java.util.function.*;
+import java.util.stream.*;
 
 import net.jqwik.api.*;
 import net.jqwik.api.state.*;
 
 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;
 	final boolean accessState;
@@ -30,19 +80,29 @@ private ShrinkableChainIteration(
 		this.precondition = precondition;
 		this.accessState = accessState;
 		this.changeState = changeState;
-		this.shrinkable = shrinkable;
+		// 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
 	public String toString() {
 		return String.format(
 			"Iteration[accessState=%s, changeState=%s, transformation=%s]",
-			accessState, changeState, shrinkable.value().transformation()
+			accessState, changeState, transformer().transformation()
 		);
 	}
 
 	boolean isEndOfChain() {
-		return shrinkable.value().equals(Transformer.END_OF_CHAIN);
+		return transformer().equals(Transformer.END_OF_CHAIN);
 	}
 
 	Optional<Predicate<T>> precondition() {

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

@@ -119,6 +119,102 @@ ActionChainArbitrary<String> xOrFailing() {
 						  .withMaxTransformations(30);
 	}
 
+	static class SetMutatingChainState {
+		final List<String> actualOps = new ArrayList<>();
+		boolean hasPrints;
+		final Set<Integer> set = new HashSet<>();
+
+		@Override
+		public String toString() {
+			return "set=" + set + ", actualOps=" + actualOps;
+		}
+	}
+
+	@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())
+			.describedAs("chain.transformations() should be the same as the list of operations in finalState.actualOps, final state is %s", finalState.set)
+			.isEqualTo(finalState.actualOps);
+	}
+
+	@Provide
+	public ActionChainArbitrary<SetMutatingChainState> setMutatingChain() {
+		return
+			ActionChain
+				.startWith(SetMutatingChainState::new)
+				// This is an action that does not depend on the state to produce the transformation
+				.addAction(
+					1,
+					Action.just("clear anyway", state -> {
+						state.actualOps.add("clear anyway");
+						state.set.clear();
+						return state;
+					})
+				)
+				// Below actions depend on the state to derive the transformations
+				.addAction(
+					1,
+					(Action.Dependent<SetMutatingChainState>)
+						state ->
+							Arbitraries
+								.just(
+									state.set.isEmpty()
+										? Transformer.noop()
+										: Transformer.<SetMutatingChainState>mutate("clear " + state.set, set -> {
+										state.actualOps.add("clear " + set.set);
+										state.set.clear();
+									})
+								)
+				)
+				.addAction(
+					4,
+					(Action.Dependent<SetMutatingChainState>)
+						state ->
+							Arbitraries
+								.integers()
+								.between(1, 10)
+								.map(i -> {
+										 if (state.set.contains(i)) {
+											 return Transformer.noop();
+										 }
+										 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(7);
+	}
+
 	@Property
 	void chainChoosesBetweenTwoActions(@ForAll("xOrY") ActionChain<String> chain) {
 		String result = chain.run();