Support creating smart pointers to UAST elements

java/java-analysis-api/java-analysis-api.iml

@@ -11,6 +11,6 @@
     <orderEntry type="module" module-name="java-psi-api" />
     <orderEntry type="module" module-name="projectModel-api" />
     <orderEntry type="library" name="KotlinJavaRuntime" level="project" />
-    <orderEntry type="module" module-name="uast-common" />
+    <orderEntry type="module" module-name="uast-common" exported="" />
   </component>
 </module>

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

@@ -0,0 +1,48 @@
+/*
+ * Copyright 2000-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.intellij.uast;
+
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.impl.smartPointers.SmartPointerAnchorProvider;
+import com.intellij.reference.SoftReference;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+import org.jetbrains.uast.UElement;
+import org.jetbrains.uast.UastContextKt;
+
+/**
+ * @author yole
+ */
+public class UastElementAnchorProvider extends SmartPointerAnchorProvider {
+  @Nullable
+  @Override
+  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));
+      }
+      return psi;
+    }
+    return null;
+  }
+
+  @Nullable
+  @Override
+  public PsiElement restoreElement(@NotNull PsiElement anchor) {
+    return (PsiElement)UastContextKt.toUElement(anchor);
+  }
+}

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

@@ -1025,6 +1025,7 @@
 
     <projectService serviceInterface="org.jetbrains.uast.UastContext" serviceImplementation="org.jetbrains.uast.UastContext"/>
     <metaLanguage implementation="com.intellij.uast.UastMetaLanguage"/>
+    <smartPointer.anchorProvider implementation="com.intellij.uast.UastElementAnchorProvider"/>
   </extensions>
 </idea-plugin>
 

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

@@ -18,7 +18,11 @@ 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
+
+val CACHED_UELEMENT_KEY = Key.create<SoftReference<UElement>>("org.jetbrains.uast.cachedElement")
 
 /**
  * Manages the UAST to PSI conversion.
@@ -51,10 +55,20 @@ 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
+        }
+
         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
+        }
+
         return findPlugin(element)?.convertElementWithParent(element, requiredType)
     }