// 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.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.function.Function;

/**
 * Provides element types of various kinds for known Python dialects.
 *
 * @author vlan
 */
@Service
public final class PythonDialectsTokenSetProvider {

  @NotNull
  public static PythonDialectsTokenSetProvider getInstance() {
    return ApplicationManager.getApplication().getService(PythonDialectsTokenSetProvider.class);
  }

  /**
   * Returns all element types of Python dialects that are subclasses of {@link com.jetbrains.python.psi.PyStatement}.
   */
  public TokenSet getStatementTokens() {
    return orSets(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);
  }

  /**
   * Returns all element types of Python dialects that are language keywords.
   */
  public TokenSet getKeywordTokens() {
    return orSets(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);
  }

  /**
   * Returns all element types of Python dialects that are subclasses of {@link com.jetbrains.python.psi.PyFunction}.
   */
  public TokenSet getFunctionDeclarationTokens() {
    return orSets(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);
  }

  /**
   * Returns all element types of Python dialects that are subclasses of {@link com.jetbrains.python.psi.PyReferenceExpression}.
   */
  public TokenSet getReferenceExpressionTokens() {
    return orSets(PythonDialectsTokenSetContributor::getReferenceExpressionTokens);
  }

  private static @NotNull TokenSet orSets(@NotNull Function<PythonDialectsTokenSetContributor, TokenSet> getter) {
    return TokenSet.orSet(
      StreamEx
        .of(PythonDialectsTokenSetContributor.EP_NAME.getExtensionList())
        .map(getter)
        .toArray(TokenSet.class)
    );
  }
}

@@ -165,16 +165,56 @@
 ython;%0A%0A
+import com.intellij.openapi.Disposable;%0A
 import c
@@ -265,16 +265,16 @@
 anager;%0A
-
 import c
@@ -430,16 +430,63 @@
 tNull;%0A%0A
+import java.util.concurrent.ConcurrentHashMap;%0A
 import j
@@ -670,17 +670,317 @@
 rovider 
-%7B
+implements Disposable %7B%0A%0A  @NotNull%0A  private final ConcurrentHashMap%3CString, TokenSet%3E myCache = new ConcurrentHashMap%3C%3E(7); // number of token types%0A%0A  public PythonDialectsTokenSetProvider() %7B%0A    PythonDialectsTokenSetContributor.EP_NAME.addExtensionPointListener(() -%3E myCache.clear(), this);%0A  %7D
 %0A%0A  @Not
@@ -1324,39 +1324,57 @@
 () %7B%0A    return 
-orSets(
+getTokenSet(%22statement%22, 
 PythonDialectsTo
@@ -1595,39 +1595,58 @@
 () %7B%0A    return 
-orSets(
+getTokenSet(%22expression%22, 
 PythonDialectsTo
@@ -1823,39 +1823,55 @@
 () %7B%0A    return 
-orSets(
+getTokenSet(%22keyword%22, 
 PythonDialectsTo
@@ -2088,39 +2088,57 @@
 () %7B%0A    return 
-orSets(
+getTokenSet(%22parameter%22, 
 PythonDialectsTo
@@ -2366,39 +2366,67 @@
 () %7B%0A    return 
-orSets(
+getTokenSet(%22functionDeclaration%22, 
 PythonDialectsTo
@@ -2669,39 +2669,72 @@
 () %7B%0A    return 
-orSets(
+getTokenSet(%22unbalancedBracesRecovery%22, 
 PythonDialectsTo
@@ -2996,23 +2996,51 @@
  return 
-orSets(
+getTokenSet(%22referenceExpression%22, 
 PythonDi
@@ -3102,16 +3102,282 @@
 );%0A  %7D%0A%0A
+  @Override%0A  public void dispose() %7B%0A    myCache.clear();%0A  %7D%0A%0A  private @NotNull TokenSet getTokenSet(@NotNull String key, @NotNull Function%3CPythonDialectsTokenSetContributor, TokenSet%3E getter) %7B%0A    return myCache.computeIfAbsent(key, __ -%3E orSets(getter));%0A  %7D%0A%0A
   privat

Add a cache to PythonDialectsTokenSetProvider (PY-41923)