Convert multi param lambdas and local method invocations to method references

baseline-error-prone/src/main/java/com/palantir/baseline/errorprone/LambdaMethodReference.java

@@ -17,6 +17,7 @@
 package com.palantir.baseline.errorprone;
 
 import com.google.auto.service.AutoService;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.google.errorprone.BugPattern;
 import com.google.errorprone.VisitorState;
@@ -27,13 +28,16 @@
 import com.google.errorprone.util.ASTHelpers;
 import com.sun.source.tree.BlockTree;
 import com.sun.source.tree.ExpressionTree;
+import com.sun.source.tree.IdentifierTree;
 import com.sun.source.tree.LambdaExpressionTree;
 import com.sun.source.tree.MethodInvocationTree;
 import com.sun.source.tree.ReturnTree;
 import com.sun.source.tree.StatementTree;
 import com.sun.source.tree.Tree;
 import com.sun.tools.javac.code.Symbol;
 import com.sun.tools.javac.code.Type;
+import java.util.List;
+import java.util.Objects;
 import java.util.Optional;
 
 @AutoService(BugChecker.class)
@@ -50,12 +54,6 @@ public final class LambdaMethodReference extends BugChecker implements BugChecke
 
     @Override
     public Description matchLambdaExpression(LambdaExpressionTree tree, VisitorState state) {
-        // Only handle simple no-arg method references for the time being, don't worry about
-        // simplifying map.forEach((k, v) -> func(k, v)) to map.forEach(this::func)
-        if (tree.getParameters().size() > 1) {
-            return Description.NO_MATCH;
-        }
-
         LambdaExpressionTree.BodyKind bodyKind = tree.getBodyKind();
         Tree body = tree.getBody();
         // n.b. These checks are meant to avoid any and all cleverness. The goal is to be confident
@@ -91,49 +89,93 @@ public Description matchLambdaExpression(LambdaExpressionTree tree, VisitorState
 
     private Description checkMethodInvocation(
             MethodInvocationTree methodInvocation, LambdaExpressionTree root, VisitorState state) {
-        if (!methodInvocation.getArguments().isEmpty()
-                || !methodInvocation.getTypeArguments().isEmpty()) {
+        Symbol.MethodSymbol methodSymbol = ASTHelpers.getSymbol(methodInvocation);
+        if (methodSymbol == null || !methodInvocation.getTypeArguments().isEmpty()) {
             return Description.NO_MATCH;
         }
-        Symbol.MethodSymbol methodSymbol = ASTHelpers.getSymbol(methodInvocation);
-        if (methodSymbol == null || shouldIgnore(methodSymbol, root, methodInvocation)) {
+
+        ExpressionTree receiver = ASTHelpers.getReceiver(methodInvocation);
+        boolean isLocal = isLocal(methodInvocation);
+        if (!isLocal && !(receiver instanceof IdentifierTree)) {
+            return Description.NO_MATCH;
+        }
+
+        if (methodInvocation.getArguments().isEmpty() && root.getParameters().size() == 1) {
+            return convertVariableInstanceMethods(methodSymbol, methodInvocation, root, state);
+        }
+
+        if (methodInvocation.getArguments().size() == root.getParameters().size()) {
+            return convertMethodInvocations(methodSymbol, methodInvocation, root, state);
+        }
+
+        return Description.NO_MATCH;
+    }
+
+    private Description convertVariableInstanceMethods(
+            Symbol.MethodSymbol methodSymbol,
+            MethodInvocationTree methodInvocation,
+            LambdaExpressionTree root,
+            VisitorState state) {
+        Symbol paramSymbol = ASTHelpers.getSymbol(Iterables.getOnlyElement(root.getParameters()));
+        Symbol receiverSymbol = ASTHelpers.getSymbol(ASTHelpers.getReceiver(methodInvocation));
+        if (!paramSymbol.equals(receiverSymbol)) {
             return Description.NO_MATCH;
         }
 
         return buildDescription(root)
                 .setMessage(MESSAGE)
-                .addFix(buildFix(methodSymbol, methodInvocation, root, state))
+                .addFix(buildFix(methodSymbol, methodInvocation, root, state, isLocal(methodInvocation)))
                 .build();
     }
 
-    private static boolean shouldIgnore(
-            Symbol.MethodSymbol methodSymbol, LambdaExpressionTree root, MethodInvocationTree methodInvocation) {
-        if (!methodSymbol.isStatic()) {
-            if (root.getParameters().size() == 1) {
-                Symbol paramSymbol = ASTHelpers.getSymbol(Iterables.getOnlyElement(root.getParameters()));
-                Symbol receiverSymbol = ASTHelpers.getSymbol(ASTHelpers.getReceiver(methodInvocation));
-                return !paramSymbol.equals(receiverSymbol);
-            }
-            return true;
+    private Description convertMethodInvocations(
+            Symbol.MethodSymbol methodSymbol,
+            MethodInvocationTree methodInvocation,
+            LambdaExpressionTree root,
+            VisitorState state) {
+        List<Symbol> methodParams = getSymbols(methodInvocation.getArguments());
+        List<Symbol> lambdaParam = getSymbols(root.getParameters());
+
+        // We are guaranteed that all of root params are symbols so equality should handle cases where methodInvocation
+        // arguments are not symbols or are out of order
+        if (!methodParams.equals(lambdaParam)) {
+            return Description.NO_MATCH;
         }
-        return !root.getParameters().isEmpty();
+
+        return buildDescription(root)
+                .setMessage(MESSAGE)
+                .addFix(buildFix(methodSymbol, methodInvocation, root, state, isLocal(methodInvocation)))
+                .build();
+    }
+
+    private static List<Symbol> getSymbols(List<? extends Tree> params) {
+        return params.stream()
+                .map(ASTHelpers::getSymbol)
+                .filter(Objects::nonNull)
+                .collect(ImmutableList.toImmutableList());
     }
 
     private static Optional<SuggestedFix> buildFix(
             Symbol.MethodSymbol symbol,
             MethodInvocationTree invocation,
             LambdaExpressionTree root,
-            VisitorState state) {
+            VisitorState state,
+            boolean isLocal) {
         SuggestedFix.Builder builder = SuggestedFix.builder();
-        return toMethodReference(qualifyTarget(symbol, invocation, builder, state))
+        return toMethodReference(qualifyTarget(symbol, invocation, builder, state, isLocal))
                 .map(qualified -> builder.replace(root, qualified).build());
     }
 
     private static String qualifyTarget(
             Symbol.MethodSymbol symbol,
             MethodInvocationTree invocation,
             SuggestedFix.Builder builder,
-            VisitorState state) {
+            VisitorState state,
+            boolean isLocal) {
+        if (!symbol.isStatic() && isLocal) {
+            return "this." + symbol.name.toString();
+        }
+
         Type receiverType = ASTHelpers.getReceiverType(invocation);
         if (receiverType == null || receiverType.getLowerBound() != null || receiverType.getUpperBound() != null) {
             return SuggestedFixes.qualifyType(state, builder, symbol);
@@ -151,4 +193,11 @@ private static Optional<String> toMethodReference(String qualifiedMethodName) {
         }
         return Optional.empty();
     }
+
+    private static boolean isLocal(MethodInvocationTree methodInvocationTree) {
+        ExpressionTree receiver = ASTHelpers.getReceiver(methodInvocationTree);
+        return receiver == null
+                || (receiver instanceof IdentifierTree
+                        && "this".equals(((IdentifierTree) receiver).getName().toString()));
+    }
 }