diff --git a/src/main/kotlin/insight/ColorLineMarkerProvider.kt b/src/main/kotlin/insight/ColorLineMarkerProvider.kt index d5c323d9a..da06f79ed 100644 --- a/src/main/kotlin/insight/ColorLineMarkerProvider.kt +++ b/src/main/kotlin/insight/ColorLineMarkerProvider.kt @@ -118,7 +118,7 @@ class ColorLineMarkerProvider : LineMarkerProvider { element, color, GutterIconNavigationHandler handler@{ _, _ -> - if (!element.isWritable) { + if (!element.isWritable || !element.isValid) { return@handler } diff --git a/src/main/kotlin/insight/ColorUtil.kt b/src/main/kotlin/insight/ColorUtil.kt index aa98533e4..2dee99a8a 100644 --- a/src/main/kotlin/insight/ColorUtil.kt +++ b/src/main/kotlin/insight/ColorUtil.kt @@ -24,20 +24,38 @@ import kotlin.math.round import org.jetbrains.uast.UCallExpression import org.jetbrains.uast.UElement import org.jetbrains.uast.UExpression +import org.jetbrains.uast.UField import org.jetbrains.uast.UIdentifier import org.jetbrains.uast.ULiteralExpression import org.jetbrains.uast.UQualifiedReferenceExpression import org.jetbrains.uast.UReferenceExpression +import org.jetbrains.uast.UResolvable import org.jetbrains.uast.generate.generationPlugin import org.jetbrains.uast.generate.replace +import org.jetbrains.uast.resolveToUElement fun UIdentifier.findColor(function: (Map, Map.Entry) -> T): T? { val parent = this.uastParent val expression = parent as? UReferenceExpression ?: return null - val type = expression.getExpressionType() ?: return null + return findColorFromExpression(expression, function) +} + +private fun findColorFromExpression( + expression: UReferenceExpression, + function: (Map, Map.Entry) -> T +): T? { + val referencedElement = expression.resolveToUElement() + if (referencedElement is UField) { + val referencedFieldInitializer = referencedElement.uastInitializer + if (referencedFieldInitializer is UReferenceExpression) { + return findColorFromExpression(referencedFieldInitializer, function) + } + } - val module = this.sourcePsi?.findModule() ?: return null + val type = expression.getExpressionType() ?: return null + val module = expression.sourcePsi?.findModule() ?: return null val facet = MinecraftFacet.getInstance(module) ?: return null + val resolvedName = expression.resolvedName ?: return null for (abstractModuleType in facet.types) { val map = abstractModuleType.classToColorMappings for (entry in map.entries) { @@ -47,7 +65,7 @@ fun UIdentifier.findColor(function: (Map, Map.Entry key.startsWith(colorClass) }, entry) } } @@ -61,21 +79,50 @@ fun UIdentifier.findColor( vectorClasses: Array? ): Pair? { val sourcePsi = this.sourcePsi ?: return null - val project = sourcePsi.project - val module = ModuleUtilCore.findModuleForPsiElement(sourcePsi) ?: return null - val facet = MinecraftFacet.getInstance(module) ?: return null - if (!facet.isOfType(moduleType)) { return null } - val methodExpression = uastParent as? UCallExpression ?: return null - if (methodExpression.resolve()?.containingClass?.qualifiedName != className) { - return null + val methodExpression = uastParent as? UCallExpression + if (methodExpression?.resolve()?.containingClass?.qualifiedName == className) { + return findColorFromCallExpression(methodExpression, vectorClasses) } + var referencedElement = (uastParent as? UResolvable)?.resolveToUElement() + while (referencedElement is UField) { + val referencedFieldInitializer: UExpression? = referencedElement.uastInitializer + if (referencedFieldInitializer is UCallExpression) { + // The field is initialized with a method call + return referencedFieldInitializer.methodIdentifier?.findColor(moduleType, className, vectorClasses) + } + + if (referencedFieldInitializer is UReferenceExpression) { + // The field is probably initialized with a reference to another field + val referenceNameElement = referencedFieldInitializer.referenceNameElement + if (referenceNameElement is UIdentifier) { + // The expression was simple enough + return referenceNameElement.findColor(moduleType, className, vectorClasses) + } else if (referenceNameElement is UResolvable) { + // The expression is complex, so we resolve it. If it is a field we're on for another round + referencedElement = referenceNameElement.resolveToUElement() + continue + } + } + + break + } + + return null +} + +private fun findColorFromCallExpression( + methodExpression: UCallExpression, + vectorClasses: Array? +): Pair? { + val project = methodExpression.sourcePsi?.project ?: return null + val arguments = methodExpression.valueArguments val types = arguments.map(UExpression::getExpressionType)