[uast] Add API for importing members on demand

GitOrigin-RevId: ab6ed31facfdb0df4204d9a1de199afd13073ded

java/java-impl/intellij.java.impl.iml

@@ -45,7 +45,6 @@
     <orderEntry type="module" module-name="intellij.java.structureView" exported="" />
     <orderEntry type="module" module-name="intellij.spellchecker" />
     <orderEntry type="module" module-name="intellij.platform.builtInServer.impl" />
-    <orderEntry type="module" module-name="intellij.java.uast" />
     <orderEntry type="library" name="gson" level="project" />
     <orderEntry type="module" module-name="intellij.java.compiler" />
     <orderEntry type="library" scope="RUNTIME" name="jaxb-runtime" level="project" />

plugins/kotlin/uast/uast-kotlin-idea/src/org/jetbrains/uast/kotlin/generate/KotlinUastCodeGenerationPlugin.kt

@@ -8,11 +8,13 @@ import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.text.StringUtil
 import com.intellij.psi.*
 import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.util.castSafelyTo
 import org.jetbrains.kotlin.idea.KotlinLanguage
 import org.jetbrains.kotlin.idea.base.fe10.codeInsight.newDeclaration.Fe10KotlinNameSuggester
 import org.jetbrains.kotlin.idea.base.psi.replaced
 import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
 import org.jetbrains.kotlin.idea.core.*
+import org.jetbrains.kotlin.idea.intentions.ImportAllMembersIntention
 import org.jetbrains.kotlin.idea.refactoring.fqName.fqName
 import org.jetbrains.kotlin.idea.references.KtSimpleNameReference
 import org.jetbrains.kotlin.references.fe10.KtFe10SimpleNameReference
@@ -87,6 +89,14 @@ class KotlinUastCodeGenerationPlugin : UastCodeGenerationPlugin {
         if (sourcePsi !is KtElement) return null
         return ShortenReferences.DEFAULT.process(sourcePsi).toUElementOfType()
     }
+
+    override fun importMemberOnDemand(reference: UQualifiedReferenceExpression): UExpression? {
+        val ktQualifiedExpression = reference.sourcePsi?.castSafelyTo<KtDotQualifiedExpression>() ?: return null
+        val selector = ktQualifiedExpression.selectorExpression ?: return null
+        val ptr = SmartPointerManager.createPointer(selector)
+        ImportAllMembersIntention().applyTo(ktQualifiedExpression, null)
+        return ptr.element?.toUElementOfType()
+    }
 }
 
 private fun hasBraces(oldPsi: KtBlockExpression): Boolean = oldPsi.lBrace != null && oldPsi.rBrace != null

uast/uast-common/src/org/jetbrains/uast/generate/UastCodeGenerationPlugin.kt

