From eb901e4f32e9886550fd7d249da5a0b4071347c2 Mon Sep 17 00:00:00 2001
From: Tianyu Geng <tgeng@google.com>
Date: Fri, 1 Oct 2021 15:42:16 -0700
Subject: [PATCH] FIR HL API: add helper for shortening ref on EDT thread

---
 .../api/components/KtReferenceShortener.kt    | 40 +++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtReferenceShortener.kt b/analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtReferenceShortener.kt
index 20a973aa695..603b9db8673 100644
--- a/analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtReferenceShortener.kt
+++ b/analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/components/KtReferenceShortener.kt
@@ -5,10 +5,14 @@
 
 package org.jetbrains.kotlin.analysis.api.components
 
+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.util.TextRange
+import org.jetbrains.kotlin.analysis.api.analyse
 import org.jetbrains.kotlin.analysis.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtClassLikeSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtEnumEntrySymbol
+import org.jetbrains.kotlin.analysis.api.tokens.HackToForceAllowRunningAnalyzeOnEDT
+import org.jetbrains.kotlin.analysis.api.tokens.hackyAllowRunningOnEdt
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 
@@ -49,6 +53,42 @@ public interface KtReferenceShortenerMixIn : KtAnalysisSessionMixIn {
             if (symbol is KtEnumEntrySymbol) ShortenOption.DO_NOT_SHORTEN
             else ShortenOption.SHORTEN_AND_IMPORT
         }
+
+        /**
+         * Shorten references in the given [element]. See [shortenReferencesInRange] for more details.
+         */
+        public fun shortenReferences(
+            element: KtElement,
+            classShortenOption: (KtClassLikeSymbol) -> ShortenOption = defaultClassShortenOption,
+            callableShortenOption: (KtCallableSymbol) -> ShortenOption = defaultCallableShortenOption
+        ): Unit = shortenReferencesInRange(
+            element.containingKtFile,
+            element.textRange,
+            classShortenOption,
+            callableShortenOption
+        )
+
+        /**
+         * Shorten references in the given [file] and [range]. The function must be invoked on EDT thread because it modifies the underlying
+         * PSI. This method analyse Kotlin code and hence could block the EDT thread for longer period of time. Hence, this method should be
+         * called only to shorten references in *newly generated code* by IDE actions. In other cases, please consider using
+         * [org.jetbrains.kotlin.analysis.api.components.KtReferenceShortenerMixIn] in a background thread to perform the analysis and then
+         * modify PSI on the EDt thread by invoking [org.jetbrains.kotlin.analysis.api.components.ShortenCommand.invokeShortening]. */
+        @OptIn(HackToForceAllowRunningAnalyzeOnEDT::class)
+        public fun shortenReferencesInRange(
+            file: KtFile,
+            range: TextRange = file.textRange,
+            classShortenOption: (KtClassLikeSymbol) -> ShortenOption = defaultClassShortenOption,
+            callableShortenOption: (KtCallableSymbol) -> ShortenOption = defaultCallableShortenOption
+        ) {
+            ApplicationManager.getApplication().assertIsDispatchThread()
+            val shortenCommand = hackyAllowRunningOnEdt {
+                analyse(file) {
+                    collectPossibleReferenceShortenings(file, range, classShortenOption, callableShortenOption)
+                }
+            }
+            shortenCommand.invokeShortening()
+        }
     }
 
     /**