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;
@@ -25,16 +26,26 @@
 import com.google.errorprone.fixes.SuggestedFixes;
 import com.google.errorprone.matchers.Description;
 import com.google.errorprone.util.ASTHelpers;
+import com.google.errorprone.util.ErrorProneToken;
 import com.sun.source.tree.BlockTree;
+import com.sun.source.tree.ClassTree;
 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.source.tree.VariableTree;
+import com.sun.tools.javac.code.Flags;
 import com.sun.tools.javac.code.Symbol;
 import com.sun.tools.javac.code.Type;
+import com.sun.tools.javac.parser.Tokens;
+import java.util.List;
+import java.util.Objects;
 import java.util.Optional;
+import java.util.Set;
+import javax.annotation.Nullable;
 
 @AutoService(BugChecker.class)
 @BugPattern(
@@ -44,18 +55,13 @@
         providesFix = BugPattern.ProvidesFix.REQUIRES_HUMAN_ATTENTION,
         severity = BugPattern.SeverityLevel.SUGGESTION,
         summary = "Lambda should be a method reference")
+@SuppressWarnings("checkstyle:CyclomaticComplexity")
 public final class LambdaMethodReference extends BugChecker implements BugChecker.LambdaExpressionTreeMatcher {
 
     private static final String MESSAGE = "Lambda should be a method reference";
 
     @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,56 +97,190 @@ 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()
+                || hasExplicitParameterTypes(root, state)) {
             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;
         }
 
-        return buildDescription(root)
-                .setMessage(MESSAGE)
-                .addFix(buildFix(methodSymbol, methodInvocation, root, state))
-                .build();
+        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 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);
+    private static boolean hasExplicitParameterTypes(LambdaExpressionTree lambda, VisitorState state) {
+        for (VariableTree varTree : lambda.getParameters()) {
+            boolean expectComma = false;
+            // Must avoid refactoring lambdas which declare explicit parameter types
+            for (ErrorProneToken token : state.getTokensForNode(varTree)) {
+                if (token.kind() == Tokens.TokenKind.EOF) {
+                    return false;
+                } else if ((token.kind() == Tokens.TokenKind.IDENTIFIER && expectComma)
+                        || (token.kind() == Tokens.TokenKind.COMMA && !expectComma)) {
+                    return true;
+                }
+                expectComma = !expectComma;
             }
-            return true;
         }
-        return !root.getParameters().isEmpty();
+        return false;
+    }
+
+    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 buildFix(methodSymbol, methodInvocation, root, state, isLocal(methodInvocation))
+                .map(fix ->
+                        buildDescription(root).setMessage(MESSAGE).addFix(fix).build())
+                .orElse(Description.NO_MATCH);
+    }
+
+    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 buildFix(methodSymbol, methodInvocation, root, state, isLocal(methodInvocation))
+                .map(fix ->
+                        buildDescription(root).setMessage(MESSAGE).addFix(fix).build())
+                .orElse(Description.NO_MATCH);
+    }
+
+    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) {
+        if (isAmbiguousMethod(symbol, ASTHelpers.getReceiver(invocation), state)) {
+            return Optional.empty();
+        }
+
         SuggestedFix.Builder builder = SuggestedFix.builder();
-        return toMethodReference(qualifyTarget(symbol, invocation, builder, state))
+        return qualifyTarget(symbol, invocation, root, builder, state, isLocal)
+                .flatMap(LambdaMethodReference::toMethodReference)
                 .map(qualified -> builder.replace(root, qualified).build());
     }
 
-    private static String qualifyTarget(
+    private static boolean isAmbiguousMethod(
+            Symbol.MethodSymbol symbol, @Nullable ExpressionTree receiver, VisitorState state) {
+        if (symbol.isStatic()) {
+            if (symbol.params().size() != 1) {
+                return false;
+            }
+            Symbol.ClassSymbol classSymbol = ASTHelpers.enclosingClass(symbol);
+            if (classSymbol == null) {
+                return false;
+            }
+            Set<Symbol.MethodSymbol> matching = ASTHelpers.findMatchingMethods(
+                    symbol.name,
+                    sym -> sym != null && !sym.isStatic() && sym.getParameters().isEmpty(),
+                    classSymbol.type,
+                    state.getTypes());
+            return !matching.isEmpty();
+        } else {
+            if (!symbol.params().isEmpty()) {
+                return false;
+            }
+            if (receiver == null) {
+                return false;
+            }
+            Type receiverType = ASTHelpers.getType(receiver);
+            if (receiverType == null) {
+                return false;
+            }
+            Set<Symbol.MethodSymbol> matching = ASTHelpers.findMatchingMethods(
+                    symbol.name,
+                    sym -> sym != null
+                            && sym.isStatic()
+                            && sym.getParameters().size() == 1
+                            && state.getTypes()
+                                    .isAssignable(
+                                            state.getTypes().erasure(receiverType),
+                                            state.getTypes()
+                                                    .erasure(sym.params().get(0).type)),
+                    receiverType,
+                    state.getTypes());
+            return !matching.isEmpty();
+        }
+    }
+
+    private static Optional<String> qualifyTarget(
             Symbol.MethodSymbol symbol,
             MethodInvocationTree invocation,
+            LambdaExpressionTree root,
             SuggestedFix.Builder builder,
-            VisitorState state) {
+            VisitorState state,
+            boolean isLocal) {
+        if (!symbol.isStatic() && isLocal) {
+            // Validate teh local method is defined in this class
+            ClassTree enclosingClass = ASTHelpers.findEnclosingNode(state.getPath(), ClassTree.class);
+            if (enclosingClass == null) {
+                return Optional.empty();
+            }
+            Type.ClassType enclosingType = ASTHelpers.getType(enclosingClass);
+            if (!ASTHelpers.findMatchingMethods(symbol.name, symbol::equals, enclosingType, state.getTypes())
+                    .isEmpty()) {
+                return Optional.of("this." + symbol.name.toString());
+            }
+            return Optional.empty();
+        }
+
+        ExpressionTree receiver = ASTHelpers.getReceiver(invocation);
         Type receiverType = ASTHelpers.getReceiverType(invocation);
         if (receiverType == null || receiverType.getLowerBound() != null || receiverType.getUpperBound() != null) {
-            return SuggestedFixes.qualifyType(state, builder, symbol);
+            return Optional.of(SuggestedFixes.qualifyType(state, builder, symbol));
+        }
+        Symbol receiverSymbol = ASTHelpers.getSymbol(receiver);
+        if (!symbol.isStatic()
+                && receiver instanceof IdentifierTree
+                && !Objects.equals(ImmutableList.of(receiverSymbol), getSymbols(root.getParameters()))) {
+            if (!isFinal(receiverSymbol)) {
+                // Not safe to replace lambdas which lazily reference values with an eager capture.
+                return Optional.empty();
+            }
+            return Optional.of(state.getSourceForNode(receiver) + '.' + symbol.name.toString());
         }
-        return SuggestedFixes.qualifyType(state, builder, state.getTypes().erasure(receiverType))
-                + '.'
-                + symbol.name.toString();
+
+        return Optional.of(
+                SuggestedFixes.qualifyType(state, builder, state.getTypes().erasure(receiverType))
+                        + '.'
+                        + symbol.name.toString());
     }
 
     private static Optional<String> toMethodReference(String qualifiedMethodName) {
@@ -151,4 +291,15 @@ 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()));
+    }
+
+    private static boolean isFinal(Symbol symbol) {
+        return (symbol.flags() & (Flags.FINAL | Flags.EFFECTIVELY_FINAL)) != 0;
+    }
 }