java: prohibit caching when using thread-local types imposed on expressions and decla...
[idea/community.git] / java / java-psi-api / src / com / intellij / psi / ThreadLocalTypes.java
1 // Copyright 2000-2019 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.
2 package com.intellij.psi;
3
4 import com.intellij.openapi.util.RecursionGuard;
5 import com.intellij.openapi.util.RecursionManager;
6 import org.jetbrains.annotations.NotNull;
7 import org.jetbrains.annotations.Nullable;
8
9 import java.util.HashMap;
10 import java.util.List;
11 import java.util.Map;
12 import java.util.function.Function;
13
14 public class ThreadLocalTypes {
15   private static final RecursionGuard<ThreadLocalTypes> ourGuard = RecursionManager.createGuard("ThreadLocalTypes");
16   private final Map<PsiElement, PsiType> myMap = new HashMap<>();
17
18   private ThreadLocalTypes() {}
19
20   @Nullable
21   public static PsiType getElementType(@NotNull PsiElement psi) {
22     List<? extends ThreadLocalTypes> stack = ourGuard.currentStack();
23     for (int i = stack.size() - 1; i >= 0; i--) {
24       ThreadLocalTypes types = stack.get(i);
25       PsiType type = types.myMap.get(psi);
26       if (type != null) {
27         ourGuard.prohibitResultCaching(types);
28         return type;
29       }
30     }
31     return null;
32   }
33
34   public static boolean hasBindingFor(@NotNull PsiElement psi) {
35     List<? extends ThreadLocalTypes> stack = ourGuard.currentStack();
36     for (int i = stack.size() - 1; i >= 0; i--) {
37       ThreadLocalTypes types = stack.get(i);
38       if (types.myMap.containsKey(psi)) {
39         ourGuard.prohibitResultCaching(types);
40         return true;
41       }
42     }
43     return false;
44   }
45
46   public static <T> T performWithTypes(@NotNull Function<ThreadLocalTypes, T> action) {
47     ThreadLocalTypes types = new ThreadLocalTypes();
48     return ourGuard.doPreventingRecursion(types, false, () -> action.apply(types));
49   }
50
51   public void forceType(@NotNull PsiElement psi, @Nullable PsiType type) {
52     myMap.put(psi, type);
53   }
54
55 }