remember transactions per weak reference on modality state (that has identity equality now)

platform/core-impl/src/com/intellij/openapi/application/TransactionGuardImpl.java

@@ -38,7 +38,7 @@ import java.util.concurrent.atomic.AtomicLong;
 public class TransactionGuardImpl extends TransactionGuard {
   private static final Logger LOG = Logger.getInstance("#com.intellij.openapi.application.TransactionGuardImpl");
   private final Queue<Transaction> myQueue = new LinkedBlockingQueue<Transaction>();
-  private final Map<ProgressIndicator, TransactionIdImpl> myProgresses = ContainerUtil.createConcurrentWeakMap();
+  private final Map<ModalityState, TransactionIdImpl> myModalities = ContainerUtil.createConcurrentWeakMap();
   private TransactionIdImpl myCurrentTransaction;
   private boolean myWritingAllowed;
 
@@ -282,18 +282,24 @@ public class TransactionGuardImpl extends TransactionGuard {
   public TransactionIdImpl getContextTransaction() {
     if (!ApplicationManager.getApplication().isDispatchThread()) {
       ProgressIndicator indicator = ProgressIndicatorProvider.getGlobalProgressIndicator();
-      return indicator != null ? myProgresses.get(indicator) : null;
+      return indicator != null ? myModalities.get(indicator.getModalityState()) : null;
     }
 
     return myWritingAllowed ? myCurrentTransaction : null;
   }
 
-  public void registerProgress(@NotNull ProgressIndicator indicator, @Nullable TransactionIdImpl contextTransaction) {
+  public void enteredModality(@NotNull ModalityState modality) {
+    TransactionIdImpl contextTransaction = getContextTransaction();
     if (contextTransaction != null) {
-      myProgresses.put(indicator, contextTransaction);
+      myModalities.put(modality, contextTransaction);
     }
   }
 
+  @Nullable
+  public TransactionIdImpl getModalityTransaction(@NotNull ModalityState modalityState) {
+    return myModalities.get(modalityState);
+  }
+
   private static class Transaction {
     @NotNull  final Runnable runnable;
     @Nullable final TransactionIdImpl mergeInto;

platform/platform-impl/src/com/intellij/openapi/application/impl/ApplicationImpl.java

@@ -748,8 +748,7 @@ public class ApplicationImpl extends PlatformComponentManagerImpl implements App
   @Override
   @NotNull
   public ModalityState getCurrentModalityState() {
-    Object[] entities = LaterInvocator.getCurrentModalEntities();
-    return entities.length > 0 ? new ModalityStateEx(entities) : getNoneModalityState();
+    return LaterInvocator.getCurrentModalityState();
   }
 
   @Override

platform/platform-impl/src/com/intellij/openapi/application/impl/LaterInvocator.java

@@ -16,11 +16,9 @@
 package com.intellij.openapi.application.impl;
 
 import com.intellij.ide.IdeEventQueue;
+import com.intellij.idea.IdeaApplication;
 import com.intellij.openapi.Disposable;
-import com.intellij.openapi.application.Application;
-import com.intellij.openapi.application.ApplicationManager;
-import com.intellij.openapi.application.ModalityState;
-import com.intellij.openapi.application.ModalityStateListener;
+import com.intellij.openapi.application.*;
 import com.intellij.openapi.diagnostic.FrequentEventDetector;
 import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.progress.ProcessCanceledException;
@@ -82,6 +80,7 @@ public class LaterInvocator {
   }
 
   private static final List<Object> ourModalEntities = ContainerUtil.createLockFreeCopyOnWriteList();
+  private static Stack<ModalityState> ourModalityStack = new Stack<>(ModalityState.NON_MODAL);
   private static final List<RunnableInfo> ourQueue = new ArrayList<RunnableInfo>(); //protected by LOCK
   private static volatile int ourQueueSkipCount; // optimization
   private static final FlushQueue ourFlushQueueRunnable = new FlushQueue();
@@ -190,6 +189,12 @@ public class LaterInvocator {
     ourModalityStateMulticaster.getMulticaster().beforeModalityStateChanged(true);
 
     ourModalEntities.add(modalEntity);
+    ourModalityStack.push(new ModalityStateEx(ArrayUtil.toObjectArray(ourModalEntities)));
+
+    TransactionGuardImpl guard = IdeaApplication.isLoaded() ? (TransactionGuardImpl)TransactionGuard.getInstance() : null;
+    if (guard != null) {
+      guard.enteredModality(ourModalityStack.peek());
+    }
   }
 
   public static void leaveModal(@NotNull Object modalEntity) {
@@ -205,7 +210,9 @@ public class LaterInvocator {
     ourModalityStateMulticaster.getMulticaster().beforeModalityStateChanged(false);
 
     boolean removed = ourModalEntities.remove(modalEntity);
+    ourModalityStack.pop();
     LOG.assertTrue(removed, modalEntity);
+    LOG.assertTrue(!ourModalityStack.isEmpty(), modalEntity);
     cleanupQueueForModal(modalEntity);
     ourQueueSkipCount = 0;
     requestFlush();
@@ -242,6 +249,11 @@ public class LaterInvocator {
     return ArrayUtil.toObjectArray(ourModalEntities);
   }
 
+  @NotNull
+  public static ModalityState getCurrentModalityState() {
+    return ourModalityStack.peek();
+  }
+
   public static boolean isInModalContext() {
     LOG.assertTrue(isDispatchThread());
     return !ourModalEntities.isEmpty();

platform/platform-impl/src/com/intellij/openapi/progress/util/ProgressWindow.java

@@ -18,8 +18,6 @@ package com.intellij.openapi.progress.util;
 import com.intellij.ide.IdeEventQueue;
 import com.intellij.openapi.Disposable;
 import com.intellij.openapi.application.ApplicationManager;
-import com.intellij.openapi.application.TransactionGuard;
-import com.intellij.openapi.application.TransactionGuardImpl;
 import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.ui.DialogWrapper;
@@ -95,9 +93,6 @@ public class ProgressWindow extends ProgressIndicatorBase implements BlockingPro
     myShouldShowCancel = shouldShowCancel;
     myCancelText = cancelText;
     setModalityProgress(shouldShowBackground ? null : this);
-    TransactionGuardImpl guard = (TransactionGuardImpl)TransactionGuard.getInstance();
-    guard.registerProgress(this, guard.getContextTransaction());
-
     myFocusTrackback = new FocusTrackback(this, WindowManager.getInstance().suggestParentWindow(project), false);
 
     Component parent = parentComponent;

platform/platform-tests/testSrc/com/intellij/openapi/application/impl/LaterInvocatorTest.java

@@ -534,4 +534,21 @@ public class LaterInvocatorTest extends PlatformTestCase {
       checkOrder(2);
     });
   }
+
+  public void testModalityStateStaysTheSameBetweenInvocations() {
+    UIUtil.invokeAndWaitIfNeeded((Runnable)() -> {
+      Object modal1 = new Object();
+      Object modal2 = new Object();
+
+      LaterInvocator.enterModal(modal1);
+      ModalityState modalityState1 = ModalityState.current();
+      assertSame(modalityState1, ModalityState.current());
+
+      LaterInvocator.enterModal(modal2);
+      assertNotSame(modalityState1, ModalityState.current());
+      LaterInvocator.leaveModal(modal2);
+
+      assertSame(modalityState1, ModalityState.current());
+    });
+  }
 }