@@ -55,10 +55,40 @@ interface UastCodeGenerationPlugin {
    * Replaces fully-qualified class names in the contents of the specified element with
    * non-qualified names and adds import statements as necessary.
    *
+   * Example:
+   * ```
+   * com.jetbrains.uast.generate.UastCodeGenerationPlugin.byLanguage(...)
+   * ```
+   * Becomes:
+   * ```
+   * import com.jetbrains.uast.generate.UastCodeGenerationPlugin
+   *
+   * UastCodeGenerationPlugin.byLanguage(...)
+   * ```
+   *
    * @param reference the element to shorten references in.
-   * @return the element in the PSI tree after the shorten references operation corresponding to the original element.
+   * @return the element after the shorten references operation corresponding to the original element.
    */
   fun shortenReference(reference: UReferenceExpression): UReferenceExpression?
+
+  /**
+   * Import the qualifier of the specified element as an on demand import (star import).
+   *
+   * Example:
+   * ```
+   * UastCodeGenerationPlugin.byLanguage(...)
+   * ```
+   * Becomes:
+   * ```
+   * import com.jetbrains.uast.generate.UastCodeGenerationPlugin.*
+   *
+   * byLanguage(...)
+   * ```
+   *
+   * @param reference the qualified element to import
+   * @return the selector part of the qualified reference after importing
+   */
+  fun importMemberOnDemand(reference: UQualifiedReferenceExpression): UExpression?
 }
 
 /**
@@ -151,6 +181,9 @@ fun UReferenceExpression.bindToElement(element: PsiElement): PsiElement? =
 fun UReferenceExpression.shortenReference(): UReferenceExpression? =
   UastCodeGenerationPlugin.byLanguage(this.lang)?.shortenReference(this)
 
+fun UQualifiedReferenceExpression.importMemberOnDemand(): UExpression? =
+  UastCodeGenerationPlugin.byLanguage(this.lang)?.importMemberOnDemand(this)
+
 @ApiStatus.Experimental
 inline fun <reified T : UElement> T.refreshed() = sourcePsi?.also {
   logger<UastCodeGenerationPlugin>().assertTrue(it.isValid,

uast/uast-java/intellij.java.uast.iml

@@ -12,5 +12,6 @@
     <orderEntry type="library" name="kotlin-stdlib-jdk8" level="project" />
     <orderEntry type="module" module-name="intellij.java.psi.impl" />
     <orderEntry type="module" module-name="intellij.java.analysis.impl" />
+    <orderEntry type="module" module-name="intellij.java.impl" />
   </component>
 </module>

uast/uast-java/src/org/jetbrains/uast/java/generate/JavaUastCodeGenerationPlugin.kt

@@ -2,6 +2,7 @@
 package org.jetbrains.uast.java.generate
 
 import com.intellij.codeInsight.BlockUtils
+import com.intellij.codeInsight.intention.impl.AddOnDemandStaticImportAction
 import com.intellij.lang.Language
 import com.intellij.lang.java.JavaLanguage
 import com.intellij.lang.jvm.JvmModifier
@@ -14,10 +15,7 @@ import com.intellij.psi.impl.PsiDiamondTypeUtil
 import com.intellij.psi.impl.source.tree.CompositeElement
 import com.intellij.psi.impl.source.tree.ElementType
 import com.intellij.psi.impl.source.tree.java.PsiLiteralExpressionImpl
-import com.intellij.psi.util.PsiTypesUtil
-import com.intellij.psi.util.PsiUtil
-import com.intellij.psi.util.elementType
-import com.intellij.psi.util.siblings
+import com.intellij.psi.util.*
 import com.intellij.util.castSafelyTo
 import com.siyeh.ig.psiutils.ParenthesesUtils
 import org.jetbrains.uast.*
@@ -103,6 +101,20 @@ internal class JavaUastCodeGenerationPlugin : UastCodeGenerationPlugin {
     val styleManager = JavaCodeStyleManager.getInstance(sourceReference.project)
     return styleManager.shortenClassReferences(sourceReference).toUElementOfType()
   }
+
+  override fun importMemberOnDemand(reference: UQualifiedReferenceExpression): UExpression? {
+    val source = reference.sourcePsi ?: return null
+    val (qualifier, selector) = when (source) {
+      is PsiMethodCallExpression -> source.methodExpression.qualifierExpression to source.methodExpression
+      is PsiReferenceExpression -> source.qualifier to source.referenceNameElement
+      else -> return null
+    }
+    if (qualifier == null || selector == null) return null
+    val ptr = SmartPointerManager.createPointer(selector)
+    val qualifierIdentifier = qualifier.childrenOfType<PsiIdentifier>().firstOrNull() ?: return null
+    AddOnDemandStaticImportAction.invoke(source.project, source.containingFile, null, qualifierIdentifier)
+    return ptr.element?.parent.toUElementOfType()
+  }
 }
 
 private fun PsiElementFactory.createExpressionStatement(expression: PsiExpression): PsiStatement? {