From b2f5993e47bf12aaea6c37b1c7646475b52278bd Mon Sep 17 00:00:00 2001
From: Dmitry Jemerov <yole@jetbrains.com>
Date: Tue, 4 Apr 2017 14:56:43 +0200
Subject: [PATCH] Support creating smart pointers to UAST elements

---
 java/java-analysis-api/java-analysis-api.iml  |  2 +-
 .../uast/UastElementAnchorProvider.java       | 48 +++++++++++++++++++
 .../src/META-INF/LangExtensions.xml           |  1 +
 .../src/org/jetbrains/uast/UastContext.kt     | 14 ++++++
 4 files changed, 64 insertions(+), 1 deletion(-)
 create mode 100644 java/java-analysis-impl/src/com/intellij/uast/UastElementAnchorProvider.java

diff --git a/java/java-analysis-api/java-analysis-api.iml b/java/java-analysis-api/java-analysis-api.iml
index c62ac8a6ea10..84a6db05d79a 100644
--- a/java/java-analysis-api/java-analysis-api.iml
+++ b/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>
\ No newline at end of file
diff --git a/java/java-analysis-impl/src/com/intellij/uast/UastElementAnchorProvider.java b/java/java-analysis-impl/src/com/intellij/uast/UastElementAnchorProvider.java
new file mode 100644
index 000000000000..3873f35ad98a
--- /dev/null
+++ b/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);
+  }
+}
diff --git a/platform/platform-resources/src/META-INF/LangExtensions.xml b/platform/platform-resources/src/META-INF/LangExtensions.xml
index 8a6c0cbbda6c..2c689c06b5ce 100644
--- a/platform/platform-resources/src/META-INF/LangExtensions.xml
+++ b/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>
 
diff --git a/uast/uast-common/src/org/jetbrains/uast/UastContext.kt b/uast/uast-common/src/org/jetbrains/uast/UastContext.kt
index 80111f822e22..5d6a54eda809 100644
--- a/uast/uast-common/src/org/jetbrains/uast/UastContext.kt
+++ b/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)
     }