Fix false positive "Collection count can be converted to size" with Iterable

#KT-34677 Fixed

idea/src/org/jetbrains/kotlin/idea/inspections/ReplaceCollectionCountWithSizeInspection.kt

@@ -11,23 +11,30 @@ import com.intellij.codeInspection.ProblemsHolder
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElementVisitor
 import org.jetbrains.kotlin.idea.KotlinBundle
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.inspections.collections.isCalling
+import org.jetbrains.kotlin.idea.inspections.collections.receiverType
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.KtCallExpression
 import org.jetbrains.kotlin.psi.KtPsiFactory
 import org.jetbrains.kotlin.psi.callExpressionVisitor
+import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 
 class ReplaceCollectionCountWithSizeInspection : AbstractKotlinInspection() {
     override fun buildVisitor(holder: ProblemsHolder, isOnTheFly: Boolean): PsiElementVisitor {
-        return callExpressionVisitor { callExpression ->
+        return callExpressionVisitor(fun(callExpression: KtCallExpression) {
-            if (callExpression.isCount()) {
+            if (callExpression.calleeExpression?.text != "count" || callExpression.valueArguments.isNotEmpty()) return
-                holder.registerProblem(
+            val context = callExpression.analyze(BodyResolveMode.PARTIAL)
-                    callExpression,
+            if (!callExpression.isCalling(FqName("kotlin.collections.count"))) return
-                    KotlinBundle.message("could.be.replaced.with.size"),
+            val receiverType = callExpression.receiverType(context) ?: return
-                    ReplaceCollectionCountWithSizeQuickFix()
+            if (KotlinBuiltIns.isIterableOrNullableIterable(receiverType)) return
-                )
+            holder.registerProblem(
-            }
+                callExpression,
-        }
+                KotlinBundle.message("could.be.replaced.with.size"),
+                ReplaceCollectionCountWithSizeQuickFix()
+            )
+        })
     }
 }
 
@@ -41,5 +48,3 @@ class ReplaceCollectionCountWithSizeQuickFix : LocalQuickFix {
         element.replace(KtPsiFactory(element).createExpression("size"))
     }
 }
-
-private fun KtCallExpression.isCount(): Boolean = valueArguments.isEmpty() && isCalling(FqName("kotlin.collections.count"))

idea/src/org/jetbrains/kotlin/idea/inspections/collections/ConvertCallChainIntoSequenceInspection.kt

@@ -148,9 +148,7 @@ private fun KtQualifiedExpression.findCallChain(): CallChain? {
     if (calls.isEmpty()) return null
 
     val lastCall = calls.last()
-    val receiverType =
+    val receiverType = lastCall.receiverType(context)
-        (lastCall.getQualifiedExpressionForSelector())?.receiverExpression?.getResolvedCall(context)?.resultingDescriptor?.returnType
-            ?: lastCall.implicitReceiver(context)?.type
     if (receiverType?.isIterable(DefaultBuiltIns.Instance) != true) return null
 
     val firstCall = calls.first()
@@ -186,10 +184,6 @@ private fun KtQualifiedExpression.collectCallExpression(context: BindingContext)
     return transformationCalls
 }
 
-private fun KtExpression.implicitReceiver(context: BindingContext): ImplicitReceiver? {
-    return getResolvedCall(context)?.getImplicitReceiverValue()
-}
-
 private fun KtCallExpression.hasReturn(): Boolean = valueArguments.any { arg ->
     arg.anyDescendantOfType<KtReturnExpression> { it.labelQualifier == null }
 }

idea/src/org/jetbrains/kotlin/idea/inspections/collections/FunctionUtils.kt

@@ -23,17 +23,20 @@ import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.KtCallExpression
+import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForSelector
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.callUtil.getType
 import org.jetbrains.kotlin.resolve.calls.model.ReceiverKotlinCallArgument
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedLambdaAtom
+import org.jetbrains.kotlin.resolve.calls.resolvedCallUtil.getImplicitReceiverValue
 import org.jetbrains.kotlin.resolve.calls.tower.NewResolvedCallImpl
 import org.jetbrains.kotlin.resolve.calls.tower.receiverValue
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import org.jetbrains.kotlin.resolve.descriptorUtil.isSubclassOf
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
+import org.jetbrains.kotlin.resolve.scopes.receivers.ImplicitReceiver
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
@@ -86,3 +89,12 @@ fun ResolvedCall<*>.hasLastFunctionalParameterWithResult(context: BindingContext
     val resultType = functionalType.arguments.lastOrNull()?.type ?: return false
     return predicate(resultType)
 }
+
+fun KtCallExpression.implicitReceiver(context: BindingContext): ImplicitReceiver? {
+    return getResolvedCall(context)?.getImplicitReceiverValue()
+}
+
+fun KtCallExpression.receiverType(context: BindingContext): KotlinType? {
+    return (getQualifiedExpressionForSelector())?.receiverExpression?.getResolvedCall(context)?.resultingDescriptor?.returnType
+        ?: implicitReceiver(context)?.type
+}

idea/testData/inspectionsLocal/replaceCollectionCountWithSize/countOfIterable.kt

@@ -0,0 +1,6 @@
+// PROBLEM: none
+// WITH_RUNTIME
+
+fun foo(iterable: Iterable<String>) {
+    iterable.<caret>count()
+}

idea/testData/inspectionsLocal/replaceCollectionCountWithSize/countOfIterable2.kt

@@ -0,0 +1,6 @@
+// PROBLEM: none
+// WITH_RUNTIME
+
+fun foo(iterable: Iterable<String>?) {
+    iterable?.<caret>count()
+}

idea/testData/inspectionsLocal/replaceCollectionCountWithSize/countOfIterable3.kt

@@ -0,0 +1,8 @@
+// PROBLEM: none
+// WITH_RUNTIME
+
+fun foo(iterable: Iterable<String>) {
+    iterable.run {
+        <caret>count()
+    }
+}

idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java

@@ -9849,6 +9849,21 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
             runTest("idea/testData/inspectionsLocal/replaceCollectionCountWithSize/countOfCollection.kt");
         }
 
+        @TestMetadata("countOfIterable.kt")
+        public void testCountOfIterable() throws Exception {
+            runTest("idea/testData/inspectionsLocal/replaceCollectionCountWithSize/countOfIterable.kt");
+        }
+
+        @TestMetadata("countOfIterable2.kt")
+        public void testCountOfIterable2() throws Exception {
+            runTest("idea/testData/inspectionsLocal/replaceCollectionCountWithSize/countOfIterable2.kt");
+        }
+
+        @TestMetadata("countOfIterable3.kt")
+        public void testCountOfIterable3() throws Exception {
+            runTest("idea/testData/inspectionsLocal/replaceCollectionCountWithSize/countOfIterable3.kt");
+        }
+
         @TestMetadata("countOfMap.kt")
         public void testCountOfMap() throws Exception {
             runTest("idea/testData/inspectionsLocal/replaceCollectionCountWithSize/countOfMap.kt");