[PY-82607] introduce a caching for all computations in PyTypingTypeProvider

GitOrigin-RevId: ad795a4fc2bed8dc7f5276eb18c80a51485dcabf

python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.java

@@ -52,7 +52,8 @@ public sealed class TypeEvalContext {
 
   protected final Map<PyTypedElement, PyType> myEvaluated = createMap();
   protected final Map<PyCallable, PyType> myEvaluatedReturn = createMap();
-
+  @ApiStatus.Internal
+  protected final Map<Pair<PyExpression, Object>, PyType> contextTypeCache = createMap();
   /**
    * AssumptionContext invariant requires that if type is in the map, 
    * it's dependencies are also in the map, so we can't use softValueMap.
@@ -330,6 +331,11 @@ public sealed class TypeEvalContext {
     return myConstraints.myOrigin;
   }
 
+  @ApiStatus.Internal
+  public @NotNull Map<Pair<PyExpression, Object>, PyType> getContextTypeCache() {
+    return contextTypeCache;
+  }
+
   /**
    * @return context constraints (see {@link TypeEvalConstraints}
    */

python/python-psi-impl/BUILD.bazel

@@ -49,6 +49,7 @@ jvm_library(
     "//python/python-parser:parser",
     "//python/python-syntax-core:syntax-core",
     "//python/impl.helperLocator:community-helpersLocator",
+    "@lib//:hash4j",
   ],
   exports = ["//python/python-syntax-core:syntax-core"],
   runtime_deps = [":psi-impl_resources"]

python/python-psi-impl/intellij.python.psi.impl.iml

@@ -45,5 +45,6 @@
     <orderEntry type="module" module-name="intellij.python.parser" />
     <orderEntry type="module" module-name="intellij.python.syntax.core" exported="" />
     <orderEntry type="module" module-name="intellij.python.community.helpersLocator" />
+    <orderEntry type="library" name="hash4j" level="project" />
   </component>
 </module>

python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java

@@ -1,6 +1,7 @@
 // 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 com.jetbrains.python.codeInsight.typing;
 
+import com.dynatrace.hash4j.hashing.HashValue128;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.intellij.openapi.util.*;
@@ -48,8 +49,11 @@ import org.jetbrains.annotations.Nullable;
 import java.util.*;
 import java.util.function.Function;
 import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
+import static com.dynatrace.hash4j.hashing.Hashing.murmur3_128;
+import static com.dynatrace.hash4j.hashing.Hashing.xxh3_128;
 import static com.intellij.openapi.util.RecursionManager.doPreventingRecursion;
 import static com.jetbrains.python.psi.PyKnownDecorator.TYPING_FINAL;
 import static com.jetbrains.python.psi.PyKnownDecorator.TYPING_FINAL_EXT;
@@ -783,9 +787,16 @@ public final class PyTypingTypeProvider extends PyTypeProviderWithCustomContext<
   }
 
   private static @Nullable Ref<PyType> getType(@NotNull PyExpression expression, @NotNull Context context) {
+    PyType type = context.getKnownType(expression);
+    if (type != null) {
+      return Ref.create(type);
+    }
     for (Pair<PyQualifiedNameOwner, PsiElement> pair : tryResolvingWithAliases(expression, context.getTypeContext())) {
       final Ref<PyType> typeRef = getTypeForResolvedElement(expression, pair.getFirst(), pair.getSecond(), context);
       if (typeRef != null) {
+        if (typeRef.get() != null) {
+          context.assumeType(expression, typeRef.get());
+        }
         return typeRef;
       }
     }
@@ -1113,8 +1124,9 @@ public final class PyTypingTypeProvider extends PyTypeProviderWithCustomContext<
           // We need this check for the type argument list because getParameterizedType() relies on getClassType() for
           // getting the type corresponding to the subscription expression operand.
           if (classLikeType instanceof PyClassType classType &&
-              isGeneric(classLikeType, typeContext) &&
-              !(typeHint.getParent() instanceof PySubscriptionExpression se && typeHint.equals(se.getOperand()))) {
+              !(getStubRetainedTypeHintContext(typeHint) instanceof PyClass) &&
+              !(typeHint.getParent() instanceof PySubscriptionExpression se && typeHint.equals(se.getOperand())) &&
+              isGeneric(classType, context.myContext)) {
             PyCollectionType parameterized = parameterizeClassDefaultAware(classType.getPyClass(), List.of(), context);
             if (parameterized != null) {
               return Ref.create(parameterized.toInstance());
@@ -2320,6 +2332,22 @@ public final class PyTypingTypeProvider extends PyTypeProviderWithCustomContext<
       return myTypeAliasStack;
     }
 
+    public @Nullable PyType getKnownType(@NotNull PyExpression expression) {
+      //noinspection SuspiciousMethodCalls
+      return myContext.getContextTypeCache().get(new Pair<>(expression, getContextStrongHashValue()));
+    }
+
+    public void assumeType(@NotNull PyExpression expression, @NotNull PyType type) {
+      myContext.getContextTypeCache().put(new Pair<>(expression, getContextStrongHashValue()), type);
+    }
+
+    private @NotNull HashValue128 getContextStrongHashValue() {
+      var result = xxh3_128().hashCharsTo128Bits(Stream.concat(Stream.of(myComputeTypeParameterScope ? "1" : "0"),
+                                                    myTypeAliasStack.stream().map(it -> it.getQualifiedName()))
+                                                   .collect(Collectors.joining("#")));
+      return result;
+    }
+
     @Override
     public boolean equals(Object o) {
       if (this == o) return true;