introduce ActionWrapperUtil

Add the ability to customize data contexts for delegates.

Session API is not used in case of customization. For now, session always uses the original data context.

GitOrigin-RevId: 4c4b7dcfe55d41ab708cba6a415272cecd8b8eb6

platform/execution-impl/api-dump-unreviewed.txt

@@ -1443,9 +1443,11 @@ f:com.intellij.execution.impl.statistics.RunConfigurationUsageTriggerCollector$R
 f:com.intellij.execution.lineMarker.ExecutorAction
 - com.intellij.openapi.actionSystem.ActionGroup
 - com.intellij.openapi.actionSystem.ActionWithDelegate
+- com.intellij.openapi.actionSystem.DataSnapshotProvider
 - sf:Companion:com.intellij.execution.lineMarker.ExecutorAction$Companion
 - b:<init>(com.intellij.openapi.actionSystem.AnAction,com.intellij.execution.Executor,I,kotlin.jvm.internal.DefaultConstructorMarker):V
 - actionPerformed(com.intellij.openapi.actionSystem.AnActionEvent):V
+- dataSnapshot(com.intellij.openapi.actionSystem.DataSink):V
 - equals(java.lang.Object):Z
 - sf:getActionList():java.util.List
 - sf:getActionList(I):java.util.List
@@ -1473,10 +1475,12 @@ c:com.intellij.execution.lineMarker.LineMarkerActionWrapper
 - com.intellij.openapi.actionSystem.ActionGroup
 - com.intellij.codeInsight.intention.PriorityAction
 - com.intellij.openapi.actionSystem.ActionWithDelegate
+- com.intellij.openapi.actionSystem.DataSnapshotProvider
 - sf:LOCATION_WRAPPER:com.intellij.openapi.util.Key
 - pf:myElement:com.intellij.psi.SmartPsiElementPointer
 - <init>(com.intellij.psi.PsiElement,com.intellij.openapi.actionSystem.AnAction):V
 - actionPerformed(com.intellij.openapi.actionSystem.AnActionEvent):V
+- dataSnapshot(com.intellij.openapi.actionSystem.DataSink):V
 - getActionUpdateThread():com.intellij.openapi.actionSystem.ActionUpdateThread
 - getChildren(com.intellij.openapi.actionSystem.AnActionEvent):com.intellij.openapi.actionSystem.AnAction[]
 - getDelegate():com.intellij.openapi.actionSystem.AnAction

platform/execution-impl/src/com/intellij/execution/lineMarker/ExecutorAction.kt

@@ -13,7 +13,7 @@ import org.jetbrains.annotations.ApiStatus
 class ExecutorAction private constructor(val origin: AnAction,
                                          val executor: Executor,
                                          val order: Int) :
