Small code tweaks on recent PR

bbakerman · bbakerman · commit d26ed000666f · 2024-02-20T22:19:59.000+11:00
diff --git a/src/main/java/graphql/execution/instrumentation/ChainedInstrumentation.java b/src/main/java/graphql/execution/instrumentation/ChainedInstrumentation.java
@@ -90,10 +90,11 @@ private <T> T chainedInstrument(InstrumentationState state, T input, ChainedInst
     }
 
     protected <T> ImmutableList<T> chainedMapAndDropNulls(InstrumentationState state, BiFunction<Instrumentation, InstrumentationState, T> mapper) {
+        ChainedInstrumentationState chainedInstrumentationState = (ChainedInstrumentationState) state;
         ImmutableList.Builder<T> result = ImmutableList.builderWithExpectedSize(instrumentations.size());
         for (int i = 0; i < instrumentations.size(); i++) {
             Instrumentation instrumentation = instrumentations.get(i);
-            InstrumentationState specificState = ((ChainedInstrumentationState) state).getState(i);
+            InstrumentationState specificState = chainedInstrumentationState.getState(i);
             T value = mapper.apply(instrumentation, specificState);
             if (value != null) {
                 result.add(value);
@@ -102,10 +103,11 @@ protected <T> ImmutableList<T> chainedMapAndDropNulls(InstrumentationState state
         return result.build();
     }
 
-    protected <T> void chainedConsume(InstrumentationState state, BiConsumer<Instrumentation, InstrumentationState> stateConsumer) {
+    protected void chainedConsume(InstrumentationState state, BiConsumer<Instrumentation, InstrumentationState> stateConsumer) {
+        ChainedInstrumentationState chainedInstrumentationState = (ChainedInstrumentationState) state;
         for (int i = 0; i < instrumentations.size(); i++) {
             Instrumentation instrumentation = instrumentations.get(i);
-            InstrumentationState specificState = ((ChainedInstrumentationState) state).getState(i);
+            InstrumentationState specificState = chainedInstrumentationState.getState(i);
             stateConsumer.accept(instrumentation, specificState);
         }
     }
@@ -358,7 +360,7 @@ public CompletableFuture<ExecutionResult> instrumentExecutionResult(ExecutionRes
         CompletableFuture<List<ExecutionResult>> resultsFuture = Async.eachSequentially(entries, (entry, prevResults) -> {
             Instrumentation instrumentation = entry.getKey();
             InstrumentationState specificState = entry.getValue();
-            ExecutionResult lastResult = prevResults.size() > 0 ? prevResults.get(prevResults.size() - 1) : executionResult;
+            ExecutionResult lastResult = !prevResults.isEmpty() ? prevResults.get(prevResults.size() - 1) : executionResult;
             return instrumentation.instrumentExecutionResult(lastResult, parameters, specificState);
         });
         return resultsFuture.thenApply((results) -> results.isEmpty() ? executionResult : results.get(results.size() - 1));
@@ -483,10 +485,8 @@ public void onCompleted(Object result, Throwable t) {
     }
 
     @FunctionalInterface
-    private interface ChainedInstrumentationFunction<I, S, A, R> {
-
-        R apply(I t, S u, A v);
-
+    private interface ChainedInstrumentationFunction<I, S, V, R> {
+        R apply(I instrumentation, S state, V value);
     }
 
 
diff --git a/src/test/groovy/graphql/execution/instrumentation/ChainedInstrumentationStateTest.groovy b/src/test/groovy/graphql/execution/instrumentation/ChainedInstrumentationStateTest.groovy
@@ -90,6 +90,8 @@ class ChainedInstrumentationStateTest extends Specification {
 
         then:
 
+        chainedInstrumentation.getInstrumentations().size() == 4
+
         a.executionList == expected
         b.executionList == expected
         c.executionList == expected
