Cache the results of PyStarImportElementImpl.multiResolveName() and ResolveImportUtil.resolveModule() (PY-19141)

These methods used to consume up to 86% of CPU for code completion in
this example that contained lots of star imports and other import
constructs:

    import numpy as np
    np.<caret>  # Invoke completion

Now we cache the results using our new
PyUtil.getParameterizedCachedValue() that stores a value in the user
data of a PsiElement for a given parameter.

python/src/com/jetbrains/python/psi/PyUtil.java

@@ -54,9 +54,7 @@ import com.intellij.openapi.wm.WindowManager;
 import com.intellij.psi.*;
 import com.intellij.psi.codeStyle.CodeStyleSettingsManager;
 import com.intellij.psi.stubs.StubElement;
-import com.intellij.psi.util.PsiTreeUtil;
-import com.intellij.psi.util.PsiUtilCore;
-import com.intellij.psi.util.QualifiedName;
+import com.intellij.psi.util.*;
 import com.intellij.ui.awt.RelativePoint;
 import com.intellij.util.*;
 import com.intellij.util.containers.ContainerUtil;
@@ -847,6 +845,22 @@ public class PyUtil {
     });
   }
 
+  public static <T, P> T getParameterizedCachedValue(@NotNull PsiElement element, @NotNull P param, @NotNull NotNullFunction<P, T> f) {
+    final Map<P, T> cache = CachedValuesManager.getCachedValue(element, new CachedValueProvider<Map<P, T>>() {
+      @Nullable
+      @Override
+      public Result<Map<P, T>> compute() {
+        return Result.create(Maps.newHashMap(), PsiModificationTracker.MODIFICATION_COUNT);
+      }
+    });
+    T result = cache.get(param);
+    if (result == null) {
+      result = f.fun(param);
+      cache.put(param, result);
+    }
+    return result;
+  }
+
   public static class KnownDecoratorProviderHolder {
     public static PyKnownDecoratorProvider[] KNOWN_DECORATOR_PROVIDERS = Extensions.getExtensions(PyKnownDecoratorProvider.EP_NAME);
 

python/src/com/jetbrains/python/psi/impl/PyStarImportElementImpl.java

@@ -41,7 +41,6 @@ import java.util.List;
  * @author dcheryasov
  */
 public class PyStarImportElementImpl extends PyBaseElementImpl<PyStarImportElementStub> implements PyStarImportElement {
-
   public PyStarImportElementImpl(ASTNode astNode) {
     super(astNode);
   }
@@ -80,7 +79,12 @@ public class PyStarImportElementImpl extends PyBaseElementImpl<PyStarImportEleme
   }
 
   @NotNull
-  public List<RatedResolveResult> multiResolveName(@NotNull final String name) {
+  public List<RatedResolveResult> multiResolveName(@NotNull String name) {
+    return PyUtil.getParameterizedCachedValue(this, name, this::calculateMultiResolveName);
+  }
+
+  @NotNull
+  private List<RatedResolveResult> calculateMultiResolveName(@NotNull String name) {
     if (PyUtil.isClassPrivateName(name)) {
       return Collections.emptyList();
     }
@@ -96,6 +100,9 @@ public class PyStarImportElementImpl extends PyBaseElementImpl<PyStarImportEleme
           final List<? extends RatedResolveResult> results = moduleType.resolveMember(name, null, AccessDirection.READ,
                                                                                       PyResolveContext.defaultContext());
           if (results != null && !results.isEmpty() && PyUtil.isStarImportableFrom(name, sourceFile)) {
+            if (results.isEmpty()) {
+              return Collections.emptyList();
+            }
             final List<RatedResolveResult> res = Lists.newArrayList();
             for (RatedResolveResult result : results) {
               res.add(result);

python/src/com/jetbrains/python/psi/resolve/ResolveImportUtil.java

@@ -191,6 +191,17 @@ public class ResolveImportUtil {
     if (qualifiedName == null || sourceFile == null) {
       return Collections.emptyList();
     }
+    final ResolveModuleParams params = new ResolveModuleParams(qualifiedName, sourceFile, importIsAbsolute, relativeLevel);
+    return PyUtil.getParameterizedCachedValue(sourceFile, params, ResolveImportUtil::calculateResolveModule);
+  }
+
+  @NotNull
+  private static List<PsiElement> calculateResolveModule(@NotNull ResolveModuleParams params) {
+    final QualifiedName qualifiedName = params.getName();
+    final int relativeLevel = params.getLevel();
+    final PsiFile sourceFile = params.getFile();
+    final boolean importIsAbsolute = params.isAbsolute();
+
     final String marker = qualifiedName + "#" + Integer.toString(relativeLevel);
     final Set<String> beingImported = ourBeingImported.get();
     if (beingImported.contains(marker)) {
@@ -486,4 +497,60 @@ public class ResolveImportUtil {
     }
     return PointInImport.NONE;
   }
+
+  private static final class ResolveModuleParams {
+    @NotNull private final QualifiedName myName;
+    @NotNull private final PsiFile myFile;
+    private final boolean myAbsolute;
+    private final int myLevel;
+
+    public ResolveModuleParams(@NotNull QualifiedName qualifiedName, @NotNull PsiFile file , boolean importIsAbsolute, int relativeLevel) {
+      myName = qualifiedName;
+      myFile = file;
+      myAbsolute = importIsAbsolute;
+      myLevel = relativeLevel;
+    }
+
+    @NotNull
+    public QualifiedName getName() {
+      return myName;
+    }
+
+    public boolean isAbsolute() {
+      return myAbsolute;
+    }
+
+    public int getLevel() {
+      return myLevel;
+    }
+
+    @NotNull
+    public PsiFile getFile() {
+      return myFile;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) return true;
+      if (o == null || getClass() != o.getClass()) return false;
+
+      ResolveModuleParams params = (ResolveModuleParams)o;
+
+      if (myAbsolute != params.myAbsolute) return false;
+      if (myLevel != params.myLevel) return false;
+      if (!myName.equals(params.myName)) return false;
+      if (!myFile.equals(params.myFile)) return false;
+
+      return true;
+    }
+
+    @Override
+    public int hashCode() {
+      int result = myName.hashCode();
+      result = 31 * result + myFile.hashCode();
+      result = 31 * result + (myAbsolute ? 1 : 0);
+      result = 31 * result + myLevel;
+      return result;
+    }
+  }
 }