83dcf43f8d9ea04df4059fa59620681459a636af
[idea/community.git] / platform / core-impl / src / com / intellij / psi / stubs / ObjectStubTree.java
1 // Copyright 2000-2020 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.stubs;
3
4 import com.intellij.openapi.util.Key;
5 import com.intellij.util.ArrayUtil;
6 import com.intellij.util.containers.ContainerUtil;
7 import gnu.trove.THashMap;
8 import gnu.trove.TObjectHashingStrategy;
9 import gnu.trove.TObjectObjectProcedure;
10 import gnu.trove.TObjectProcedure;
11 import org.jetbrains.annotations.NotNull;
12 import org.jetbrains.annotations.Nullable;
13
14 import java.util.ArrayList;
15 import java.util.List;
16 import java.util.Map;
17 import java.util.function.Function;
18
19 /**
20  * @author Dmitry Avdeev
21  */
22 public class ObjectStubTree<T extends Stub> {
23   private static final Key<ObjectStubTree<?>> STUB_TO_TREE_REFERENCE = Key.create("stub to tree reference");
24   protected final ObjectStubBase<?> myRoot;
25   private String myDebugInfo;
26   private boolean myHasBackReference;
27   private final List<T> myPlainList;
28
29   public ObjectStubTree(final @NotNull ObjectStubBase<?> root, boolean withBackReference) {
30     myRoot = root;
31     myPlainList = enumerateStubs(root);
32     if (withBackReference) {
33       myRoot.putUserData(STUB_TO_TREE_REFERENCE, this); // This will prevent soft references to stub tree to be collected before all of the stubs are collected.
34     }
35   }
36
37   public @NotNull Stub getRoot() {
38     return myRoot;
39   }
40
41   public @NotNull List<T> getPlainList() {
42     return myPlainList;
43   }
44
45   @NotNull
46   List<T> getPlainListFromAllRoots() {
47     return getPlainList();
48   }
49
50   @Deprecated
51   public @NotNull Map<StubIndexKey<?, ?>, Map<Object, int[]>> indexStubTree() {
52     return indexStubTree(key -> ContainerUtil.canonicalStrategy());
53   }
54
55   public @NotNull Map<StubIndexKey<?, ?>, Map<Object, int[]>> indexStubTree(@NotNull Function<StubIndexKey<?, ?>, TObjectHashingStrategy<?>> keyHashingStrategyFunction) {
56     StubIndexSink sink = new StubIndexSink(keyHashingStrategyFunction);
57     final List<T> plainList = getPlainListFromAllRoots();
58     for (int i = 0, plainListSize = plainList.size(); i < plainListSize; i++) {
59       final Stub stub = plainList.get(i);
60       sink.myStubIdx = i;
61       StubSerializationUtil.getSerializer(stub).indexStub(stub, sink);
62     }
63
64     return sink.getResult();
65   }
66
67   protected @NotNull List<T> enumerateStubs(@NotNull Stub root) {
68     List<T> result = new ArrayList<>();
69     //noinspection unchecked
70     enumerateStubsInto(root, (List)result);
71     return result;
72   }
73
74   private static void enumerateStubsInto(@NotNull Stub root, @NotNull List<? super Stub> result) {
75     ((ObjectStubBase)root).id = result.size();
76     result.add(root);
77     List<? extends Stub> childrenStubs = root.getChildrenStubs();
78     //noinspection ForLoopReplaceableByForEach
79     for (int i = 0; i < childrenStubs.size(); i++) {
80       Stub child = childrenStubs.get(i);
81       enumerateStubsInto(child, result);
82     }
83   }
84
85   public void setDebugInfo(@NotNull String info) {
86     ObjectStubTree<?> ref = getStubTree(myRoot);
87     if (ref != null) {
88       assert ref == this;
89     }
90     myHasBackReference = ref != null;
91     myDebugInfo = info;
92   }
93
94   public static @Nullable ObjectStubTree getStubTree(@NotNull ObjectStubBase root) {
95     return root.getUserData(STUB_TO_TREE_REFERENCE);
96   }
97
98   public String getDebugInfo() {
99     return myHasBackReference ? myDebugInfo + "; with backReference" : myDebugInfo;
100   }
101
102   @Override
103   public String toString() {
104     return getClass().getSimpleName() + "{myDebugInfo='" + getDebugInfo() + '\'' + ", myRoot=" + myRoot + '}' + hashCode();
105   }
106
107   private static final class StubIndexSink implements IndexSink, TObjectProcedure<Map<Object, int[]>>, TObjectObjectProcedure<Object,int[]> {
108     private final THashMap<StubIndexKey<?, ?>, Map<Object, int[]>> myResult = new THashMap<>();
109     private final Function<StubIndexKey<?, ?>, TObjectHashingStrategy<?>> myHashingStrategyFunction;
110     private int myStubIdx;
111     private Map<Object, int[]> myProcessingMap;
112
113     private StubIndexSink(@NotNull Function<StubIndexKey<?, ?>, TObjectHashingStrategy<?>> hashingStrategyFunction) {
114       myHashingStrategyFunction = hashingStrategyFunction;
115     }
116
117     @Override
118     public void occurrence(final @NotNull StubIndexKey indexKey, final @NotNull Object value) {
119       Map<Object, int[]> map = myResult.get(indexKey);
120       if (map == null) {
121         //noinspection unchecked
122         map = new THashMap<>((TObjectHashingStrategy<Object>)myHashingStrategyFunction.apply(indexKey));
123         myResult.put(indexKey, map);
124       }
125
126       int[] list = map.get(value);
127       if (list == null) {
128         map.put(value, new int[] {myStubIdx});
129       }
130       else {
131         int lastNonZero = ArrayUtil.lastIndexOfNot(list, 0);
132         if (lastNonZero >= 0 && list[lastNonZero] == myStubIdx) {
133           // second and subsequent occurrence calls for the same value are no op
134           return;
135         }
136         int lastZero = lastNonZero + 1;
137
138         if (lastZero == list.length) {
139           list = ArrayUtil.realloc(list, Math.max(4, list.length << 1));
140           map.put(value, list);
141         }
142         list[lastZero] = myStubIdx;
143       }
144     }
145
146     public @NotNull Map<StubIndexKey<?, ?>, Map<Object, int[]>> getResult() {
147       myResult.forEachValue(this);
148       return myResult;
149     }
150
151     @Override
152     public boolean execute(Map<Object, int[]> object) {
153       myProcessingMap = object;
154       ((THashMap<Object, int[]>)object).forEachEntry(this);
155       return true;
156     }
157
158     @Override
159     public boolean execute(Object a, int[] b) {
160       if (b.length == 1) return true;
161       int firstZero = ArrayUtil.indexOf(b, 0);
162       if (firstZero != -1) {
163         int[] shorterList = ArrayUtil.realloc(b, firstZero);
164         myProcessingMap.put(a, shorterList);
165       }
166       return true;
167     }
168   }
169 }