ElementColorProvider#setColorTo only works once (with same Color-Picker open) when replacing parent element

Hello,

I'm currently working on a custom Language-Plugin with utilizing the ElementColorProvider.

Everything works fine so far, except that when I'm trying to change the Color using the Color-Picker: It works the first time but not a second time while using the same Color-Picker-Instance (not reopening the picker). I think that is due to my usage of the ElementColorProvider:

As LineMarkers should/must be attached to a LeafNode, I'm using the lowest possible Node and traverse up to get the full context (e.g. “rgba(12, 12, 12, 0.5)” - whereas 12 is the LeafNode). The consequence of that is, that I'm modifying the parent node of the leaf node which invalidates further calls to getParent (to access the full context). Is there any way to mitigate that problem?

A snippet of the affected Psi-Tree which represents a color:

LssPropertyImpl(PROPERTY)(17,58)
      PsiElement(LssTokenType.KEY)('background-color')(17,33)
      PsiElement(LssTokenType.SEPARATOR)(':')(33,34)
      PsiWhiteSpace(' ')(34,35)
      LssValueImpl(VALUE)(35,57)
        LssRawValueImpl(RAW_VALUE)(35,57)
          LssFuncValueImpl(FUNC_VALUE)(35,57)
            PsiElement(LssTokenType.FUNC_NAME)('rgba')(35,39)
            PsiElement(LssTokenType.FUNCTION_LPARANTHESIS)('(')(39,40)
            LssRawValueImpl(RAW_VALUE)(40,43)
              LssFixedValueImpl(FIXED_VALUE)(40,43)
                PsiElement(LssTokenType.FIXED_VALUE_TYPE)('194')(40,43)
            PsiElement(LssTokenType.FUNCTION_PARAMETER_SEPARATOR)(',')(43,44)
            PsiWhiteSpace(' ')(44,45)
            LssRawValueImpl(RAW_VALUE)(45,47)
              LssFixedValueImpl(FIXED_VALUE)(45,47)
                PsiElement(LssTokenType.FIXED_VALUE_TYPE)('50')(45,47)
            PsiElement(LssTokenType.FUNCTION_PARAMETER_SEPARATOR)(',')(47,48)
            PsiWhiteSpace(' ')(48,49)
            LssRawValueImpl(RAW_VALUE)(49,51)
              LssFixedValueImpl(FIXED_VALUE)(49,51)
                PsiElement(LssTokenType.FIXED_VALUE_TYPE)('50')(49,51)
            PsiElement(LssTokenType.FUNCTION_PARAMETER_SEPARATOR)(',')(51,52)
            PsiWhiteSpace(' ')(52,53)
            LssRawValueImpl(RAW_VALUE)(53,56)
              LssFixedValueImpl(FIXED_VALUE)(53,56)
                PsiElement(LssTokenType.FIXED_VALUE_TYPE)('191')(53,56)
            PsiElement(LssTokenType.FUNCTION_RPARANTHESIS)(')')(56,57)
      PsiElement(LssTokenType.SEMICOLON)(';')(57,58)

In that particular example I'm attaching the Color to the first leaf node of the complex Value-Node (PsiElement(LssTokenType.FIXED_VALUE_TYPE)('194')(40,43)), traverse up to the LssProperty, find the property key and value and use those for further processing. When using setColor I'm replacing the whole LssFuncValue which breaks another traverse up. 

In case it helps, the full code:

public class ElementColorProvider implements com.intellij.openapi.editor.ElementColorProvider {

    @Override
    public @Nullable Color getColorFrom(@NotNull PsiElement element) {
        if (element.getNode().getElementType().equals(LssTypes.FIXED_VALUE_TYPE) &&
                element.getParent().getNode().getElementType().equals(LssTypes.FIXED_VALUE)) {
            Color color = getColor((LssFixedValue) element.getParent());
            if (color != null) {
                return new JBColor(color, color);
            }
        }
        return null;
    }

