Add validation of parameters in @BeforeTemplates and @AfterTemplates

core/src/main/java/com/google/errorprone/refaster/RefasterRule.java

@@ -16,6 +16,7 @@
 
 package com.google.errorprone.refaster;
 
+import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkState;
 
 import com.google.auto.value.AutoValue;
@@ -26,6 +27,7 @@
 import com.google.common.collect.ImmutableClassToInstanceMap;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
+import com.google.common.collect.Sets;
 import com.google.errorprone.CodeTransformer;
 import com.google.errorprone.DescriptionListener;
 import com.google.errorprone.SubContext;
@@ -38,6 +40,7 @@
 import java.lang.annotation.Annotation;
 import java.util.Collection;
 import java.util.List;
+import java.util.Set;
 import javax.tools.JavaFileManager;
 
 /**
@@ -83,7 +86,21 @@ public abstract class RefasterRule<M extends TemplateMatch, T extends Template<M
           beforeTemplate.getClass(),
           qualifiedTemplateClass);
     }
+
     for (Template<?> afterTemplate : afterTemplates) {
+      Set<String> missingArguments =
+          Sets.difference(
+              afterTemplate.expressionArgumentTypes().keySet(),
+              beforeTemplates.stream()
+                  .<Set<String>>map(t -> t.expressionArgumentTypes().keySet())
+                  .reduce(Sets::intersection)
+                  .get());
+      checkArgument(
+          missingArguments.isEmpty(),
+          "@AfterTemplate of %s defines arguments that are not present in all @BeforeTemplate methods: %s",
+          qualifiedTemplateClass,
+          missingArguments);
+
       checkState(
           afterTemplate.getClass().equals(templateType),
           "Expected all templates to be of type %s but found template of type %s in %s",

core/src/test/java/com/google/errorprone/refaster/TemplateIntegrationTest.java

@@ -88,6 +88,16 @@ private void runTest(String testName) throws IOException {
     expectTransforms(transformer, input, output);
   }
 
+  @Test
+  public void keyBindingError() {
+    IllegalArgumentException failure =
+        assertThrows(IllegalArgumentException.class, () -> runTest("KeyBindingErrorTemplate"));
+    assertThat(failure)
+        .hasMessageThat()
+        .matches(
+            "@AfterTemplate of .*\\.KeyBindingErrorTemplate defines arguments that are not present in all @BeforeTemplate methods: \\[b, c\\]");
+  }
+
   @Test
   public void binary() throws IOException {
     runTest("BinaryTemplate");

core/src/test/java/com/google/errorprone/refaster/testdata/template/KeyBindingErrorTemplate.java

@@ -0,0 +1,43 @@
+/*
+ * Copyright 2021 The Error Prone Authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.errorprone.refaster.testdata.template;
+
+import com.google.errorprone.refaster.annotation.AfterTemplate;
+import com.google.errorprone.refaster.annotation.BeforeTemplate;
+
+/**
+ * Refaster template which triggers an {@link IllegalArgumentException} upon compilation.
+ *
+ * <p>The {@code @AfterTemplate} defines parameters that are not present in all
+ * {@code @BeforeTemplate} methods.
+ */
+public class KeyBindingErrorTemplate {
+  @BeforeTemplate
+  public int before1(int a, int b) {
+    return a + b;
+  }
+
+  @BeforeTemplate
+  public int before2(int a, int c) {
+    return a + c;
+  }
+
+  @AfterTemplate
+  public int after(int a, int b, int c) {
+    return a + b + c;
+  }
+}