WIP provenance works

Author: prdoyle

bosk-junit/src/main/java/works/bosk/junit/ParameterInjectionContextProvider.java

@@ -6,10 +6,13 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.HashSet;
 import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 import java.util.stream.Stream;
 import org.junit.jupiter.api.extension.Extension;
 import org.junit.jupiter.api.extension.ExtensionContext;
@@ -35,31 +38,53 @@ public boolean supportsTestTemplate(ExtensionContext context) {
 
 	@Override
 	public Stream<TestTemplateInvocationContext> provideTestTemplateInvocationContexts(ExtensionContext context) {
-		// TODO: This loop is not quite right.
-		// We should never call the same test method twice with the same parameter values,
-		// but if we expand branches for all injectors, we'll end up building up a cartesian product
-		// of parameter values for injectors that aren't even used by the test method.
-		// We should start with the test method parameters as the "root set",
-		// find their injectors, then find the parameters of those injectors,
-		// and so on, rather than just blindly expanding for all injectors.
-		//
- 		// For the time being, we sometimes end up calling the same test method
-		// multiple times with the same parameter values.
-		List<Branch> branches = List.of(Branch.empty());
-		for (var injectorClass : getAllInjectorClasses(context)) {
-			branches = expandedBranches(branches, injectorClass);
+		List<Parameter> requiredParameters = new ArrayList<>();
+		requiredParameters.addAll(asList(context.getRequiredTestMethod().getParameters()));
+		requiredParameters.addAll(asList(context.getRequiredTestClass().getDeclaredConstructors()[0].getParameters()));
+
+		List<Branch> allPossibleBranches = List.of(Branch.empty());
+		var allInjectorClasses = getAllInjectorClasses(context);
+		for (var injectorClass : allInjectorClasses) {
+			allPossibleBranches = expandedBranches(allPossibleBranches, injectorClass);
+		}
+
+		if (allPossibleBranches.isEmpty()) {
+			return Stream.empty();
 		}
 
-		List<Parameter> allParameters = new ArrayList<>();
-		allParameters.addAll(asList(context.getRequiredTestMethod().getParameters()));
-		allParameters.addAll(asList(context.getRequiredTestClass().getDeclaredConstructors()[0].getParameters()));
+		// At this stage, we have a list of branches that have instantiated
+		// all the injectors we could possibly have needed for any parameter,
+		// but some injectors might be for parameters that aren't used
+		// by the test method or the class constructor.
+		//
+		// If we don't prune out the unneeded injectors,
+		// we will end up calling the test method with the same parameters
+		// multiple times, varying only the values of parameters
+		// that aren't even used.
+		//
+		// Let's determine which injector classes we actually needed.
+		// We can do this by picking any Branch (they all have the same types of injectors)
+		// and seeing which injectors are needed
+		// to provide values for the requiredParameters.
+
+		var neededInjectorClasses = new HashSet<Class<? extends ParameterInjector>>();
+		Branch someBranch = allPossibleBranches.getFirst();
+		requiredParameters.forEach(p -> getNeededInjectorClasses(p, someBranch, neededInjectorClasses));
 
-		return branches.stream().flatMap(branch -> {
+		// And now we can recalculate the branch list, expanding only the required injector classes
+		List<Branch> neededBranches = List.of(Branch.empty());
+		for (var injectorClass : allInjectorClasses) { // The order matters here
+			if (neededInjectorClasses.contains(injectorClass)) {
+				neededBranches = expandedBranches(neededBranches, injectorClass);
+			}
+		}
+
+		return neededBranches.stream().flatMap(branch -> {
 			var valuesByInjector = new LinkedHashMap<ParameterInjector, List<Object>>();
-			allParameters.forEach(p -> {
+			requiredParameters.forEach(p -> {
 				ParameterInjector injector = branch.injectorFor(p);
 				if (injector != null) {
-					valuesByInjector.computeIfAbsent(injector, branch.toInject::get);
+					valuesByInjector.computeIfAbsent(injector, key -> branch.toInject.get(key).values());
 				}
 			});
 
@@ -69,7 +94,7 @@ public Stream<TestTemplateInvocationContext> provideTestTemplateInvocationContex
 			return combinations.stream().map(combo -> {
 				// Swizzle the combo into a useful map from parameter to value
 				var paramValueMap = new LinkedHashMap<Parameter, Object>();
-				allParameters.forEach(parameter -> {
+				requiredParameters.forEach(parameter -> {
 					// TODO: There's essentially a copy of this in Branch.withInjectors
 					ParameterInjector pi = branch.injectorFor(parameter);
 					if (pi != null) {
@@ -105,6 +130,9 @@ public Object resolveParameter(ParameterContext pc, ExtensionContext ec) throws
 		});
 	}
 
+	/**
+	 * @return the injector classes in the order they should be instantiated
+	 */
 	private static List<Class<? extends ParameterInjector>> getAllInjectorClasses(ExtensionContext context) {
 		List<Class<? >> bottomUp = new ArrayList<>();
 		for (var c = context.getRequiredTestClass(); c != Object.class; c = c.getSuperclass()) {
@@ -118,6 +146,22 @@ private static List<Class<? extends ParameterInjector>> getAllInjectorClasses(Ex
 		}
 		return allInjectors;
 	}
+
+	private void getNeededInjectorClasses(
+		Parameter param,
+		Branch branch,
+		Set<Class<? extends ParameterInjector>> needed
+	) {
+		ParameterInjector pi = branch.injectorFor(param);
+		if (pi != null && needed.add(pi.getClass())) {
+			var provenance = branch.toInject.get(pi).provenance();
+			var list = provenance.stream()
+				.map(ParameterInjector::getClass)
+				.toList();
+			needed.addAll(list);
+		}
+	}
+
 	/**
 	 * Expand each branch by instantiating all injectors of the given type.
 	 */
@@ -147,6 +191,19 @@ private static List<List<Object>> cartesianProduct(Collection<List<Object>> inpu
 		return result;
 	}
 
+	/**
+	 * @param values the subset of {@link ParameterInjector#values()} to be injected in this scenario
+	 * @param provenance the set of injectors required, directly or indirectly, to produce these values, with no guarantees on the order
+	 */
+	record Superposition(
+		List<Object> values,
+		Set<ParameterInjector> provenance
+	){
+		Superposition collapsed(Object singleValue) {
+			return new Superposition(List.of(singleValue), provenance);
+		}
+	}
+
 	/**
 	 * A list of injectors that have been instantiated so far.
 	 * Represents one "scenario" for cartesian product expansion of parameter values.
@@ -161,7 +218,7 @@ private static List<List<Object>> cartesianProduct(Collection<List<Object>> inpu
 	 *                 this map will contain just the one value used to construct that injector on this branch.
 	 */
 	record Branch(
-		Map<ParameterInjector, List<Object>> toInject
+		Map<ParameterInjector, Superposition> toInject
 	) {
 		static Branch empty() {
 			return new Branch(Map.of());
@@ -186,31 +243,45 @@ List<Branch> withInjectors(Class<? extends ParameterInjector> injectorType) {
 			// and will not further expand the cartesian product.)
 			List<List<Object>> valueLists = injectorsToUse.stream()
 				.map(toInject::get)
+				.map(Superposition::values)
 				.toList();
 
 			List<Branch> result = new ArrayList<>();
 			for (List<Object> combos : cartesianProduct(valueLists)) {
 				try {
 					List<Object> args = new ArrayList<>();
+					var injectorsUsed = new LinkedHashSet<ParameterInjector>();
 					for (var p: ctor.getParameters()) {
 						var pi = injectorFor(p);
+						injectorsUsed.add(pi);
 						var index = injectorsToUse.indexOf(pi);
+						assert index >= 0: "Internal error: injector not found for parameter " + p + " of constructor " + ctor;
 						args.add(combos.get(index));
 					}
 
-					// Instantiate the injector and add it to the map
+					// Instantiate the injector
 					var injector = (ParameterInjector) ctor.newInstance(args.toArray());
+
+					// Compute the provenance
+					var provenance = new HashSet<ParameterInjector>();
+					injectorsUsed.forEach(pi -> {
+						provenance.add(pi);
+						provenance.addAll(toInject.get(pi).provenance());
+					});
+					provenance.add(injector);
+
+					// Compute the new toInject map
 					var map = new LinkedHashMap<>(this.toInject);
-					map.put(injector, injector.values());
+					map.put(injector, new Superposition(injector.values(), provenance));
 
 					// Collapse the wavefunction for injectors that provided a value
 					int i = 0;
 					for (var pi: injectorsToUse) {
-						var possibleValues = toInject.get(pi);
+						var possibleValues = toInject.get(pi).values();
 						if (possibleValues.size() >= 2) {
 							// pi has provided a value on this branch.
 							// Record that fact for future uses of the same injector.
-							map.put(pi, List.of(args.get(i)));
+							map.put(pi, map.get(pi).collapsed(args.get(i)));
 						}
 						++i;
 					}