11package org.kdb.inside.brains.lang;
22
33import com.intellij.codeInsight.daemon.RainbowVisitor;
4- import com.intellij.codeInsight.daemon.impl.HighlightInfo;
54import com.intellij.codeInsight.daemon.impl.HighlightVisitor;
5+ import com.intellij.openapi.editor.colors.TextAttributesKey;
66import com.intellij.psi.PsiElement;
77import com.intellij.psi.PsiFile;
88import com.intellij.psi.PsiReference;
9+ import com.intellij.psi.tree.IElementType;
910import org.jetbrains.annotations.NotNull;
10- import org.jetbrains.annotations.Nullable;
1111import org.kdb.inside.brains.QFileType;
1212import org.kdb.inside.brains.psi.*;
1313
14+ import java.util.function.Predicate;
15+
1416public class QRainbowVisitor extends RainbowVisitor {
1517 @Override
1618 public boolean suitableForFile(@NotNull PsiFile psiFile) {
@@ -20,59 +22,96 @@ public boolean suitableForFile(@NotNull PsiFile psiFile) {
2022 @Override
2123 public void visit(@NotNull PsiElement psiElement) {
2224 if (psiElement instanceof @NotNull QVariable var) {
23- final HighlightInfo info = createVariableInfo(var);
24- if (info != null) {
25- addInfo(info);
25+ createVariableInfo(var);
26+ } else {
27+ final IElementType type = psiElement.getNode().getElementType();
28+ if (type == QTypes.BRACE_OPEN || type == QTypes.BRACE_CLOSE) {
29+ createBraceInfo(psiElement);
30+ } else if (type == QTypes.PAREN_OPEN || type == QTypes.PAREN_CLOSE) {
31+ createParenInfo(psiElement);
32+ } else if (type == QTypes.BRACKET_OPEN || type == QTypes.BRACKET_CLOSE) {
33+ createBracketInfo(psiElement);
2634 }
2735 }
2836 }
2937
30- private @Nullable HighlightInfo createVariableInfo(@NotNull QVariable var) {
38+ private void createBraceInfo(@NotNull PsiElement psiElement) {
39+ createBracketsInfo(psiElement, "q_rainbow_brace", QSyntaxHighlighter.BRACES, e -> e instanceof QLambdaExpr);
40+ }
41+
42+ private void createParenInfo(@NotNull PsiElement psiElement) {
43+ createBracketsInfo(psiElement, "q_rainbow_paren", QSyntaxHighlighter.PARENTHESES, e -> e instanceof QParenthesesExpr);
44+ }
45+
46+ private void createBracketInfo(@NotNull PsiElement psiElement) {
47+ createBracketsInfo(psiElement, "q_rainbow_bracket", QSyntaxHighlighter.BRACKETS, e -> e instanceof QConditionExpr || e instanceof QArguments || e instanceof QControlExpr || e instanceof QParameters || e instanceof QGroupingExpr);
48+ }
49+
50+ private void createVariableInfo(@NotNull QVariable var) {
3151 final QLambda lambda = var.getContext(QLambda.class);
3252 if (lambda == null) {
33- return null;
34- }
35-
36- if (var instanceof QVarDeclaration d) {
37- return getInfo(var, d, lambda);
53+ return;
3854 }
3955
4056 if (lambda.isImplicitDeclaration(var)) {
41- return createInfo (lambda, var);
42- }
43-
44- final PsiReference[] references = var.getReferences();
45- for (PsiReference reference : references) {
46- final PsiElement resolve = reference.resolve();
47- if (resolve instanceof QVarDeclaration d) {
48- final HighlightInfo info = getInfo(var, d, lambda);
49- if (info != null) {
50- return info ;
57+ createVariableInfo (lambda, var);
58+ } else if (var instanceof QVarDeclaration d) {
59+ createDeclarationInfo(lambda, d, var);
60+ } else {
61+ final PsiReference[] references = var.getReferences();
62+ for (PsiReference reference : references) {
63+ final PsiElement resolve = reference.resolve();
64+ if (resolve instanceof QVarDeclaration d) {
65+ createDeclarationInfo(lambda, d, var);
66+ break ;
5167 }
5268 }
5369 }
54- return null;
5570 }
5671
57- private @Nullable HighlightInfo getInfo (@NotNull QVariable var , @NotNull QVarDeclaration dec, @NotNull QLambda lambda ) {
72+ private void createDeclarationInfo (@NotNull QLambda lambda , @NotNull QVarDeclaration dec, @NotNull QVariable var ) {
5873 if (QPsiUtil.isGlobalDeclaration(dec)) {
59- return null ;
74+ return;
6075 }
6176
6277 final ElementScope scope = dec.getVariableContext().getScope();
6378 // only root variables, no tables, dicts and so on
6479 if (scope == ElementScope.PARAMETERS || scope == ElementScope.LAMBDA) {
65- return createInfo(lambda, var);
80+ createVariableInfo(lambda, var);
81+ }
82+ }
83+
84+ private void createVariableInfo(@NotNull QLambda lambda, @NotNull QVariable var) {
85+ addInfo(getInfo(lambda, var, var.getName(), QSyntaxHighlighter.VARIABLE));
86+ }
87+
88+ private void createBracketsInfo(PsiElement psiElement, String name, TextAttributesKey key, Predicate<PsiElement> predicate) {
89+ final DepthInfo info = calculateDepth(psiElement, predicate);
90+ if (info.context == null || info.depth < 0) {
91+ return;
92+ }
93+ addInfo(getInfo(info.context, psiElement, name + info.depth, key));
94+ }
95+
96+ private DepthInfo calculateDepth(@NotNull PsiElement psiElement, Predicate<PsiElement> predicate) {
97+ int depth = 0;
98+ PsiElement context = null;
99+ PsiElement parent = psiElement.getParent();
100+ while (parent != null) {
101+ if (predicate.test(parent)) {
102+ context = parent;
103+ depth++;
104+ }
105+ parent = parent.getParent();
66106 }
67- return null ;
107+ return new DepthInfo(context, depth - 1) ;
68108 }
69109
70110 @Override
71111 public @NotNull HighlightVisitor clone() {
72112 return new QRainbowVisitor();
73113 }
74114
75- private @NotNull HighlightInfo createInfo(@NotNull QLambda lambda, @NotNull QVariable var) {
76- return getInfo(lambda, var, var.getName(), QSyntaxHighlighter.VARIABLE);
115+ private record DepthInfo(PsiElement context, int depth) {
77116 }
78117}
0 commit comments