From e3d816402ab871d7db1bd1ee580de321d27f7c6f Mon Sep 17 00:00:00 2001
From: Daniil Ovchinnikov <daniil.ovchinnikov@jetbrains.com>
Date: Tue, 27 Sep 2016 23:39:04 +0300
Subject: [PATCH] [groovy] cache base class reference within stub (follow up
 IDEA-CR-14096)

---
 .../GrAnonymousClassDefinitionImpl.java       | 21 +++++++------------
 .../lang/psi/stubs/GrTypeDefinitionStub.java  | 18 ++++++++++++++++
 2 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/psi/impl/statements/typedef/GrAnonymousClassDefinitionImpl.java b/plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/psi/impl/statements/typedef/GrAnonymousClassDefinitionImpl.java
index 2a44e6dd2752..98f8fb7cbb09 100644
--- a/plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/psi/impl/statements/typedef/GrAnonymousClassDefinitionImpl.java
+++ b/plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/psi/impl/statements/typedef/GrAnonymousClassDefinitionImpl.java
@@ -17,7 +17,6 @@
 package org.jetbrains.plugins.groovy.lang.psi.impl.statements.typedef;
 
 import com.intellij.lang.ASTNode;
-import com.intellij.openapi.util.NullableLazyValue;
 import com.intellij.psi.*;
 import com.intellij.psi.scope.PsiScopeProcessor;
 import com.intellij.reference.SoftReference;
@@ -26,7 +25,6 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.plugins.groovy.lang.parser.GroovyElementTypes;
 import org.jetbrains.plugins.groovy.lang.psi.GroovyElementVisitor;
-import org.jetbrains.plugins.groovy.lang.psi.GroovyPsiElementFactory;
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.arguments.GrArgumentList;
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.typedef.GrAnonymousClassDefinition;
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.typedef.GrExtendsClause;
@@ -45,15 +43,7 @@ import org.jetbrains.plugins.groovy.lang.psi.util.GroovyCommonClassNames;
  */
 public class GrAnonymousClassDefinitionImpl extends GrTypeDefinitionImpl implements GrAnonymousClassDefinition {
 
-  private SoftReference<PsiClassType> myCachedBaseType = null;
-
-  private NullableLazyValue<GrCodeReferenceElement> myStubBaseReference = NullableLazyValue.createValue(() -> {
-    GrTypeDefinitionStub stub = getStub();
-    if (stub == null) return null;
-    String baseClassName = stub.getBaseClassName();
-    assert baseClassName != null;
-    return GroovyPsiElementFactory.getInstance(getProject()).createReferenceElementFromText(baseClassName, this);
-  });
+  private SoftReference<PsiClassType> myCachedBaseType;
 
   public GrAnonymousClassDefinitionImpl(@NotNull ASTNode node) {
     super(node);
@@ -70,8 +60,13 @@ public class GrAnonymousClassDefinitionImpl extends GrTypeDefinitionImpl impleme
   @Override
   @NotNull
   public GrCodeReferenceElement getBaseClassReferenceGroovy() {
-    GrCodeReferenceElement stubReference = myStubBaseReference.getValue();
-    if (stubReference != null) return stubReference;
+    GrTypeDefinitionStub stub = getStub();
+    if (stub != null) {
+      GrCodeReferenceElement reference = stub.getBaseClassReference();
+      assert reference != null;
+      return reference;
+    }
+
     return findNotNullChildByClass(GrCodeReferenceElement.class); //not null because of definition =)
   }
 
diff --git a/plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/psi/stubs/GrTypeDefinitionStub.java b/plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/psi/stubs/GrTypeDefinitionStub.java
index 33444fd2496d..2d24b807868d 100644
--- a/plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/psi/stubs/GrTypeDefinitionStub.java
+++ b/plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/psi/stubs/GrTypeDefinitionStub.java
@@ -20,11 +20,14 @@ import com.intellij.psi.stubs.IStubElementType;
 import com.intellij.psi.stubs.NamedStub;
 import com.intellij.psi.stubs.StubBase;
 import com.intellij.psi.stubs.StubElement;
+import com.intellij.reference.SoftReference;
 import com.intellij.util.io.StringRef;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.plugins.groovy.lang.psi.GroovyPsiElementFactory;
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.typedef.GrAnonymousClassDefinition;
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.typedef.GrTypeDefinition;
+import org.jetbrains.plugins.groovy.lang.psi.api.types.GrCodeReferenceElement;
 
 /**
  * @author ilyas
@@ -43,6 +46,8 @@ public class GrTypeDefinitionStub extends StubBase<GrTypeDefinition> implements
   private final String[] myAnnotations;
   private final byte myFlags;
 
+  private SoftReference<GrCodeReferenceElement> myStubBaseReference;
+
   public GrTypeDefinitionStub(StubElement parent,
                                   final String name,
                                   @Nullable final String baseClassName,
@@ -63,6 +68,19 @@ public class GrTypeDefinitionStub extends StubBase<GrTypeDefinition> implements
     return myBaseClassName;
   }
 
+  @Nullable
+  public GrCodeReferenceElement getBaseClassReference() {
+    String baseClassName = getBaseClassName();
+    if (baseClassName == null) return null;
+
+    GrCodeReferenceElement reference = SoftReference.dereference(myStubBaseReference);
+    if (reference == null) {
+      reference = GroovyPsiElementFactory.getInstance(getProject()).createReferenceElementFromText(baseClassName, getPsi());
+      myStubBaseReference = new SoftReference<>(reference);
+    }
+    return reference;
+  }
+
   @Override
   public String getName() {
     return StringRef.toString(myName);