From 0b48416a1eeaf5c2c95716e109b7495da36a506b Mon Sep 17 00:00:00 2001
From: Roman Golyshev <roman.golyshev@jetbrains.com>
Date: Wed, 23 Dec 2020 13:45:05 +0300
Subject: [PATCH] FIR IDE: Unwrap nullable types

---
 .../shortenRefs/FirShortenRefsTestGenerated.java      |  5 +++++
 .../api/fir/components/KtFirReferenceShortener.kt     | 11 +++++++++--
 .../shortenRefsFir/types/ParameterTypeNullableType.kt |  6 ++++++
 .../types/ParameterTypeNullableType.kt.after          |  6 ++++++
 4 files changed, 26 insertions(+), 2 deletions(-)
 create mode 100644 idea/testData/shortenRefsFir/types/ParameterTypeNullableType.kt
 create mode 100644 idea/testData/shortenRefsFir/types/ParameterTypeNullableType.kt.after

diff --git a/idea/idea-fir/tests/org/jetbrains/kotlin/shortenRefs/FirShortenRefsTestGenerated.java b/idea/idea-fir/tests/org/jetbrains/kotlin/shortenRefs/FirShortenRefsTestGenerated.java
index 728f2a8a6ff..33a57420cff 100644
--- a/idea/idea-fir/tests/org/jetbrains/kotlin/shortenRefs/FirShortenRefsTestGenerated.java
+++ b/idea/idea-fir/tests/org/jetbrains/kotlin/shortenRefs/FirShortenRefsTestGenerated.java
@@ -96,6 +96,11 @@ public class FirShortenRefsTestGenerated extends AbstractFirShortenRefsTest {
             runTest("idea/testData/shortenRefsFir/types/ParameterTypeNotImportedNestedClass.kt");
         }
 
+        @TestMetadata("ParameterTypeNullableType.kt")
+        public void testParameterTypeNullableType() throws Exception {
+            runTest("idea/testData/shortenRefsFir/types/ParameterTypeNullableType.kt");
+        }
+
         @TestMetadata("ParameterTypeStarImportedTypeLoses.kt")
         public void testParameterTypeStarImportedTypeLoses() throws Exception {
             runTest("idea/testData/shortenRefsFir/types/ParameterTypeStarImportedTypeLoses.kt");
diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirReferenceShortener.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirReferenceShortener.kt
index 17b2cfb1597..de77008cae5 100644
--- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirReferenceShortener.kt
+++ b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirReferenceShortener.kt
@@ -53,7 +53,11 @@ internal class KtFirReferenceShortener(
 
         firFile.acceptChildren(TypesCollectingVisitor(typesToImport, typesToShorten))
 
-        return ShortenCommandImpl(file, typesToImport, typesToShorten.map { it.createSmartPointer() })
+        return ShortenCommandImpl(
+            file,
+            typesToImport.distinct(),
+            typesToShorten.distinct().map { it.createSmartPointer() }
+        )
     }
 
     private fun findFirstClassifierInScopesByName(positionScopes: List<FirScope>, targetClassName: Name): ClassId? {
@@ -134,7 +138,7 @@ internal class KtFirReferenceShortener(
             val wholeTypeReference = resolvedTypeRef.psi as? KtTypeReference ?: return
 
             val wholeClassifierId = resolvedTypeRef.type.classId ?: return
-            val wholeTypeElement = wholeTypeReference.typeElement as? KtUserType ?: return
+            val wholeTypeElement = wholeTypeReference.typeElement.unwrapNullable() as? KtUserType ?: return
 
             if (wholeTypeElement.qualifier == null) return
 
@@ -197,3 +201,6 @@ private class ShortenCommandImpl(
         }
     }
 }
+
+private tailrec fun KtTypeElement?.unwrapNullable(): KtTypeElement? =
+    if (this is KtNullableType) this.innerType.unwrapNullable() else this
diff --git a/idea/testData/shortenRefsFir/types/ParameterTypeNullableType.kt b/idea/testData/shortenRefsFir/types/ParameterTypeNullableType.kt
new file mode 100644
index 00000000000..817a9e3732c
--- /dev/null
+++ b/idea/testData/shortenRefsFir/types/ParameterTypeNullableType.kt
@@ -0,0 +1,6 @@
+// FIR_COMPARISON
+package test
+
+class T
+
+<selection>fun foo(t: test.T???) {}</selection>
diff --git a/idea/testData/shortenRefsFir/types/ParameterTypeNullableType.kt.after b/idea/testData/shortenRefsFir/types/ParameterTypeNullableType.kt.after
new file mode 100644
index 00000000000..e1b4d75b9c9
--- /dev/null
+++ b/idea/testData/shortenRefsFir/types/ParameterTypeNullableType.kt.after
@@ -0,0 +1,6 @@
+// FIR_COMPARISON
+package test
+
+class T
+
+fun foo(t: T???) {}