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