Suggestions

documentation-support/src/main/java/tech/picnic/errorprone/documentation/BugPatternExtractor.java

@@ -16,7 +16,6 @@
 import com.sun.source.tree.ClassTree;
 import com.sun.tools.javac.code.Attribute;
 import com.sun.tools.javac.code.Symbol.ClassSymbol;
-import com.sun.tools.javac.util.Context;
 import javax.lang.model.element.AnnotationValue;
 import tech.picnic.errorprone.documentation.BugPatternExtractor.BugPatternDocumentation;
 
@@ -26,7 +25,7 @@
 @Immutable
 final class BugPatternExtractor implements Extractor<BugPatternDocumentation> {
   @Override
-  public BugPatternDocumentation extract(ClassTree tree, Context context) {
+  public BugPatternDocumentation extract(ClassTree tree, VisitorState state) {
     ClassSymbol symbol = ASTHelpers.getSymbol(tree);
     BugPattern annotation = symbol.getAnnotation(BugPattern.class);
     requireNonNull(annotation, "BugPattern annotation must be present");

documentation-support/src/main/java/tech/picnic/errorprone/documentation/BugPatternTestExtractor.java

@@ -1,25 +1,30 @@
 package tech.picnic.errorprone.documentation;
 
+import static com.google.errorprone.matchers.Matchers.allOf;
+import static com.google.errorprone.matchers.Matchers.anything;
+import static com.google.errorprone.matchers.Matchers.argument;
+import static com.google.errorprone.matchers.Matchers.classLiteral;
 import static com.google.errorprone.matchers.Matchers.hasAnnotation;
 import static com.google.errorprone.matchers.Matchers.instanceMethod;
+import static com.google.errorprone.matchers.Matchers.toType;
 import static com.google.errorprone.matchers.method.MethodMatchers.staticMethod;
 
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableList;
 import com.google.errorprone.VisitorState;
 import com.google.errorprone.annotations.Immutable;
-import com.google.errorprone.annotations.Var;
 import com.google.errorprone.matchers.Matcher;
 import com.google.errorprone.util.ASTHelpers;
 import com.sun.source.tree.ClassTree;
 import com.sun.source.tree.ExpressionTree;
 import com.sun.source.tree.MemberSelectTree;
 import com.sun.source.tree.MethodInvocationTree;
 import com.sun.source.tree.MethodTree;
+import com.sun.source.tree.Tree;
 import com.sun.source.util.TreeScanner;
-import com.sun.tools.javac.util.Context;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
 import org.jspecify.annotations.Nullable;
 import tech.picnic.errorprone.documentation.BugPatternTestExtractor.BugPatternTestDocumentation;
 
@@ -29,21 +34,20 @@
  */
 @Immutable
 final class BugPatternTestExtractor implements Extractor<BugPatternTestDocumentation> {
-  private static final Matcher<MethodTree> JUNIT_TEST_METHOD =
-      hasAnnotation("org.junit.jupiter.api.Test");
+  private static final Matcher<Tree> JUNIT_TEST_METHOD =
+      toType(MethodTree.class, hasAnnotation("org.junit.jupiter.api.Test"));
 
   // XXX: Improve support for correctly extracting multiple sources from a single
   // `{BugCheckerRefactoring,Compilation}TestHelper` test.
   @Override
-  public BugPatternTestDocumentation extract(ClassTree tree, Context context) {
-    VisitorState state = VisitorState.createForUtilityPurposes(context);
-    CollectBugPatternTests scanner = new CollectBugPatternTests(state);
+  public BugPatternTestDocumentation extract(ClassTree tree, VisitorState state) {
+    CollectBugPatternTests scanner = new CollectBugPatternTests();
 
-    tree.getMembers().stream()
-        .filter(MethodTree.class::isInstance)
-        .map(MethodTree.class::cast)
-        .filter(m -> JUNIT_TEST_METHOD.matches(m, state))
-        .forEach(m -> scanner.scan(m, null));
+    for (Tree m : tree.getMembers()) {
+      if (JUNIT_TEST_METHOD.matches(m, state)) {
+        scanner.scan(m, state);
+      }
+    }
 
     String className = tree.getSimpleName().toString();
     return new AutoValue_BugPatternTestExtractor_BugPatternTestDocumentation(
@@ -66,13 +70,16 @@ public boolean canExtract(ClassTree tree, VisitorState state) {
     return scanBugPatternTest.hasTestUsingClassInstance(bugPatternName);
   }
 
+  // XXX: Consider replacing this type with an anonymous class in a method. Possibly also below.
   private static final class ScanBugPatternTest extends TreeScanner<@Nullable Void, VisitorState> {
-    private static final Matcher<ExpressionTree> BUG_PATTERN_TEST_METHOD =
-        staticMethod()
-            .onDescendantOfAny(
-                "com.google.errorprone.CompilationTestHelper",
-                "com.google.errorprone.BugCheckerRefactoringTestHelper")
-            .named("newInstance");
+    private static final Matcher<MethodInvocationTree> BUG_PATTERN_TEST_METHOD =
+        allOf(
+            staticMethod()
+                .onDescendantOfAny(
+                    "com.google.errorprone.CompilationTestHelper",
+                    "com.google.errorprone.BugCheckerRefactoringTestHelper")
+                .named("newInstance"),
+            argument(0, classLiteral(anything())));
 
     private final List<String> encounteredClasses = new ArrayList<>();
 
@@ -86,12 +93,13 @@ boolean hasTestUsingClassInstance(String clazz) {
         MemberSelectTree firstArgumentTree = (MemberSelectTree) node.getArguments().get(0);
         encounteredClasses.add(firstArgumentTree.getExpression().toString());
       }
+
       return super.visitMethodInvocation(node, state);
     }
   }
 
   private static final class CollectBugPatternTests
-      extends TreeScanner<@Nullable Void, @Nullable Void> {
+      extends TreeScanner<@Nullable Void, VisitorState> {
     private static final Matcher<ExpressionTree> IDENTIFICATION_SOURCE_LINES =
         instanceMethod()
             .onDescendantOf("com.google.errorprone.CompilationTestHelper")
@@ -105,16 +113,9 @@ private static final class CollectBugPatternTests
             .onDescendantOf("com.google.errorprone.BugCheckerRefactoringTestHelper.ExpectOutput")
             .named("addOutputLines");
 
-    private final VisitorState state;
     private final List<String> identificationTests = new ArrayList<>();
     private final List<BugPatternReplacementTestDocumentation> replacementTests = new ArrayList<>();
 
-    @Var private String replacementOutputLines = "";
-
-    CollectBugPatternTests(VisitorState state) {
-      this.state = state;
-    }
-
     public ImmutableList<String> getIdentificationTests() {
       return ImmutableList.copyOf(identificationTests);
     }
@@ -123,39 +124,55 @@ public ImmutableList<BugPatternReplacementTestDocumentation> getReplacementTests
       return ImmutableList.copyOf(replacementTests);
     }
 
+    // XXX: Consider:
+    // - Whether to omit or handle differently identification tests without `// BUG: Diagnostic
+    //   (contains|matches)` markers.
+    // - Whether to omit or handle differently replacement tests with identical input and output.
+    //   (Though arguably we should have a separate checker which replaces such cases with
+    //   `.expectUnchanged()`.)
+    // - Whether to track `.expectUnchanged()` test cases.
     @Override
-    public @Nullable Void visitMethodInvocation(MethodInvocationTree node, @Nullable Void unused) {
+    public @Nullable Void visitMethodInvocation(MethodInvocationTree node, VisitorState state) {
       if (IDENTIFICATION_SOURCE_LINES.matches(node, state)) {
-        identificationTests.add(getSourceLines(node));
-      } else if (REPLACEMENT_INPUT.matches(node, state)) {
-        /* The visitor starts with `addOutputLines` and in the next visit it will go over the `addInputLines`. */
-        replacementTests.add(
-            BugPatternReplacementTestDocumentation.create(
-                getSourceLines(node), replacementOutputLines));
+        getSourceCode(node).ifPresent(identificationTests::add);
       } else if (REPLACEMENT_OUTPUT.matches(node, state)) {
-        replacementOutputLines = getSourceLines(node);
+        ExpressionTree receiver = ASTHelpers.getReceiver(node);
+        // XXX: Make this code nicer.
+        if (REPLACEMENT_INPUT.matches(receiver, state)) {
+          getSourceCode(node)
+              .ifPresent(
+                  output ->
+                      getSourceCode((MethodInvocationTree) receiver)
+                          .ifPresent(
+                              input ->
+                                  replacementTests.add(
+                                      new AutoValue_BugPatternTestExtractor_BugPatternReplacementTestDocumentation(
+                                          input, output))));
+        }
       }
-      return super.visitMethodInvocation(node, unused);
+
+      return super.visitMethodInvocation(node, state);
     }
 
-    // XXX: Duplicate from `ErrorProneTestSourceFormat`, should we move this to `SourceCode` util?
-    private static String getSourceLines(MethodInvocationTree tree) {
+    // XXX: Duplicated from `ErrorProneTestSourceFormat`. Can we do better?
+    private static Optional<String> getSourceCode(MethodInvocationTree tree) {
       List<? extends ExpressionTree> sourceLines =
           tree.getArguments().subList(1, tree.getArguments().size());
       StringBuilder source = new StringBuilder();
 
       for (ExpressionTree sourceLine : sourceLines) {
-        Object value = ASTHelpers.constValue(sourceLine);
+        String value = ASTHelpers.constValue(sourceLine, String.class);
         if (value == null) {
-          return "";
+          return Optional.empty();
         }
         source.append(value).append('\n');
       }
 
-      return source.toString();
+      return Optional.of(source.toString());
     }
   }
 
+  // XXX: Rename?
   @AutoValue
   abstract static class BugPatternTestDocumentation {
     abstract String name();
@@ -165,13 +182,9 @@ abstract static class BugPatternTestDocumentation {
     abstract ImmutableList<BugPatternReplacementTestDocumentation> replacementTests();
   }
 
+  // XXX: Rename?
   @AutoValue
   abstract static class BugPatternReplacementTestDocumentation {
-    static BugPatternReplacementTestDocumentation create(String sourceLines, String outputLines) {
-      return new AutoValue_BugPatternTestExtractor_BugPatternReplacementTestDocumentation(
-          sourceLines, outputLines);
-    }
-
     abstract String inputLines();
 
     abstract String outputLines();

documentation-support/src/main/java/tech/picnic/errorprone/documentation/DocumentationGeneratorTaskListener.java

@@ -7,9 +7,11 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.errorprone.VisitorState;
 import com.sun.source.tree.ClassTree;
+import com.sun.source.tree.CompilationUnitTree;
 import com.sun.source.util.TaskEvent;
 import com.sun.source.util.TaskEvent.Kind;
 import com.sun.source.util.TaskListener;
+import com.sun.source.util.TreePath;
 import com.sun.tools.javac.api.JavacTrees;
 import com.sun.tools.javac.util.Context;
 import java.io.File;
@@ -52,19 +54,24 @@ public void finished(TaskEvent taskEvent) {
       return;
     }
 
-    ClassTree classTree = JavacTrees.instance(context).getTree(taskEvent.getTypeElement());
     JavaFileObject sourceFile = taskEvent.getSourceFile();
-    if (classTree == null || sourceFile == null) {
+    CompilationUnitTree compilationUnit = taskEvent.getCompilationUnit();
+    ClassTree classTree = JavacTrees.instance(context).getTree(taskEvent.getTypeElement());
+    if (sourceFile == null || compilationUnit == null || classTree == null) {
       return;
     }
 
-    ExtractorType.findMatchingType(classTree, VisitorState.createForUtilityPurposes(context))
+    VisitorState state =
+        VisitorState.createForUtilityPurposes(context)
+            .withPath(new TreePath(new TreePath(compilationUnit), classTree));
+
+    ExtractorType.findMatchingType(classTree, state)
         .ifPresent(
             extractorType ->
                 writeToFile(
                     extractorType.getIdentifier(),
                     getSimpleClassName(sourceFile.toUri()),
-                    extractorType.getExtractor().extract(classTree, context)));
+                    extractorType.getExtractor().extract(classTree, state)));
   }
 
   private void createDocsDirectory() {

documentation-support/src/main/java/tech/picnic/errorprone/documentation/Extractor.java

@@ -3,7 +3,6 @@
 import com.google.errorprone.VisitorState;
 import com.google.errorprone.annotations.Immutable;
 import com.sun.source.tree.ClassTree;
-import com.sun.tools.javac.util.Context;
 
 /**
  * Interface implemented by classes that define how to extract data of some type {@link T} from a
@@ -17,18 +16,18 @@ interface Extractor<T> {
    * Extracts and returns an instance of {@link T} using the provided arguments.
    *
    * @param tree The {@link ClassTree} to analyze and from which to extract instances of {@link T}.
-   * @param context The {@link Context} in which the current compilation takes place.
+   * @param state A {@link VisitorState} describing the context in which the given {@link ClassTree}
+   *     is found.
    * @return A non-null instance of {@link T}.
    */
-  // XXX: Drop `Context` parameter unless used.
-  T extract(ClassTree tree, Context context);
+  T extract(ClassTree tree, VisitorState state);
 
   /**
    * Tells whether this {@link Extractor} can extract documentation content from the given {@link
    * ClassTree}.
    *
    * @param tree The {@link ClassTree} of interest.
-   * @param state A {@link VisitorState} describes the context in which the given {@link ClassTree}
+   * @param state A {@link VisitorState} describing the context in which the given {@link ClassTree}
    *     is found.
    * @return {@code true} iff data extraction is supported.
    */

documentation-support/src/test/java/tech/picnic/errorprone/documentation/BugPatternExtractorTest.java

@@ -141,7 +141,7 @@ public static final class TestChecker extends BugChecker implements ClassTreeMat
     public Description matchClass(ClassTree tree, VisitorState state) {
       BugPatternExtractor extractor = new BugPatternExtractor();
 
-      assertThatThrownBy(() -> extractor.extract(tree, state.context))
+      assertThatThrownBy(() -> extractor.extract(tree, state))
           .isInstanceOf(NullPointerException.class)
           .hasMessage("BugPattern annotation must be present");
 

documentation-support/src/test/resources/tech/picnic/errorprone/documentation/bugpattern-test-documentation-identification-and-replacement.json

@@ -9,4 +9,4 @@
       "outputLines": "class A {}\n"
     }
   ]
-}
+}

documentation-support/src/test/resources/tech/picnic/errorprone/documentation/bugpattern-test-documentation-replacement-expect-unchanged.json

@@ -1,10 +1,5 @@
 {
   "name": "TestChecker",
   "identificationTests": [],
-  "replacementTests": [
-    {
-      "inputLines": "class A {}\n",
-      "outputLines": ""
-    }
-  ]
+  "replacementTests": []
 }

error-prone-contrib/src/main/java/tech/picnic/errorprone/bugpatterns/ErrorProneTestHelperSourceFormat.java

@@ -152,7 +152,7 @@ private static Optional<String> getConstantSourceCode(
     StringBuilder source = new StringBuilder();
 
     for (ExpressionTree sourceLine : sourceLines) {
-      Object value = ASTHelpers.constValue(sourceLine);
+      String value = ASTHelpers.constValue(sourceLine, String.class);
       if (value == null) {
         return Optional.empty();
       }