CodeContextProvider.java

CodeContextProvider.java

import com.intellij.openapi.editor.Document;
import com.intellij.openapi.project.Project;
import com.intellij.psi.PsiDocumentManager;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiFile;
import com.intellij.psi.PsiManager;
import com.intellij.psi.PsiReference;
import com.intellij.psi.search.searches.ReferencesSearch;
import com.intellij.openapi.vfs.VirtualFile;
import com.intellij.openapi.vfs.LocalFileSystem;
import com.intellij.psi.util.PsiTreeUtil;
import com.tencent.code.intel.bean.item.Position;
import com.tencent.code.intel.bean.item.TextDocumentItem;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;

import org.jetbrains.annotations.NotNull;

public class LanguageProvider implements CodeContextProvider {

    private Project project;

    public LanguageProvider(Project project) {
        this.project = project;
    }

    @Override
    public List<TextDocumentItem> getDefinitions(String uri, Position position) {
        return getTextDocumentItems(uri, position);
    }

    @Override
    public List<TextDocumentItem> getImplementations(String uri, Position position) {
        return getTextDocumentItems(uri, position);
    }

    @NotNull
    private List<TextDocumentItem> getTextDocumentItems(String uri, Position position) {
        PsiElement element = findElementAtPosition(uri, position);
        if (element == null) {
            return List.of();
        }

        if ((element instanceof PsiReference)) {
            PsiReference elementReference = (PsiReference) element;
            element = elementReference.resolve();
        }

        return List.of(createTextDocumentItem(element));
    }

    @Override
    public List<TextDocumentItem> getReferences(String uri, Position position) {
        PsiElement element = findElementAtPosition(uri, position);
        if (element == null) {
            return List.of();
        }

        return findReferences(element);
    }

    private PsiElement findElementAtPosition(String uri, Position position) {
        VirtualFile virtualFile = getVirtualFileFromUri(uri);
        if (virtualFile == null) {
            return null;
        }

        PsiFile psiFile = PsiManager.getInstance(project).findFile(virtualFile);
        if (psiFile == null) {
            return null;
        }

        int offset = getOffset(psiFile, position);
        PsiElement elementAt = psiFile.findElementAt(offset);
        if (elementAt == null) {
            return null;
        }

        return PsiTreeUtil.getParentOfType(elementAt, PsiElement.class);
    }

    private List<TextDocumentItem> findReferences(PsiElement element) {
        Collection<PsiReference> references = ReferencesSearch.search(element).findAll();
        List<TextDocumentItem> items = new ArrayList<>();
        for (PsiReference reference : references) {
            PsiElement usage = reference.getElement();
            items.add(createTextDocumentItem(usage));
        }
        return items;
    }

    private int getOffset(PsiFile psiFile, Position position) {
        Document document = PsiDocumentManager.getInstance(project).getDocument(psiFile);
        if (document != null) {
            return document.getLineStartOffset(position.getLine()) + position.getCharacter();
        }
        return 0;
    }

    private VirtualFile getVirtualFileFromUri(String uri) {
        if (uri.startsWith("file://")) {
            String path = uri.substring("file://".length());
            return LocalFileSystem.getInstance().findFileByPath(path);
        }
        return null;
    }

    private TextDocumentItem createTextDocumentItem(PsiElement element) {
        TextDocumentItem documentItem = new TextDocumentItem();
        PsiFile psiFile = element.getContainingFile();
        documentItem.setUri(psiFile.getVirtualFile().getPath());
        documentItem.setLanguageId(psiFile.getLanguage().getID());
        documentItem.setText(psiFile.getText());
        return documentItem;
    }
}