sd-plugin/src/main/java/org/intellij/sdk/language/findUsages/SdFindUsagesHandler.java


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75

package org.intellij.sdk.language.findUsages;

import com.intellij.find.findUsages.FindUsagesHandler;
import com.intellij.find.findUsages.FindUsagesOptions;
import com.intellij.openapi.application.ReadAction;
import com.intellij.openapi.application.ReadActionProcessor;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiFile;
import com.intellij.psi.PsiReference;
import com.intellij.psi.search.GlobalSearchScope;
import com.intellij.psi.search.SearchScope;
import com.intellij.psi.search.searches.ReferencesSearch;
import com.intellij.usageView.UsageInfo;
import com.intellij.util.Processor;
import org.intellij.sdk.language.SdUtil;
import org.intellij.sdk.language.psi.SdFile;
import org.intellij.sdk.language.psi.SdFunctionDefinition;
import org.jetbrains.annotations.NotNull;

import java.util.HashMap;
import java.util.List;

public class SdFindUsagesHandler extends FindUsagesHandler {
    
    protected HashMap<String, List<PsiElement>> macrosMap;
    
    protected SdFindUsagesHandler(@NotNull PsiElement psiElement) {
        super(psiElement);
        PsiFile file = psiElement.getContainingFile();
        macrosMap = file instanceof SdFile ? SdUtil.createMacrosMap((SdFile) psiElement.getContainingFile()) : new HashMap<>();
    }
    
    @Override
    public boolean processElementUsages(@NotNull final PsiElement elementToSearch,
                                        @NotNull final Processor<? super UsageInfo> processor,
                                        @NotNull final FindUsagesOptions options) {
        final ReadActionProcessor<PsiReference> refProcessor = new ReadActionProcessor<>() {
            @Override
            public boolean processInReadAction(final PsiReference ref) {
                return processor.process(new UsageInfo(ref));
            }
        };
        
        final SearchScope scope = options.searchScope;
        
        final boolean searchText = options.isSearchForTextOccurrences && scope instanceof GlobalSearchScope;
        
        if (options.isUsages) {
            if (!(elementToSearch instanceof SdFunctionDefinition)) {
                boolean success =
                    ReferencesSearch.search(createSearchParameters(elementToSearch, scope, options)).forEach(refProcessor);
                if (!success) return false;
            } else {
                String macroName = ReadAction.compute( ((SdFunctionDefinition) elementToSearch)::getName);
                
                for (PsiElement macroImpl : macrosMap.get(macroName)) {
                    boolean success =
                        ReferencesSearch.search(createSearchParameters(macroImpl, scope, options)).forEach(refProcessor);
                    if (!success) return false;
                }
            }
        }
        if (searchText) {
            if (options.fastTrack != null) {
                options.fastTrack.searchCustom(consumer -> processUsagesInText(elementToSearch, processor, (GlobalSearchScope)scope));
            }
            else {
                return processUsagesInText(elementToSearch, processor, (GlobalSearchScope)scope);
            }
        }
        
        return true;
    }
    
}