From e0ecc967cebab12d4c883edfa0217bfe14e92da8 Mon Sep 17 00:00:00 2001
From: Semyon Proshev <Semyon.Proshev@jetbrains.com>
Date: Thu, 23 Apr 2020 23:47:16 +0300
Subject: [PATCH] Add a cache to PythonDialectsTokenSetProvider (PY-41923)

GitOrigin-RevId: 0bc6ab1cbffc795afc371046e4b322d83865348c
---
 .../PythonDialectsTokenSetProvider.java       | 34 ++++++++++++++-----
 1 file changed, 26 insertions(+), 8 deletions(-)

diff --git a/python/python-psi-impl/src/com/jetbrains/python/PythonDialectsTokenSetProvider.java b/python/python-psi-impl/src/com/jetbrains/python/PythonDialectsTokenSetProvider.java
index 17179a47d61d..50e993fab65c 100644
--- a/python/python-psi-impl/src/com/jetbrains/python/PythonDialectsTokenSetProvider.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/PythonDialectsTokenSetProvider.java
@@ -1,12 +1,14 @@
 // Copyright 2000-2018 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;
 
+import com.intellij.openapi.Disposable;
 import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.openapi.components.Service;
 import com.intellij.psi.tree.TokenSet;
 import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.NotNull;
 
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Function;
 
 /**
@@ -15,7 +17,14 @@ import java.util.function.Function;
  * @author vlan
  */
 @Service
-public final class PythonDialectsTokenSetProvider {
+public final class PythonDialectsTokenSetProvider implements Disposable {
+
+  @NotNull
+  private final ConcurrentHashMap<String, TokenSet> myCache = new ConcurrentHashMap<>(7); // number of token types
+
+  public PythonDialectsTokenSetProvider() {
+    PythonDialectsTokenSetContributor.EP_NAME.addExtensionPointListener(() -> myCache.clear(), this);
+  }
 
   @NotNull
   public static PythonDialectsTokenSetProvider getInstance() {
@@ -26,49 +35,58 @@ public final class PythonDialectsTokenSetProvider {
    * Returns all element types of Python dialects that are subclasses of {@link com.jetbrains.python.psi.PyStatement}.
    */
   public TokenSet getStatementTokens() {
-    return orSets(PythonDialectsTokenSetContributor::getStatementTokens);
+    return getTokenSet("statement", PythonDialectsTokenSetContributor::getStatementTokens);
   }
 
   /**
    * Returns all element types of Python dialects that are subclasses of {@link com.jetbrains.python.psi.PyExpression}.
    */
   public TokenSet getExpressionTokens() {
-    return orSets(PythonDialectsTokenSetContributor::getExpressionTokens);
+    return getTokenSet("expression", PythonDialectsTokenSetContributor::getExpressionTokens);
   }
 
   /**
    * Returns all element types of Python dialects that are language keywords.
    */
   public TokenSet getKeywordTokens() {
-    return orSets(PythonDialectsTokenSetContributor::getKeywordTokens);
+    return getTokenSet("keyword", PythonDialectsTokenSetContributor::getKeywordTokens);
   }
 
   /**
    * Returns all element types of Python dialects that are subclasses of {@link com.jetbrains.python.psi.PyParameter}.
    */
   public TokenSet getParameterTokens() {
-    return orSets(PythonDialectsTokenSetContributor::getParameterTokens);
+    return getTokenSet("parameter", PythonDialectsTokenSetContributor::getParameterTokens);
   }
 
   /**
    * Returns all element types of Python dialects that are subclasses of {@link com.jetbrains.python.psi.PyFunction}.
    */
   public TokenSet getFunctionDeclarationTokens() {
-    return orSets(PythonDialectsTokenSetContributor::getFunctionDeclarationTokens);
+    return getTokenSet("functionDeclaration", PythonDialectsTokenSetContributor::getFunctionDeclarationTokens);
   }
 
   /**
    * Returns all element types of Python dialects that can be used as unbalanced braces recovery tokens in the lexer.
    */
   public TokenSet getUnbalancedBracesRecoveryTokens() {
-    return orSets(PythonDialectsTokenSetContributor::getUnbalancedBracesRecoveryTokens);
+    return getTokenSet("unbalancedBracesRecovery", PythonDialectsTokenSetContributor::getUnbalancedBracesRecoveryTokens);
   }
 
   /**
    * Returns all element types of Python dialects that are subclasses of {@link com.jetbrains.python.psi.PyReferenceExpression}.
    */
   public TokenSet getReferenceExpressionTokens() {
-    return orSets(PythonDialectsTokenSetContributor::getReferenceExpressionTokens);
+    return getTokenSet("referenceExpression", PythonDialectsTokenSetContributor::getReferenceExpressionTokens);
+  }
+
+  @Override
+  public void dispose() {
+    myCache.clear();
+  }
+
+  private @NotNull TokenSet getTokenSet(@NotNull String key, @NotNull Function<PythonDialectsTokenSetContributor, TokenSet> getter) {
+    return myCache.computeIfAbsent(key, __ -> orSets(getter));
   }
 
   private static @NotNull TokenSet orSets(@NotNull Function<PythonDialectsTokenSetContributor, TokenSet> getter) {