make modality stateForComponent work correctly for not yet displayed dialogs (IDEA-181253)

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

@@ -54,6 +54,11 @@ public class ModalityStateEx extends ModalityState {
     return new ModalityStateEx(list.toArray());
   }
 
+  void forceModalEntities(List<Object> entities) {
+    myModalEntities.clear();
+    myModalEntities.addAll(entities);
+  }
+
   @Override
   public boolean dominates(@NotNull ModalityState anotherState){
     if (anotherState == ModalityState.any()) return false;

platform/platform-api/src/com/intellij/openapi/ui/DialogWrapper.java

@@ -23,7 +23,6 @@ import com.intellij.openapi.Disposable;
 import com.intellij.openapi.MnemonicHelper;
 import com.intellij.openapi.actionSystem.*;
 import com.intellij.openapi.actionSystem.ex.ActionUtil;
-import com.intellij.openapi.application.Application;
 import com.intellij.openapi.application.ApplicationInfo;
 import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.openapi.application.ModalityState;
@@ -1401,17 +1400,11 @@ public abstract class DialogWrapper {
 
     if (getValidationThreadToUse() == Alarm.ThreadToUse.SWING_THREAD) {
       // null if headless
-      Application app = ApplicationManager.getApplication();
-      if (app == null) {
-        myValidationAlarm.addRequest(validateRequest, myValidationDelay, null);
-      }
-      else {
-        app.invokeLater(() -> {
-          JRootPane rootPane = getRootPane();
-          myValidationAlarm.addRequest(validateRequest, myValidationDelay,
-                                       rootPane == null ? ModalityState.current() : ModalityState.stateForComponent(rootPane));
-        }, ModalityState.any(), __ -> myDisposed);
-      }
+      JRootPane rootPane = getRootPane();
+      myValidationAlarm.addRequest(validateRequest, myValidationDelay,
+                                   ApplicationManager.getApplication() == null
+                                   ? null
+                                   : rootPane == null ? ModalityState.current() : ModalityState.stateForComponent(rootPane));
     }
     else {
       myValidationAlarm.addRequest(validateRequest, myValidationDelay);

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

@@ -40,6 +40,7 @@ import javax.swing.*;
 import java.awt.*;
 import java.util.*;
 import java.util.List;
+import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 @SuppressWarnings("SSBasedInspection")
@@ -102,15 +103,26 @@ public class LaterInvocator {
   public static void removeModalityStateListener(@NotNull ModalityStateListener listener) {
     ourModalityStateMulticaster.removeListener(listener);
   }
+  
+  private static final ConcurrentMap<Window, ModalityStateEx> ourWindowModalities = ContainerUtil.createConcurrentWeakMap();
 
   @NotNull
   static ModalityStateEx modalityStateForWindow(@NotNull Window window) {
-    for (ModalityStateEx state : ourModalityStack) {
-      if (state.getModalEntities().contains(window)) {
-        return state;
+    return ourWindowModalities.computeIfAbsent(window, __ -> {
+      for (ModalityStateEx state : ourModalityStack) {
+        if (state.getModalEntities().contains(window)) {
+          return state;
+        }
       }
-    }
-    return (ModalityStateEx)ModalityState.NON_MODAL;
+
+      Window owner = window.getOwner();
+      ModalityStateEx ownerState = owner == null ? (ModalityStateEx)ModalityState.NON_MODAL : modalityStateForWindow(owner);
+      return isModalDialog(window) ? ownerState.appendEntity(window) : ownerState;
+    });
+  }
+
+  private static boolean isModalDialog(@NotNull Object window) {
+    return window instanceof Dialog && ((Dialog)window).isModal();
   }
 
   @NotNull
@@ -187,7 +199,13 @@ public class LaterInvocator {
   }
 
   public static void enterModal(@NotNull Object modalEntity) {
-    enterModal(modalEntity, getCurrentModalityState().appendEntity(modalEntity));
+    ModalityStateEx state = getCurrentModalityState().appendEntity(modalEntity);
+    if (isModalDialog(modalEntity)) {
+      List<Object> currentEntities = state.getModalEntities();
+      state = modalityStateForWindow((Window)modalEntity);
+      state.forceModalEntities(currentEntities);
+    }
+    enterModal(modalEntity, state);
   }
 
   public static void enterModal(@NotNull Object modalEntity, @NotNull ModalityStateEx appendedState) {
@@ -206,6 +224,9 @@ public class LaterInvocator {
     if (guard != null) {
       guard.enteredModality(appendedState);
     }
+
+    reincludeSkippedItems();
+    requestFlush();
   }
 
   public static void enterModal(Project project, Dialog dialog) {

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

@@ -660,4 +660,35 @@ public class LaterInvocatorTest extends PlatformTestCase {
     }).assertTiming();
 
   }
+
+  private final JDialog myModalDialog = new JDialog((Dialog)null, true);
+
+  public void testModalityStateForNonDisplayedDialogGetsActualizedWhenItIsDisplayed() {
+    ApplicationManager.getApplication().invokeAndWait(() -> {
+      ModalityState state = ModalityState.stateForComponent(myModalDialog);
+      AtomicBoolean invoked = new AtomicBoolean();
+      ApplicationManager.getApplication().invokeLater(() -> invoked.set(true), state);
+      
+      LaterInvocator.enterModal("some object");
+
+      UIUtil.dispatchAllInvocationEvents();
+      assertFalse(invoked.get());
+      
+      LaterInvocator.enterModal(myModalDialog);
+
+      UIUtil.dispatchAllInvocationEvents();
+      assertTrue(invoked.get());
+    });
+  }
+
+  public void testModalityStateWorksImmediately() {
+    ApplicationManager.getApplication().invokeAndWait(() -> {
+      ModalityState state = ModalityState.stateForComponent(myModalDialog);
+      AtomicBoolean invoked = new AtomicBoolean();
+      ApplicationManager.getApplication().invokeLater(() -> invoked.set(true), state);
+      
+      UIUtil.dispatchAllInvocationEvents();
+      assertTrue(invoked.get());
+    });
+  }
 }