From 447a7d089146b2b4d09f39434af286e8e4322627 Mon Sep 17 00:00:00 2001 From: "Vladimir.Koshelev" Date: Fri, 1 Aug 2025 15:43:57 +0200 Subject: [PATCH] [PY-82607] introduce a caching for all computations in PyTypingTypeProvider GitOrigin-RevId: ad795a4fc2bed8dc7f5276eb18c80a51485dcabf --- .../python/psi/types/TypeEvalContext.java | 8 ++++- python/python-psi-impl/BUILD.bazel | 1 + .../intellij.python.psi.impl.iml | 1 + .../typing/PyTypingTypeProvider.java | 32 +++++++++++++++++-- 4 files changed, 39 insertions(+), 3 deletions(-) diff --git a/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.java b/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.java index 0772fe262ed3..0ee89aaca784 100644 --- a/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.java +++ b/python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.java @@ -52,7 +52,8 @@ public sealed class TypeEvalContext { protected final Map myEvaluated = createMap(); protected final Map myEvaluatedReturn = createMap(); - + @ApiStatus.Internal + protected final Map, 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, PyType> getContextTypeCache() { + return contextTypeCache; + } + /** * @return context constraints (see {@link TypeEvalConstraints} */ diff --git a/python/python-psi-impl/BUILD.bazel b/python/python-psi-impl/BUILD.bazel index e14a4129f68c..58b0c05037d1 100644 --- a/python/python-psi-impl/BUILD.bazel +++ b/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"] diff --git a/python/python-psi-impl/intellij.python.psi.impl.iml b/python/python-psi-impl/intellij.python.psi.impl.iml index 834c3024683e..95128ab9f127 100644 --- a/python/python-psi-impl/intellij.python.psi.impl.iml +++ b/python/python-psi-impl/intellij.python.psi.impl.iml @@ -45,5 +45,6 @@ + \ No newline at end of file diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java index 6add02ed41fb..85bef5e9885e 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java +++ b/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 getType(@NotNull PyExpression expression, @NotNull Context context) { + PyType type = context.getKnownType(expression); + if (type != null) { + return Ref.create(type); + } for (Pair pair : tryResolvingWithAliases(expression, context.getTypeContext())) { final Ref 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;