-  ActionGroup(), ActionWithDelegate<AnAction> {
+  ActionGroup(), DataSnapshotProvider, ActionWithDelegate<AnAction> {
   init {
     copyFrom(origin)
     if (origin !is ActionGroup) {
@@ -22,14 +22,18 @@ class ExecutorAction private constructor(val origin: AnAction,
     }
   }
 
-  override fun getActionUpdateThread() = origin.actionUpdateThread
-
-  override fun getDelegate(): AnAction {
-    return origin
-  }
+  override fun getDelegate(): AnAction = origin
 
   override fun isDumbAware() = origin.isDumbAware
 
+  override fun getActionUpdateThread() = ActionWrapperUtil.getActionUpdateThread(this, origin)
+
+  override fun dataSnapshot(sink: DataSink) {
+    if (order != 0) {
+      sink[orderKey] = order
+    }
+  }
+
   companion object {
 
     @JvmStatic
@@ -52,18 +56,10 @@ class ExecutorAction private constructor(val origin: AnAction,
             ExecutorAction(it, executor, order)
           }
         }
-      if (createAction != null) {
-        result.add(object : ActionGroupWrapper(createAction as ActionGroup) {
-          override fun update(e: AnActionEvent) {
-            delegate.update(wrapEvent(e, order))
-          }
-
-          override fun actionPerformed(e: AnActionEvent) {
-            delegate.actionPerformed(wrapEvent(e, order))
-          }
-
-          override fun getChildren(e: AnActionEvent?): Array<AnAction> {
-            return delegate.getChildren(e?.let { wrapEvent(e, order) })
+      if (createAction is ActionGroup) {
+        result.add(object : ActionGroupWrapper(createAction), DataSnapshotProvider {
+          override fun dataSnapshot(sink: DataSink) {
+            sink[orderKey] = order
           }
 
           override fun equals(other: Any?): Boolean {
@@ -78,13 +74,6 @@ class ExecutorAction private constructor(val origin: AnAction,
       return result
     }
 
-    private fun wrapEvent(e: AnActionEvent, order: Int): AnActionEvent {
-      return if (order == 0) e
-      else e.withDataContext(CustomizedDataContext.withSnapshot(e.dataContext) { sink ->
-        sink[orderKey] = order
-      })
-    }
-
     @ApiStatus.Internal
     @JvmStatic
     fun wrap(runContextAction: AnAction, executor: Executor, order: Int): AnAction {
@@ -93,19 +82,19 @@ class ExecutorAction private constructor(val origin: AnAction,
   }
 
   override fun update(e: AnActionEvent) {
-    origin.update(wrapEvent(e, order))
+    ActionWrapperUtil.update(e, this, origin)
   }
 
   override fun getChildren(e: AnActionEvent?): Array<AnAction> {
     if (origin !is ActionGroup) return EMPTY_ARRAY
-    if (e == null) return origin.getChildren(null)
-    return origin.getChildren(wrapEvent(e, order))
+    return ActionWrapperUtil.getChildren(e, this, origin)
   }
 
   override fun actionPerformed(e: AnActionEvent) {
-    origin.actionPerformed(wrapEvent(e, order))
+    ActionWrapperUtil.actionPerformed(e, this, origin)
   }
 
+  // for EquatableTooltipProvider
   override fun equals(other: Any?): Boolean {
     if (this === other) {
       return true

platform/execution-impl/src/com/intellij/execution/lineMarker/LineMarkerActionWrapper.java

@@ -6,7 +6,6 @@ import com.intellij.execution.Location;
 import com.intellij.execution.PsiLocation;
 import com.intellij.execution.actions.ExecutorGroupActionGroup;
 import com.intellij.execution.actions.RunContextAction;
-import com.intellij.ide.DataManager;
 import com.intellij.openapi.actionSystem.*;
 import com.intellij.openapi.util.Key;
 import com.intellij.openapi.util.Pair;
@@ -17,12 +16,13 @@ import com.intellij.util.containers.ContainerUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-import javax.swing.*;
-
 /**
  * @author Dmitry Avdeev
  */
-public class LineMarkerActionWrapper extends ActionGroup implements PriorityAction, ActionWithDelegate<AnAction> {
+public class LineMarkerActionWrapper extends ActionGroup
+  implements DataSnapshotProvider, PriorityAction, ActionWithDelegate<AnAction> {
+
+  @Deprecated(forRemoval = true)
   public static final Key<Pair<PsiElement, DataContext>> LOCATION_WRAPPER = Key.create("LOCATION_WRAPPER");
 
   protected final SmartPsiElementPointer<PsiElement> myElement;
@@ -44,7 +44,25 @@ public class LineMarkerActionWrapper extends ActionGroup implements PriorityActi
 
   @Override
   public @NotNull ActionUpdateThread getActionUpdateThread() {
-    return myOrigin.getActionUpdateThread();
+    return ActionWrapperUtil.getActionUpdateThread(this, myOrigin);
+  }
+
+  @Override
+  public @NotNull AnAction getDelegate() {
+    return myOrigin;
+  }
+
+  @Override
+  public @NotNull Priority getPriority() {
+    return Priority.TOP;
+  }
+
+  @Override
+  public void dataSnapshot(@NotNull DataSink sink) {
+    sink.lazy(Location.DATA_KEY, () -> {
+      PsiElement e = myElement.getElement();
+      return e != null && e.isValid() ? new PsiLocation<>(e) : null;
+    });
   }
 
   @Override
@@ -58,7 +76,7 @@ public class LineMarkerActionWrapper extends ActionGroup implements PriorityActi
           action, ((RunContextAction)action).getExecutor(), order)));
     }
     if (myOrigin instanceof ActionGroup o) {
-      return o.getChildren(e == null ? null : wrapEvent(e));
+      return ActionWrapperUtil.getChildren(e, this, o);
     }
     return EMPTY_ARRAY;
   }
@@ -70,46 +88,11 @@ public class LineMarkerActionWrapper extends ActionGroup implements PriorityActi
 
   @Override
   public void update(@NotNull AnActionEvent e) {
-    AnActionEvent wrapped = wrapEvent(e);
-    myOrigin.update(wrapped);
-    Icon icon = wrapped.getPresentation().getIcon();
-    if (icon != null) {
-      e.getPresentation().setIcon(icon);
-    }
-  }
-
-  private @NotNull AnActionEvent wrapEvent(@NotNull AnActionEvent e) {
-    return e.withDataContext(wrapContext(e.getDataContext()));
-  }
-
-  private @NotNull DataContext wrapContext(@NotNull DataContext dataContext) {
-    Pair<PsiElement, DataContext> pair = DataManager.getInstance().loadFromDataContext(dataContext, LOCATION_WRAPPER);
-    PsiElement element = myElement.getElement();
-    if (pair == null || pair.first != element) {
-      DataContext wrapper = CustomizedDataContext.withSnapshot(dataContext, sink -> {
-        sink.lazy(Location.DATA_KEY, () -> {
-          PsiElement e = myElement.getElement();
-          return e != null && e.isValid() ? new PsiLocation<>(e) : null;
-        });
-      });
-      pair = Pair.pair(element, wrapper);
-      DataManager.getInstance().saveInDataContext(dataContext, LOCATION_WRAPPER, pair);
-    }
-    return pair.second;
+    ActionWrapperUtil.update(e, this, myOrigin);
   }
 
   @Override
   public void actionPerformed(@NotNull AnActionEvent e) {
-    myOrigin.actionPerformed(wrapEvent(e));
-  }
-
-  @Override
-  public @NotNull Priority getPriority() {
-    return Priority.TOP;
-  }
-
-  @Override
-  public @NotNull AnAction getDelegate() {
-    return myOrigin;
+    ActionWrapperUtil.actionPerformed(e, this, myOrigin);
   }
 }

platform/lang-impl/src/com/intellij/tools/BaseExternalToolsGroup.java

@@ -34,8 +34,6 @@ public abstract class BaseExternalToolsGroup<T extends Tool> extends ActionGroup
     ActionGroup delegate = (ActionGroup)(e != null ? e.getActionManager() : ActionManager.getInstance())
       .getAction(getDelegateGroupId());
     if (delegate == null) return EMPTY_ARRAY;
-    UpdateSession session = e != null ? e.getUpdateSession() : UpdateSession.EMPTY;
-    return session == UpdateSession.EMPTY ? delegate.getChildren(e) :
-           session.children(delegate).toArray(EMPTY_ARRAY);
+    return ActionWrapperUtil.getChildren(e, this, delegate);
   }
 }

platform/platform-api/api-dump-unreviewed.txt

@@ -1628,6 +1628,11 @@ e:com.intellij.openapi.actionSystem.ActionToolbarPosition
 - sf:TOP:com.intellij.openapi.actionSystem.ActionToolbarPosition
 - s:valueOf(java.lang.String):com.intellij.openapi.actionSystem.ActionToolbarPosition
 - s:values():com.intellij.openapi.actionSystem.ActionToolbarPosition[]
+f:com.intellij.openapi.actionSystem.ActionWrapperUtil
+- s:actionPerformed(com.intellij.openapi.actionSystem.AnActionEvent,com.intellij.openapi.actionSystem.AnAction,com.intellij.openapi.actionSystem.AnAction):V
+- s:getActionUpdateThread(com.intellij.openapi.actionSystem.AnAction,com.intellij.openapi.actionSystem.AnAction):com.intellij.openapi.actionSystem.ActionUpdateThread
+- s:getChildren(com.intellij.openapi.actionSystem.AnActionEvent,com.intellij.openapi.actionSystem.ActionGroup,com.intellij.openapi.actionSystem.ActionGroup):com.intellij.openapi.actionSystem.AnAction[]
+- s:update(com.intellij.openapi.actionSystem.AnActionEvent,com.intellij.openapi.actionSystem.AnAction,com.intellij.openapi.actionSystem.AnAction):V
 com.intellij.openapi.actionSystem.AlwaysVisibleActionGroup
 com.intellij.openapi.actionSystem.AnActionHolder
 - a:getAction():com.intellij.openapi.actionSystem.AnAction

platform/platform-api/src/com/intellij/openapi/actionSystem/ActionGroupWrapper.java

@@ -27,21 +27,17 @@ public class ActionGroupWrapper extends ActionGroup implements ActionWithDelegat
 
   @Override
   public @NotNull ActionUpdateThread getActionUpdateThread() {
-    return ActionUpdateThread.BGT;
+    return ActionWrapperUtil.getActionUpdateThread(this, myDelegate);
   }
 
   @Override
   public AnAction @NotNull [] getChildren(@Nullable AnActionEvent e) {
-    UpdateSession session = e != null ? e.getUpdateSession() : UpdateSession.EMPTY;
-    return session == UpdateSession.EMPTY ? myDelegate.getChildren(e) :
-           session.children(myDelegate).toArray(AnAction.EMPTY_ARRAY);
+    return ActionWrapperUtil.getChildren(e, this, myDelegate);
   }
 
   @Override
   public void update(@NotNull AnActionEvent e) {
-    UpdateSession session = e.getUpdateSession();
-    if (session == UpdateSession.EMPTY) myDelegate.update(e);
-    else e.getPresentation().copyFrom(session.presentation(myDelegate), null, true);
+    ActionWrapperUtil.update(e, this, myDelegate);
   }
 
   @Override

platform/platform-api/src/com/intellij/openapi/actionSystem/ActionWrapperUtil.java

@@ -0,0 +1,53 @@
+// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.openapi.actionSystem;
+
+import com.intellij.openapi.actionSystem.ex.CustomComponentAction;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import javax.swing.*;
+
+public final class ActionWrapperUtil {
+  private ActionWrapperUtil() {
+  }
+
+  public static void update(@NotNull AnActionEvent e,
+                            @NotNull AnAction wrapper,
+                            @NotNull AnAction delegate) {
+    UpdateSession session = e.getUpdateSession();
+    JComponent customComponent =
+      session == UpdateSession.EMPTY || !(wrapper instanceof CustomComponentAction) ? null :
+      e.getPresentation().getClientProperty(CustomComponentAction.COMPONENT_KEY);
+    AnActionEvent event = customizeEvent(e, wrapper);
+    if (session == UpdateSession.EMPTY || e != event) delegate.update(event);
+    else e.getPresentation().copyFrom(session.presentation(delegate), customComponent, true);
+  }
+
+  public static AnAction @NotNull [] getChildren(@Nullable AnActionEvent e,
+                                                 @NotNull ActionGroup wrapper,
+                                                 @NotNull ActionGroup delegate) {
+    if (e == null) return delegate.getChildren(null);
+    UpdateSession session = e.getUpdateSession();
+    AnActionEvent event = customizeEvent(e, wrapper);
+    return session == UpdateSession.EMPTY || e != event ? delegate.getChildren(event) :
+           session.children(delegate).toArray(AnAction.EMPTY_ARRAY);
+  }
+
+  public static void actionPerformed(@NotNull AnActionEvent e,
+                                     @NotNull AnAction wrapper,
+                                     @NotNull AnAction delegate) {
+    delegate.actionPerformed(customizeEvent(e, wrapper));
+  }
+
+  public static @NotNull ActionUpdateThread getActionUpdateThread(@NotNull AnAction wrapper,
+                                                                  @NotNull AnAction delegate) {
+    if (wrapper instanceof DataSnapshotProvider) return delegate.getActionUpdateThread();
+    return ActionUpdateThread.BGT;
+  }
+
+  private static @NotNull AnActionEvent customizeEvent(@NotNull AnActionEvent e,
+                                                       @NotNull AnAction wrapper) {
+    return wrapper instanceof DataSnapshotProvider o ?
+           e.withDataContext(CustomizedDataContext.withSnapshot(e.getDataContext(), o)) : e;
+  }
+}

platform/platform-api/src/com/intellij/openapi/actionSystem/AnActionWrapper.java

@@ -1,11 +1,8 @@
 // Copyright 2000-2022 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.intellij.openapi.actionSystem;
 
-import com.intellij.openapi.actionSystem.ex.CustomComponentAction;
 import org.jetbrains.annotations.NotNull;
 
-import javax.swing.*;
-
 public class AnActionWrapper extends AnAction implements ActionWithDelegate<AnAction>, PerformWithDocumentsCommitted {
   private final AnAction myDelegate;
 
@@ -22,12 +19,7 @@ public class AnActionWrapper extends AnAction implements ActionWithDelegate<AnAc
 
   @Override
   public void update(@NotNull AnActionEvent e) {
-    UpdateSession session = e.getUpdateSession();
-    JComponent customComponent =
-      session == UpdateSession.EMPTY || !(this instanceof CustomComponentAction) ? null :
-      e.getPresentation().getClientProperty(CustomComponentAction.COMPONENT_KEY);
-    if (session == UpdateSession.EMPTY) myDelegate.update(e);
-    else e.getPresentation().copyFrom(session.presentation(myDelegate), customComponent, true);
+    ActionWrapperUtil.update(e, this, myDelegate);
   }
 
   @Override
@@ -37,7 +29,7 @@ public class AnActionWrapper extends AnAction implements ActionWithDelegate<AnAc
 
   @Override
   public void actionPerformed(@NotNull AnActionEvent e) {
-    myDelegate.actionPerformed(e);
+    ActionWrapperUtil.actionPerformed(e, this, myDelegate);
   }
 
   @Override
@@ -47,7 +39,7 @@ public class AnActionWrapper extends AnAction implements ActionWithDelegate<AnAc
 
   @Override
   public @NotNull ActionUpdateThread getActionUpdateThread() {
-    return ActionUpdateThread.BGT;
+    return ActionWrapperUtil.getActionUpdateThread(this, myDelegate);
   }
 
   @Override