Uast: caching UElements as SemElements (IDEA-180552)

java/java-analysis-impl/src/com/intellij/uast/UastElementAnchorProvider.java

@@ -17,10 +17,11 @@ package com.intellij.uast;
 
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.impl.smartPointers.SmartPointerAnchorProvider;
-import com.intellij.reference.SoftReference;
+import com.intellij.semantic.SemService;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.uast.UElement;
+import org.jetbrains.uast.UElementSemContributorKt;
 import org.jetbrains.uast.UastContextKt;
 
 /**
@@ -32,9 +33,8 @@ public class UastElementAnchorProvider extends SmartPointerAnchorProvider {
   public PsiElement getAnchor(@NotNull PsiElement element) {
     if (element instanceof UElement) {
       PsiElement psi = ((UElement)element).getPsi();
-      if (psi != null) {
-        psi.putUserData(UastContextKt.getCACHED_UELEMENT_KEY(), new SoftReference<>((UElement)element));
-      }
+      SemService.getSemService(psi.getProject())
+        .setCachedSemElement(UElementSemContributorKt.UAST_SEM_KEY, psi, (UElement)element);
       return psi;
     }
     return null;

platform/platform-resources/src/META-INF/LangExtensions.xml

@@ -1050,6 +1050,8 @@
     <automaticRenamerFactory implementation="com.intellij.refactoring.rename.RelatedFilesRenamerFactory"/>
     <refactoring.moveHandler implementation="com.intellij.refactoring.move.moveFilesOrDirectories.MoveRelatedFilesHandler"
                              order="before moveFileOrDir"/>
+
+    <semContributor implementation="org.jetbrains.uast.UElementSemContributor"/>
   </extensions>
 </idea-plugin>
 

uast/uast-common/src/org/jetbrains/uast/UElementSemContributor.kt

@@ -0,0 +1,23 @@
+// Copyright 2000-2017 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+
+package org.jetbrains.uast
+
+import com.intellij.openapi.components.ServiceManager
+import com.intellij.patterns.PlatformPatterns
+import com.intellij.psi.PsiElement
+import com.intellij.semantic.SemContributor
+import com.intellij.semantic.SemKey
+import com.intellij.semantic.SemRegistrar
+
+@JvmField
+val UAST_SEM_KEY = SemKey.createKey<UElement>("org.jetbrains.uast.uElement.semKey")!!
+
+class UElementSemContributor : SemContributor() {
+  override fun registerSemProviders(registrar: SemRegistrar) {
+    registrar.registerSemElementProvider(UAST_SEM_KEY, PlatformPatterns.psiElement(), { sourceElement: PsiElement ->
+
+      ServiceManager.getService(sourceElement.project, UastContext::class.java).findPlugin(sourceElement)
+        ?.convertElementWithParent(sourceElement, null)
+    })
+  }
+}

uast/uast-common/src/org/jetbrains/uast/UastContext.kt

@@ -18,11 +18,9 @@ package org.jetbrains.uast
 import com.intellij.lang.Language
 import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.project.Project
-import com.intellij.openapi.util.Key
 import com.intellij.psi.*
-import com.intellij.reference.SoftReference
+import com.intellij.semantic.SemService
 
-val CACHED_UELEMENT_KEY = Key.create<SoftReference<UElement>>("org.jetbrains.uast.cachedElement")
 
 /**
  * Manages the UAST to PSI conversion.
@@ -32,6 +30,8 @@ class UastContext(val project: Project) : UastLanguagePlugin {
     private val CONTEXT_LANGUAGE = object : Language("UastContextLanguage") {}
   }
 
+  private val mySemService = SemService.getSemService(project)
+
   override val language: Language
     get() = CONTEXT_LANGUAGE
 
@@ -55,18 +55,23 @@ class UastContext(val project: Project) : UastLanguagePlugin {
   fun getClass(clazz: PsiClass): UClass = convertWithParent<UClass>(clazz)!!
 
   override fun convertElement(element: PsiElement, parent: UElement?, requiredType: Class<out UElement>?): UElement? {
-    val cachedElement = element.getUserData(CACHED_UELEMENT_KEY)?.get()
-    if (cachedElement != null) {
-      return if (requiredType == null || requiredType.isInstance(cachedElement)) cachedElement else null
+    if (parent != null) {
+      val cachedElement = mySemService.getCachedSemElements(UAST_SEM_KEY, element)?.firstOrNull()
+        ?.takeIf { requiredType?.isAssignableFrom(it.javaClass) ?: true }
+
+      if (cachedElement != null && cachedElement.uastParent === parent) {
+        return cachedElement
+      }
     }
 
     return findPlugin(element)?.convertElement(element, parent, requiredType)
   }
 
   override fun convertElementWithParent(element: PsiElement, requiredType: Class<out UElement>?): UElement? {
-    val cachedElement = element.getUserData(CACHED_UELEMENT_KEY)?.get()
-    if (cachedElement != null) {
-      return if (requiredType == null || requiredType.isInstance(cachedElement)) cachedElement else null
+    val defaultUElement = mySemService.getSemElement(UAST_SEM_KEY, element)
+
+    if (defaultUElement == null || requiredType == null || requiredType.isAssignableFrom(defaultUElement.javaClass)) {
+      return defaultUElement
     }
 
     return findPlugin(element)?.convertElementWithParent(element, requiredType)

uast/uast-common/src/org/jetbrains/uast/baseElements/UElement.kt

@@ -16,6 +16,7 @@
 package org.jetbrains.uast
 
 import com.intellij.psi.PsiElement
+import com.intellij.semantic.SemElement
 import org.jetbrains.annotations.ApiStatus
 import org.jetbrains.uast.visitor.UastTypedVisitor
 import org.jetbrains.uast.visitor.UastVisitor
@@ -23,7 +24,7 @@ import org.jetbrains.uast.visitor.UastVisitor
 /**
  * The common interface for all Uast elements.
  */
-interface UElement {
+interface UElement : SemElement {
   /**
    * Returns the element parent.
    */

uast/uast-common/uast-common.iml

@@ -10,5 +10,6 @@
     <orderEntry type="library" name="KotlinJavaRuntime" level="project" />
     <orderEntry type="module" module-name="java-psi-api" exported="" />
     <orderEntry type="module" module-name="java-psi-impl" />
+    <orderEntry type="module" module-name="lang-api" exported="" />
   </component>
 </module>

uast/uast-tests/test/org/jetbrains/uast/test/java/JavaUastPerformanceTest.kt

@@ -49,6 +49,13 @@ class JavaUastPerformanceTest : LightCodeInsightFixtureTestCase() {
       }.attempts(1).assertTiming()
     }
 
+    val cachedWalk = EachPsiToUastWalker().apply {
+      PlatformTestUtil.startPerformanceTest("convert each element to uast second time", 400) {
+        clazz.accept(this)
+        TestCase.assertEquals(expectedUElementsCount, totalCount)
+      }.attempts(1).assertTiming()
+    }
+
   }