    @Override
    public void setColorTo(@NotNull PsiElement element, @NotNull Color color) {
        // element is always a FIXED_VALUE - we need to determine the root argument which represents the properties color
        PsiElement declarationElement = getDeclarationElement(element);
        if (declarationElement == null) {
            System.out.println("DeclEl null - " + element);
            return;
        }
        if (declarationElement.getNode().getElementType().equals(LssTypes.PROPERTY)) {
            LssProperty property = (LssProperty) declarationElement;
            // We can't have complex property values for colors
            if (property.getValue().getComplexPropertyValue() != null) {
                return;
            }
            IValueElement toReplace;
            LssRawValue rawValue = property.getValue().getRawValue();
            assert rawValue != null : "Neither a complex or raw value";
            if (rawValue.getFuncValue() != null) {
                LssFuncValue funcValue = rawValue.getFuncValue();
                if (funcValue.getFunctionIdentifier().equals("var")) {
                    // TODO: update declared variable or replace var call with fixed value?
                    return;
                }
                toReplace = rawValue.getFuncValue();
            } else {
                toReplace = rawValue.getFixedValue();
            }
            assert toReplace != null : "No element for replacement found";
            Document document = PsiDocumentManager.getInstance(element.getProject()).getDocument(element.getContainingFile());
            if (color.getAlpha() != 255) {
                CommandProcessor.getInstance().executeCommand(
                        element.getProject(),
                        () -> toReplace.replace(LssElementFactory.createFuncValue(
                                element.getProject(),
                                "rgba",
                                color.getRed(), color.getGreen(), color.getBlue(), color.getAlpha()
                        )),
                        JavaBundle.message("change.color.command.text"),
                        null, document
                );
                return;
            }
            CommandProcessor.getInstance().executeCommand(
                    element.getProject(),
                    () -> toReplace.replace(LssElementFactory.createFuncValue(
                            element.getProject(),
                            "rgb",
                            color.getRed(), color.getGreen(), color.getBlue()
                    )),
                    JavaBundle.message("change.color.command.text"),
                    null, document
            );
        }
    }

    private PsiElement getDeclarationElement(PsiElement element) {
        return PsiTreeUtil.findFirstParent(element, psiElement ->
                psiElement.getNode().getElementType().equals(LssTypes.PROPERTY) ||
                        psiElement.getNode().getElementType().equals(LssTypes.VAR_DECLARATION));
    }

    private @Nullable Color getColor(LssFixedValue value) {
        PsiElement root = getDeclarationElement(value);
        if (root == null) {
            return null;
        }
        if (root.getNode().getElementType().equals(LssTypes.PROPERTY)) {
            return getColor((LssProperty) root, value);
        }
        return getColor((LssVarDeclaration) root, value);
    }

    private @Nullable Color getColor(LssProperty property, LssFixedValue triggeredBy) {
        if (!isFirstInValueList(property.getValue(), triggeredBy)) {
            return null;
        }
        UnvaluedLssProperty<?> knownProperty = ApplicationManager.getApplication().getService(PropertyRegistry.class).getKnownProperty(property.getFirstChild().getText());
        return getColor(knownProperty, property.getValue());
    }

    private @Nullable Color getColor(LssVarDeclaration declaration, LssFixedValue triggeredBy) {
        return null; // TODO
    }

    private @Nullable Color getColor(@Nullable UnvaluedLssProperty<?> property, @NotNull IValueElement element) {
        if (property == null || property.valueType() != de.pierreschwang.labymod.language.lss.property.types.Color.class) {
            return null;
        }
        try {
            ProcessedObject<de.pierreschwang.labymod.language.lss.property.types.Color>[] processedObjects = element.computeValue(((UnvaluedLssProperty<de.pierreschwang.labymod.language.lss.property.types.Color>) property));
            if (processedObjects == null || processedObjects.length == 0 || processedObjects[0].value() == null) {
                return null;
            }

            de.pierreschwang.labymod.language.lss.property.types.Color value = processedObjects[0].value();
            if (value == null) {
                return null;
            }
            return value.nativeColor();
        } catch (Exception e) {
            Logger.getInstance(ElementColorProvider.class).error(e);
            return null;
        }
    }

    private boolean isFirstInValueList(LssValue value, LssFixedValue fixedValue) {
        if (value.getComplexPropertyValue() != null) {
            return isFirstOfComplexPropertyValue(value.getComplexPropertyValue(), fixedValue);
        }
        return isFirstOfRawPropertyValue(Objects.requireNonNull(value.getRawValue()), fixedValue);
    }

    private boolean isFirstOfComplexPropertyValue(LssComplexPropertyValue value, LssFixedValue fixedValue) {
        return isFirstOfRawPropertyValue(value.getRawValueList().get(0), fixedValue);
    }

    private boolean isFirstOfRawPropertyValue(LssRawValue value, LssFixedValue fixedValue) {
        if (value.getFuncValue() != null) {
            return isFirstOfFunctionValue(value.getFuncValue(), fixedValue);
        }
        if (value.getFixedValue() == null) {
            return false;
        }
        return value.getFixedValue().getTextOffset() == fixedValue.getTextOffset();
    }

    private boolean isFirstOfFunctionValue(LssFuncValue value, LssFixedValue fixedValue) {
        return isFirstOfRawPropertyValue(value.getRawValueList().get(0), fixedValue);
    }

}