Fix issues with color editing.

The color picker now auto-completes to color constants and works even if the previous version of a color is a color constant. flutter#5796 flutter#5780
jacob314 · Mar 30, 2022 · d10f38c · d10f38c
1 parent 0e65e6c
commit d10f38c
Show file tree

Hide file tree

Showing 4 changed files with 133 additions and 20 deletions.
diff --git a/flutter-idea/src/io/flutter/dart/DartPsiUtil.java b/flutter-idea/src/io/flutter/dart/DartPsiUtil.java
@@ -6,6 +6,7 @@
 package io.flutter.dart;
 
 import com.intellij.psi.PsiElement;
+import com.intellij.psi.tree.IElementType;
 import com.jetbrains.lang.dart.DartTokenTypes;
 import com.jetbrains.lang.dart.psi.DartArgumentList;
 import com.jetbrains.lang.dart.psi.DartArguments;
@@ -36,6 +37,17 @@ public static PsiElement getNewExprFromType(PsiElement element) {
     return parent;
   }
 
+  @Nullable
+  public static PsiElement getSurroundingNewOrCallExpression(PsiElement element) {
+    while (element != null) {
+      IElementType type = element.getNode().getElementType();
+      if (type == DartTokenTypes.NEW_EXPRESSION || type == DartTokenTypes.CALL_EXPRESSION) return element;
+      element = element.getParent();
+    }
+    return null;
+  }
+
+
   @Nullable
   public static String getValueOfPositionalArgument(@NotNull DartArguments arguments, int index) {
     final DartExpression expression = getPositionalArgument(arguments, index);
@@ -80,8 +92,13 @@ public static PsiElement getNamedArgumentExpression(@NotNull DartArguments argum
   @Nullable
   public static PsiElement topmostReferenceExpression(@NotNull PsiElement element) {
     final PsiElement id = element.getParent();
-    if (id == null || id.getNode().getElementType() != DartTokenTypes.ID) return null;
-    PsiElement refExpr = id.getParent();
+    if (id == null) return null;
+    PsiElement refExpr = null;
+    if (id.getNode().getElementType() == DartTokenTypes.ID) {
+      refExpr = id.getParent();
+    } else if (id.getNode().getElementType() == DartTokenTypes.REFERENCE_EXPRESSION) {
+      refExpr = id;
+    }
     if (refExpr == null || refExpr.getNode().getElementType() != DartTokenTypes.REFERENCE_EXPRESSION) return null;
 
     PsiElement parent = refExpr.getParent();

diff --git a/flutter-idea/src/io/flutter/editor/ColorField.java b/flutter-idea/src/io/flutter/editor/ColorField.java
@@ -109,17 +109,7 @@ private static String buildColorExpression(@Nullable Color color) {
     if (color == null) {
       return "";
     }
-
-    final String flutterColorName = FlutterColors.getColorName(color);
-    if (flutterColorName != null) {
-      // TODO(jacobr): only apply this conversion if the material library is already imported in the
-      // library being edited. We also need to be able to handle cases where the material library is
-      // imported with a prefix.
-      return "Colors." + flutterColorName;
-    }
-
-    return String.format(
-      "Color(0x%02x%02x%02x%02x)", color.getAlpha(), color.getRed(), color.getGreen(), color.getBlue());
+    return FlutterColors.buildColorExpression(color);
   }
 
   public void addTextFieldListeners(String name, JBTextField field) {

diff --git a/flutter-idea/src/io/flutter/editor/FlutterColorProvider.java b/flutter-idea/src/io/flutter/editor/FlutterColorProvider.java
@@ -18,7 +18,9 @@
 import com.jetbrains.lang.dart.psi.DartArguments;
 import com.jetbrains.lang.dart.psi.DartExpression;
 import com.jetbrains.lang.dart.psi.DartLiteralExpression;
+import com.jetbrains.lang.dart.util.DartElementGenerator;
 import io.flutter.FlutterBundle;
+import io.flutter.dart.DartPsiUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -29,18 +31,73 @@
 import static io.flutter.dart.DartPsiUtil.getNewExprFromType;
 import static io.flutter.dart.DartPsiUtil.topmostReferenceExpression;
 
+// This class is required to resolve
+// https://github.com/flutter/flutter-intellij/issues/5796
+// TODO(jacobr): track down a possible bug in Dart PsiElement implementation
+// that results in Dart PsiElement objects showing up as equal when they
+// are not safe to treat as equal for purposes of resolving PsiElement
+// markers. If that issue can be resolved then this hack can be removed.
+/**
+ * Color class that enables creating colors that are visually identical but
+ * are only equal if the colors have the same associated PsiElement.
+ * <p>
+ * This class is used as a hack to avoid a bug where Dart color icon markers
+ * can fail to update after making small code changes that are enough to
+ * invalidate the PsiElement but not enough to break equality checks for the
+ * PsiElement objects.
+ */
+class PsiElementColor extends Color {
+
+  PsiElementColor(int r, int g, int b, int a, PsiElement psiElement) {
+    super(r, g, b, a);
+    this.psiElement = psiElement;
+  }
+
+  private PsiElement psiElement;
+
+  public boolean equals(Object obj) {
+    if (!(obj instanceof PsiElementColor)) return false;
+    PsiElementColor other = (PsiElementColor)obj;
+    return other.getRGB() == this.getRGB() && other.psiElement == psiElement;
+  }
+}
+
 public class FlutterColorProvider implements ElementColorProvider {
 
+  /**
+   * When we replace the target PsiElement as part of tweaking a color, we
+   * continue to get back the old PsiElement even after it is invalid.
+   * To handle this case we track the orignal and replacement elements so
+   * that we can swap elements as needed.
+   */
+  PsiElement originalElement;
+  PsiElement replacementElement;
+
   @Nullable
   @Override
   public Color getColorFrom(@NotNull PsiElement element) {
+    final Color color = getColorFromHelper(element);
+    if (color == null) return null;
+    return new PsiElementColor(color.getRed(), color.getGreen(), color.getBlue(), color.getAlpha(), element);
+  }
+
+  void cleanupInvalidCache() {
+    if (replacementElement != null && !replacementElement.isValid()) {
+      originalElement = null;
+      replacementElement = null;
+    }
+  }
+
+  public Color getColorFromHelper(@NotNull PsiElement element) {
+    cleanupInvalidCache();
     // This must return null for non-leaf nodes and any language other than Dart.
     if (element.getNode().getElementType() != DartTokenTypes.IDENTIFIER) return null;
+    if (element.getFirstChild() != null) return null;
 
     final String name = element.getText();
     if (!(name.equals("Colors") || name.equals("CupertinoColors") || name.equals("Color"))) return null;
 
-    final PsiElement refExpr = topmostReferenceExpression(element);
+    final PsiElement refExpr = DartPsiUtil.topmostReferenceExpression(element);
     if (refExpr == null) return null;
     PsiElement parent = refExpr.getParent();
     if (parent == null) return null;
@@ -123,17 +180,38 @@ private Color parseColorText(@NotNull String text, @NotNull String platform) {
   }
 
   @Override
-  public void setColorTo(@NotNull PsiElement element, @NotNull Color color) {
-    // Not trying to look up Material or Cupertino colors.
-    // Unfortunately, there is no way to prevent the color picker from showing (if clicked) for those expressions.
-    if (!element.getText().equals("Color")) return;
+  public void setColorTo(@NotNull PsiElement targetElement, @NotNull Color color) {
+    cleanupInvalidCache();
+    PsiElement element;
+    if (!targetElement.isValid() && originalElement == targetElement && replacementElement != null) {
+      element = replacementElement;
+    }
+    else {
+      element = targetElement;
+    }
+    if (!element.isValid()) return;
+
     final Document document = PsiDocumentManager.getInstance(element.getProject()).getDocument(element.getContainingFile());
     final Runnable command = () -> {
       final PsiElement refExpr = topmostReferenceExpression(element);
       if (refExpr == null) return;
       PsiElement parent = refExpr.getParent();
       if (parent == null) return;
-      if (parent.getNode().getElementType() == DartTokenTypes.CALL_EXPRESSION) {
+      final PsiElement targetForReplacement = getTargetForReplacement(element, refExpr);
+      if (!element.getText().equals("Color") || FlutterColors.getColorName(color) != null) {
+        // Generate an expression for the Color from scratch rather than
+        // reusing the existing Color constructor expression because the
+        // existing expression is not a color constructor call or there is a
+        // Flutter color name that matches the exact color value.
+        final String colorExpression = FlutterColors.buildColorExpression(color);
+        final PsiFileFactoryImpl factory = new PsiFileFactoryImpl(element.getManager());
+        PsiElement newPsi = DartElementGenerator.createExpressionFromText(element.getProject(), colorExpression);
+        if (newPsi != null) {
+          originalElement = targetElement;
+          replacementElement = targetForReplacement.replace(newPsi).getFirstChild().getFirstChild().getFirstChild();
+        }
+      }
+      else if (parent.getNode().getElementType() == DartTokenTypes.CALL_EXPRESSION) {
         // foo(Color.fromRGBO(0, 255, 0, 0.5))
         replaceColor(parent, refExpr, color);
       }
@@ -149,6 +227,21 @@ else if (parent.getNode().getElementType() == DartTokenTypes.SIMPLE_TYPE) {
       .executeCommand(element.getProject(), command, FlutterBundle.message("change.color.command.text"), null, document);
   }
 
+  private PsiElement getTargetForReplacement(PsiElement element, PsiElement refExpr) {
+    final String name = element.getText();
+    if (element.getText().equals("Color")) {
+      // Parent of refExpr is either a CALL_EXPRESSION or SIMPLE_TYPE that creates the instance of the Color object.
+      return DartPsiUtil.getSurroundingNewOrCallExpression(refExpr);
+    }
+    else {
+      PsiElement parent = refExpr.getParent();
+      // Handle Colors.grey[300] case.
+      if (parent != null && parent.getNode().getElementType() == DartTokenTypes.ARRAY_ACCESS_EXPRESSION) return parent;
+      // Colors.grey and Colors.grey.shade300 case.
+      return refExpr;
+    }
+  }
+
   private void replaceColor(@NotNull PsiElement parent, @NotNull PsiElement refExpr, Color color) {
     final PsiElement selectorNode = refExpr.getLastChild();
     if (selectorNode == null) return;

diff --git a/flutter-idea/src/io/flutter/editor/FlutterColors.java b/flutter-idea/src/io/flutter/editor/FlutterColors.java
@@ -93,7 +93,7 @@ else if (colors.containsKey(key + primarySuffix)) {
   }
 
   /**
-   * Returns the the shortest material color name matching a color if one exists.
+   * Returns the shortest material color name matching a color if one exists.
    */
   @Nullable
   public static String getColorName(@Nullable Color color) {
@@ -133,4 +133,17 @@ private static Color getColorValue(String name) {
     final String hexValue = colors.getProperty(name);
     return parseColor(hexValue);
   }
+
+  public static String buildColorExpression(Color color) {
+    final String flutterColorName = FlutterColors.getColorName(color);
+    if (flutterColorName != null) {
+      // TODO(jacobr): only apply this conversion if the material library is already imported in the
+      // library being edited. We also need to be able to handle cases where the material library is
+      // imported with a prefix.
+      return "Colors." + flutterColorName;
+    }
+
+    return String.format(
+      "Color(0x%02x%02x%02x%02x)", color.getAlpha(), color.getRed(), color.getGreen(), color.getBlue());
+  }
 }