KT-17608 Kotlin plugin says missing apply() call for SharedPreferences.Editor in With expression (#1164)

idea/testData/android/lint/sharedPrefs.kt

@@ -29,6 +29,35 @@ class SharedPrefsText(context: Context) : Activity() {
         }
     }
 
+    // OK using with lambda
+    fun withLambda() {
+        val preferences = PreferenceManager.getDefaultSharedPreferences(this)
+        with(preferences.edit()) {
+            putString("foo", "bar")
+            putInt("bar", 42)
+            apply()
+        }
+    }
+
+    // OK using apply lambda
+    fun testApplyLambda() {
+        PreferenceManager.getDefaultSharedPreferences(this).edit().apply {
+            putString("foo", "bar")
+            putInt("bar", 42)
+            apply()
+        }
+    }
+
+    // OK using also lambda
+    fun testAlsoLambda() {
+        val preferences = PreferenceManager.getDefaultSharedPreferences(this)
+        preferences.edit().also {
+            it.putString("foo", "bar")
+            it.putInt("bar", 42)
+            it.apply()
+        }
+    }
+
     // Not a bug
     fun test(foo: Foo) {
         val bar1 = foo.edit()
@@ -61,6 +90,30 @@ class SharedPrefsText(context: Context) : Activity() {
         editor.putInt("bar", 42)
     }
 
+    // Bug missing commit in apply lambda
+    fun applyLambdaMissingCommit() {
+        PreferenceManager.getDefaultSharedPreferences(this).<warning descr="`SharedPreferences.edit()` without a corresponding `commit()` or `apply()` call">edit()</warning>.apply {
+            putString("foo", "bar")
+            putInt("bar", 42)
+        }
+    }
+
+    // Bug missing commit in also lambda
+    fun alsoLambdaMissingCommit() {
+        PreferenceManager.getDefaultSharedPreferences(this).<warning descr="`SharedPreferences.edit()` without a corresponding `commit()` or `apply()` call">edit()</warning>.also {
+            it.putString("foo", "bar")
+            it.putInt("bar", 42)
+        }
+    }
+
+    // Bug missing commit in with lambda
+    fun withLambdaMissingCommit() {
+        with(PreferenceManager.getDefaultSharedPreferences(this).<warning descr="`SharedPreferences.edit()` without a corresponding `commit()` or `apply()` call">edit()</warning>) {
+            putString("foo", "bar")
+            putInt("bar", 42)
+        }
+    }
+
     init {
         val preferences = PreferenceManager.getDefaultSharedPreferences(context)
         val editor = preferences.<warning descr="`SharedPreferences.edit()` without a corresponding `commit()` or `apply()` call"><warning descr="`SharedPreferences.edit()` without a corresponding `commit()` or `apply()` call">edit()</warning></warning>

plugins/lint/lint-checks/src/com/android/tools/klint/checks/CleanupDetector.java

@@ -370,11 +370,50 @@ public class CleanupDetector extends Detector implements Detector.UastScanner {
                 if (isCleanupCallPredicate.apply(call)) {
                     return true;
                 }
+
+                //If function taking lambda
+                for (UExpression argument : call.getValueArguments()) {
+                    if (argument instanceof ULambdaExpression) {
+                        if (isCleanedUpInLambda((ULambdaExpression) argument, isCleanupCallPredicate)) {
+                            return true;
+                        }
+                    }
+                }
             }
         }
         return false;
     }
 
+    private static boolean isCleanedUpInLambda(ULambdaExpression lambda, Predicate<UCallExpression> isCleanupCallPredicate) {
+        UExpression body = lambda.getBody();
+        if (body instanceof UBlockExpression) {
+            List<UExpression> expressions = ((UBlockExpression) body).getExpressions();
+            for (UExpression expression : expressions) {
+                //Might be a single cleanup call
+                if (expression instanceof UCallExpression) {
+                    if (isCleanupCallPredicate.apply((UCallExpression) expression)) {
+                        return true;
+                    }
+                } else {
+                    //Might be a chain containing the cleanup call
+                    List<UExpression> chain = getQualifiedChain(getOutermostQualified(expression));
+                    if (!chain.isEmpty()) {
+                        for (UExpression e : chain) {
+                            if (e instanceof UCallExpression) {
+                                UCallExpression call = (UCallExpression) e;
+                                if (isCleanupCallPredicate.apply(call)) {
+                                    return true;
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        return false;
+    }
+
     private static void checkRecycled(@NonNull final JavaContext context, @NonNull UCallExpression node,
             @NonNull final String recycleType, @NonNull final String recycleName) {
 
@@ -555,11 +594,11 @@ public class CleanupDetector extends Detector implements Detector.UastScanner {
     private static void checkEditorApplied(@NonNull JavaContext context,
             @NonNull UCallExpression node, @NonNull PsiMethod calledMethod) {
         if (isSharedEditorCreation(context, calledMethod)) {
-            PsiVariable boundVariable = getVariableElement(node, true);
             if (isEditorCommittedInChainedCalls(context, node)) {
                 return;
             }
 
+            PsiVariable boundVariable = getVariableElement(node, true);
             if (boundVariable != null) {
                 UMethod method = getParentOfType(node, UMethod.class, true);
                 if (method == null) {
@@ -612,6 +651,8 @@ public class CleanupDetector extends Detector implements Detector.UastScanner {
             } else if (UastUtils.getParentOfType(node, UReturnExpression.class) != null) {
                 // Allocation is in a return statement
                 return;
+            } else if (isEditorCommittedInParentLambda(context, node)) {
+                return;
             }
 
             String message = "`SharedPreferences.edit()` without a corresponding `commit()` or "
@@ -643,6 +684,26 @@ public class CleanupDetector extends Detector implements Detector.UastScanner {
         });
     }
 
+    private static boolean isEditorCommittedInParentLambda(@NonNull JavaContext context, @NonNull UCallExpression node) {
+        UCallExpression call = UastUtils.getParentOfType(node, UCallExpression.class, true);
+        if (call != null) {
+            //If function taking lambda
+            for (UExpression argument : call.getValueArguments()) {
+                if (argument instanceof ULambdaExpression) {
+                    if (isCleanedUpInLambda((ULambdaExpression) argument, new Predicate<UCallExpression>() {
+                        @Override
+                        public boolean apply(@org.jetbrains.annotations.Nullable UCallExpression call) {
+                            return isEditorCommitMethodCall(context, call) || isEditorApplyMethodCall(context, call);
+                        }
+                    })) {
+                        return true;
+                    }
+                }
+            }
+        }
+        return false;
+    }
+
     private static boolean isEditorCommitMethodCall(@NonNull JavaContext context,
             @NonNull UCallExpression call) {
         String methodName = call.getMethodName();