From 1d9247ae0f8cdbbd31a9fd8839d70a1c04987ef0 Mon Sep 17 00:00:00 2001
From: Mark Punzalan <markpun@google.com>
Date: Sat, 1 May 2021 05:33:45 +0000
Subject: [PATCH] FIR IDE: Offer AddExclExclCallFix for nullable types with an
 `iterator()` function that does NOT have `operator` modifier.

This is different from FE1.0. Adding `!!` will then surface the error
that `operator` modifier needs to be added (with corresponding fix).
---
 .../idea/quickfix/fixes/AddExclExclCallFixFactories.kt |  6 ++++--
 .../addExclExclCall/iterableWithBadIterator.fir.kt     | 10 ++++++++++
 .../iterableWithBadIterator.fir.kt.after               | 10 ++++++++++
 3 files changed, 24 insertions(+), 2 deletions(-)
 create mode 100644 idea/testData/quickfix/addExclExclCall/iterableWithBadIterator.fir.kt
 create mode 100644 idea/testData/quickfix/addExclExclCall/iterableWithBadIterator.fir.kt.after

diff --git a/idea/idea-fir/src/org/jetbrains/kotlin/idea/quickfix/fixes/AddExclExclCallFixFactories.kt b/idea/idea-fir/src/org/jetbrains/kotlin/idea/quickfix/fixes/AddExclExclCallFixFactories.kt
index 9407da4e99d..1777a5fe69f 100644
--- a/idea/idea-fir/src/org/jetbrains/kotlin/idea/quickfix/fixes/AddExclExclCallFixFactories.kt
+++ b/idea/idea-fir/src/org/jetbrains/kotlin/idea/quickfix/fixes/AddExclExclCallFixFactories.kt
@@ -81,9 +81,11 @@ object AddExclExclCallFixFactories {
         val type = expression.getKtType()
         if (!type.canBeNull) return@diagnosticFixFactory emptyList()
 
+        // NOTE: This is different from FE1.0 in that we offer the fix even if the function does NOT have the `operator` modifier.
+        // Adding `!!` will then surface the error that `operator` should be added (with corresponding fix).
         val typeScope = type.getTypeScope() ?: return@diagnosticFixFactory emptyList()
         val hasValidIterator = typeScope.getCallableSymbols { it == OperatorNameConventions.ITERATOR }
-            .filter { it is KtFunctionSymbol && it.isOperator && it.valueParameters.isEmpty() }.singleOrNull() != null
+            .filter { it is KtFunctionSymbol && it.valueParameters.isEmpty() }.singleOrNull() != null
         if (hasValidIterator) {
             listOfNotNull(expression.asAddExclExclCallFix())
         } else {
@@ -93,4 +95,4 @@ object AddExclExclCallFixFactories {
 }
 
 internal fun PsiElement?.asAddExclExclCallFix(hasImplicitReceiver: Boolean = false) =
-    this?.let { AddExclExclCallFix(it, hasImplicitReceiver) } ?: null
+    this?.let { AddExclExclCallFix(it, hasImplicitReceiver) }
diff --git a/idea/testData/quickfix/addExclExclCall/iterableWithBadIterator.fir.kt b/idea/testData/quickfix/addExclExclCall/iterableWithBadIterator.fir.kt
new file mode 100644
index 00000000000..746ec38c6b2
--- /dev/null
+++ b/idea/testData/quickfix/addExclExclCall/iterableWithBadIterator.fir.kt
@@ -0,0 +1,10 @@
+// "Add non-null asserted (!!) call" "true"
+
+class Some {
+    fun iterator(): Iterator<Int> = null!!
+}
+
+fun foo() {
+    val test: Some? = Some()
+    for (i in <caret>test) { }
+}
diff --git a/idea/testData/quickfix/addExclExclCall/iterableWithBadIterator.fir.kt.after b/idea/testData/quickfix/addExclExclCall/iterableWithBadIterator.fir.kt.after
new file mode 100644
index 00000000000..02eda3ed5b3
--- /dev/null
+++ b/idea/testData/quickfix/addExclExclCall/iterableWithBadIterator.fir.kt.after
@@ -0,0 +1,10 @@
+// "Add non-null asserted (!!) call" "true"
+
+class Some {
+    fun iterator(): Iterator<Int> = null!!
+}
+
+fun foo() {
+    val test: Some? = Some()
+    for (i in test!!) { }
+}