From b9ef2f71b3ad694dfb66ae990e10482955dae7bd Mon Sep 17 00:00:00 2001
From: Masayuki Nakano <masayuki@d-toybox.com>
Date: Tue, 27 Aug 2024 00:39:49 +0000
Subject: [PATCH] Bug 1911010 - Make `IMEContentObserver` observe
 `ParentChainChanged` and let `IMEStateManager` know that r=smaug

`mozAutoDocUpdate` does not make it in a document change when container node
of `insertBefore` has already been removed from the tree.  Therefore, once
`IMEContentObserver::mRootElement` is removed from the DOM tree without a
focus move, `IMEContentObserver` is notified mutations not in a document change.

Similar situations are handled in `IMEStateManager::OnRemoveContent` with
emulating a focus change and that will destroy the active `IMEContentObserver`.
Therefore, if `IMEContentObserver::mRootElement` is removed, we should emulate
a focus move when `IMEStateManager` does not have focused element but there
is active `IMEContentObserver` (that means it is/was in the design mode).

However, checking whether the removed node contains the observing node of the
active `IMEContentObserver` may be expensive.  So, doing expensive things in
`IMEStateManager::OnRemoveContent` may make mutations slower.  Therefore, this
patch makes `IMEContentObserver` observe `ParentChainChanged` and it let
`IMEStateManager` know that with calling its
`OnParentChainChangedOfObservingElement`.  Finally, it calls
`IMEStateManager::OnRemoveContent` to emulate "blur" (and refocus if it's
required).

Differential Revision: https://phabricator.services.mozilla.com/D218696
---
 dom/base/FragmentOrElement.cpp                |   2 +-
 dom/base/nsINode.cpp                          |  77 +++++---
 dom/events/IMEContentObserver.cpp             |  41 +++++
 dom/events/IMEContentObserver.h               |   2 +
 dom/events/IMEStateManager.cpp                |  53 +++++-
 dom/events/IMEStateManager.h                  |   7 +
 dom/events/TextComposition.cpp                |   8 +-
 editor/libeditor/HTMLEditor.cpp               |  39 ++--
 ...body_in_design_mode_during_focus_move.html |  37 ++++
 widget/tests/browser/browser.toml             |   3 +
 ...y_removed_and_reconnected_in_designMode.js | 168 ++++++++++++++++++
 11 files changed, 386 insertions(+), 51 deletions(-)
 create mode 100644 testing/web-platform/tests/editing/crashtests/remove_body_in_design_mode_during_focus_move.html
 create mode 100644 widget/tests/browser/browser_test_ime_state_after_body_removed_and_reconnected_in_designMode.js

diff --git a/dom/base/FragmentOrElement.cpp b/dom/base/FragmentOrElement.cpp
index b79a9971988e..6fb2ca6bbdf9 100644
--- a/dom/base/FragmentOrElement.cpp
+++ b/dom/base/FragmentOrElement.cpp
@@ -185,7 +185,7 @@ HTMLSlotElement* nsIContent::GetAssignedSlotByMode() const {
 }
 
 nsIContent::IMEState nsIContent::GetDesiredIMEState() {
-  if (!IsEditable()) {
+  if (!IsEditable() || !IsInComposedDoc()) {
     // Check for the special case where we're dealing with elements which don't
     // have the editable flag set, but are readwrite (such as text controls).
     if (!IsElement() ||
diff --git a/dom/base/nsINode.cpp b/dom/base/nsINode.cpp
index 431288b3ee6d..f698b5921b4d 100644
--- a/dom/base/nsINode.cpp
+++ b/dom/base/nsINode.cpp
@@ -591,46 +591,73 @@ nsIContent* nsINode::GetSelectionRootContent(PresShell* aPresShell,
                                              bool aAllowCrossShadowBoundary) {
   NS_ENSURE_TRUE(aPresShell, nullptr);
 
-  if (IsDocument()) return AsDocument()->GetRootElement();
-  if (!IsContent()) return nullptr;
+  const bool isContent = IsContent();
 
-  if (GetComposedDoc() != aPresShell->GetDocument()) {
+  if (!isContent && !IsDocument()) {
     return nullptr;
   }
 
-  if (AsContent()->HasIndependentSelection() || IsInNativeAnonymousSubtree()) {
-    // This node should be an inclusive descendant of input/textarea editor.
-    // In that case, the anonymous <div> for TextEditor should be always the
-    // selection root.
-    // FIXME: If Selection for the document is collapsed in <input> or
-    // <textarea>, returning anonymous <div> may make the callers confused.
-    // Perhaps, we should do this only when this is in the native anonymous
-    // subtree unless the callers explicitly want to retrieve the anonymous
-    // <div> from a text control element.
-    if (Element* anonymousDivElement = GetAnonymousRootElementOfTextEditor()) {
-      return anonymousDivElement;
+  if (isContent) {
+    if (GetComposedDoc() != aPresShell->GetDocument()) {
+      return nullptr;
+    }
+
+    if (AsContent()->HasIndependentSelection() ||
+        IsInNativeAnonymousSubtree()) {
+      // This node should be an inclusive descendant of input/textarea editor.
+      // In that case, the anonymous <div> for TextEditor should be always the
+      // selection root.
+      // FIXME: If Selection for the document is collapsed in <input> or
+      // <textarea>, returning anonymous <div> may make the callers confused.
+      // Perhaps, we should do this only when this is in the native anonymous
+      // subtree unless the callers explicitly want to retrieve the anonymous
+      // <div> from a text control element.
+      if (Element* anonymousDivElement =
+              GetAnonymousRootElementOfTextEditor()) {
+        return anonymousDivElement;
+      }
     }
   }
 
-  nsPresContext* presContext = aPresShell->GetPresContext();
-  if (presContext) {
-    HTMLEditor* htmlEditor = nsContentUtils::GetHTMLEditor(presContext);
-    if (htmlEditor) {
-      // This node is in HTML editor.
+  if (nsPresContext* presContext = aPresShell->GetPresContext()) {
+    if (HTMLEditor* htmlEditor = nsContentUtils::GetHTMLEditor(presContext)) {
+      // When there is an HTMLEditor, selection root should be one of focused
+      // editing host, <body> or root of the (sub)tree which this node belong.
+
+      // If this node is in design mode or this node is not editable, selection
+      // root should be the <body> if this node is not in any subtrees and there
+      // is a <body> or the root of the shadow DOM if this node is in a shadow
+      // or the document element.
+      // XXX If this node is not connected, it seems that this should return
+      // nullptr because this node is not selectable.
       if (!IsInComposedDoc() || IsInDesignMode() ||
           !HasFlag(NODE_IS_EDITABLE)) {
-        nsIContent* editorRoot = htmlEditor->GetRoot();
-        NS_ENSURE_TRUE(editorRoot, nullptr);
-        return nsContentUtils::IsInSameAnonymousTree(this, editorRoot)
-                   ? editorRoot
+        Element* const bodyOrDocumentElement = [&]() -> Element* {
+          if (Element* const bodyElement = OwnerDoc()->GetBodyElement()) {
+            return bodyElement;
+          }
+          return OwnerDoc()->GetDocumentElement();
+        }();
+        NS_ENSURE_TRUE(bodyOrDocumentElement, nullptr);
+        return nsContentUtils::IsInSameAnonymousTree(this,
+                                                     bodyOrDocumentElement)
+                   ? bodyOrDocumentElement
                    : GetRootForContentSubtree(AsContent());
       }
-      // If the document isn't editable but this is editable, this is in
-      // contenteditable.  Use the editing host element for selection root.
+      // If this node is editable but not in the design mode, this is always an
+      // editable node in an editing host of contenteditable.  In this case,
+      // let's use the editing host element as selection root.
+      MOZ_ASSERT(IsEditable());
+      MOZ_ASSERT(!IsInDesignMode());
+      MOZ_ASSERT(IsContent());
       return static_cast<nsIContent*>(this)->GetEditingHost();
     }
   }
 
+  if (!isContent) {
+    return nullptr;
+  }
+
   RefPtr<nsFrameSelection> fs = aPresShell->FrameSelection();
   nsCOMPtr<nsIContent> content = fs->GetLimiter();
   if (!content) {
diff --git a/dom/events/IMEContentObserver.cpp b/dom/events/IMEContentObserver.cpp
index 2dcfca707592..d209bbfe2393 100644
--- a/dom/events/IMEContentObserver.cpp
+++ b/dom/events/IMEContentObserver.cpp
@@ -213,6 +213,10 @@ void IMEContentObserver::OnIMEReceivedFocus() {
 bool IMEContentObserver::InitWithEditor(nsPresContext& aPresContext,
                                         Element* aElement,
                                         EditorBase& aEditorBase) {
+  // mEditableNode is one of
+  // - Anonymous <div> in <input> or <textarea>
+  // - Editing host if it's not in the design mode
+  // - Document if it's in the design mode
   mEditableNode = IMEStateManager::GetRootEditableNode(aPresContext, aElement);
   if (NS_WARN_IF(!mEditableNode)) {
     return false;
@@ -260,11 +264,18 @@ bool IMEContentObserver::InitWithEditor(nsPresContext& aPresContext,
       return false;
     }
 
+    // If an editing host has focus, mRootElement is it.
+    // Otherwise, if we're in the design mode, mRootElement is the <body> if
+    // there is and startContainer is not outside of the <body>.  Otherwise, the
+    // document element is used instead.
     nsCOMPtr<nsINode> startContainer = selRange->GetStartContainer();
     mRootElement = Element::FromNodeOrNull(
         startContainer->GetSelectionRootContent(presShell));
   } else {
     MOZ_ASSERT(!mIsTextControl);
+    // If an editing host has focus, mRootElement is it.
+    // Otherwise, if we're in the design mode, mRootElement is the <body> if
+    // there is.  Otherwise, the document element is used instead.
     nsCOMPtr<nsINode> editableNode = mEditableNode;
     mRootElement = Element::FromNodeOrNull(
         editableNode->GetSelectionRootContent(presShell));
@@ -325,6 +336,10 @@ void IMEContentObserver::ObserveEditableNode() {
     mEditorBase->SetIMEContentObserver(this);
   }
 
+  MOZ_LOG(sIMECOLog, LogLevel::Info,
+          ("0x%p ObserveEditableNode(), starting to observe 0x%p (%s)", this,
+           mRootElement.get(), ToString(*mRootElement).c_str()));
+
   mRootElement->AddMutationObserver(this);
   // If it's in a document (should be so), we can use document observer to
   // reduce redundant computation of text change offsets.
@@ -380,6 +395,12 @@ void IMEContentObserver::UnregisterObservers() {
   if (!mIsObserving) {
     return;
   }
+
+  MOZ_LOG(sIMECOLog, LogLevel::Info,
+          ("0x%p UnregisterObservers(), stop observing 0x%p (%s)", this,
+           mRootElement.get(),
+           mRootElement ? ToString(*mRootElement).c_str() : "nullptr"));
+
   mIsObserving = false;
 
   if (mEditorBase) {
@@ -1236,6 +1257,26 @@ void IMEContentObserver::ContentRemoved(nsIContent* aChild,
   MaybeNotifyIMEOfTextChange(data);
 }
 
+MOZ_CAN_RUN_SCRIPT_BOUNDARY void IMEContentObserver::ParentChainChanged(
+    nsIContent* aContent) {
+  // When the observing element itself is directly removed from the document
+  // without a focus move, i.e., it's the root of the removed document fragment
+  // and the editor was handling the design mode, we have already stopped
+  // observing the element because IMEStateManager::OnRemoveContent() should
+  // have already been called for it and the instance which was observing the
+  // node has already been destroyed.  Therefore, this is called only when
+  // this is observing the <body> in the design mode and it's disconnected from
+  // the tree by an <html> element removal.  Even in this case, IMEStateManager
+  // never gets a focus change notification, but we need to notify IME of focus
+  // change because we cannot interact with IME anymore due to no editable
+  // content.  Therefore, this method notifies IMEStateManager of the
+  // disconnection of the observing node to emulate a blur from the editable
+  // content.
+  MOZ_ASSERT(mIsObserving);
+  OwningNonNull<IMEContentObserver> observer(*this);
+  IMEStateManager::OnParentChainChangedOfObservingElement(observer);
+}
+
 void IMEContentObserver::OnTextControlValueChangedWhileNotObservable(
     const nsAString& aNewValue) {
   MOZ_ASSERT(mEditorBase);
diff --git a/dom/events/IMEContentObserver.h b/dom/events/IMEContentObserver.h
index b31f52abc1ed..419fb4099982 100644
--- a/dom/events/IMEContentObserver.h
+++ b/dom/events/IMEContentObserver.h
@@ -15,6 +15,7 @@
 #include "nsCOMPtr.h"
 #include "nsCycleCollectionParticipant.h"
 #include "nsIDocShell.h"  // XXX Why does only this need to be included here?
+#include "nsIMutationObserver.h"
 #include "nsIReflowObserver.h"
 #include "nsIScrollObserver.h"
 #include "nsIWidget.h"
@@ -59,6 +60,7 @@ class IMEContentObserver final : public nsStubMutationObserver,
   NS_DECL_NSIMUTATIONOBSERVER_CONTENTAPPENDED
   NS_DECL_NSIMUTATIONOBSERVER_CONTENTINSERTED
   NS_DECL_NSIMUTATIONOBSERVER_CONTENTREMOVED
+  NS_DECL_NSIMUTATIONOBSERVER_PARENTCHAINCHANGED
   NS_DECL_NSIREFLOWOBSERVER
 
   // nsIScrollObserver
diff --git a/dom/events/IMEStateManager.cpp b/dom/events/IMEStateManager.cpp
index 90946871b1b7..53576110735a 100644
--- a/dom/events/IMEStateManager.cpp
+++ b/dom/events/IMEStateManager.cpp
@@ -415,8 +415,7 @@ nsresult IMEStateManager::OnRemoveContent(nsPresContext& aPresContext,
 
     if (compositionInContent) {
       MOZ_LOG(sISMLog, LogLevel::Debug,
-              ("  OnRemoveContent(), "
-               "composition is in the content"));
+              ("  OnRemoveContent(), composition is in the content"));
 
       // Try resetting the native IME state.  Be aware, typically, this method
       // is called during the content being removed.  Then, the native
@@ -430,8 +429,15 @@ nsresult IMEStateManager::OnRemoveContent(nsPresContext& aPresContext,
     }
   }
 
-  if (!sFocusedPresContext || !sFocusedElement ||
-      !sFocusedElement->IsInclusiveDescendantOf(&aElement)) {
+  if (!sFocusedPresContext ||
+      // If focused element is a text control or an editing host, we need to
+      // emulate "blur" on it when it's removed.
+      (sFocusedElement && sFocusedElement != &aElement) ||
+      // If it is (or was) in design mode, we need to emulate "blur" on the
+      // document when the observing element (typically, <body>) is removed.
+      (!sFocusedElement &&
+       (!sActiveIMEContentObserver ||
+        sActiveIMEContentObserver->GetObservingElement() != &aElement))) {
     return NS_OK;
   }
   MOZ_ASSERT(sFocusedPresContext == &aPresContext);
@@ -463,19 +469,48 @@ nsresult IMEStateManager::OnRemoveContent(nsPresContext& aPresContext,
   SetIMEState(newState, &aPresContext, nullptr, textInputHandlingWidget, action,
               origin);
   if (sFocusedPresContext != &aPresContext || sFocusedElement) {
-    return NS_OK;  // Some body must have set focus
+    return NS_OK;  // Somebody already has focus, don't steal it.
   }
 
   if (IsIMEObserverNeeded(newState)) {
-    if (RefPtr<HTMLEditor> htmlEditor =
-            nsContentUtils::GetHTMLEditor(&aPresContext)) {
-      CreateIMEContentObserver(*htmlEditor, nullptr);
-    }
+    // Initializing IMEContentObserver instance requires Selection, but its
+    // ranges have not been adjusted for this removal.  Therefore, we need to
+    // wait a moment.
+    nsContentUtils::AddScriptRunner(NS_NewRunnableFunction(
+        "IMEStateManager::RecreateIMEContentObserverWhenContentRemoved",
+        [presContext = OwningNonNull{aPresContext}]() {
+          MOZ_ASSERT(sFocusedPresContext == presContext);
+          MOZ_ASSERT(!sFocusedElement);
+          if (RefPtr<HTMLEditor> htmlEditor =
+                  nsContentUtils::GetHTMLEditor(presContext)) {
+            CreateIMEContentObserver(*htmlEditor, nullptr);
+          }
+        }));
   }
 
   return NS_OK;
 }
 
+void IMEStateManager::OnParentChainChangedOfObservingElement(
+    IMEContentObserver& aObserver) {
+  if (!sFocusedPresContext || sActiveIMEContentObserver != &aObserver) {
+    return;
+  }
+  RefPtr<nsPresContext> presContext = aObserver.GetPresContext();
+  RefPtr<Element> element = aObserver.GetObservingElement();
+  if (NS_WARN_IF(!presContext) || NS_WARN_IF(!element)) {
+    return;
+  }
+  MOZ_LOG(sISMLog, LogLevel::Info,
+          ("OnParentChainChangedOfObservingElement(aObserver=0x%p), "
+           "sFocusedPresContext=0x%p, sFocusedElement=0x%p, "
+           "aObserver->GetPresContext()=0x%p, "
+           "aObserver->GetObservingElement()=0x%p",
+           &aObserver, sFocusedPresContext.get(), sFocusedElement.get(),
+           presContext.get(), element.get()));
+  OnRemoveContent(*presContext, *element);
+}
+
 // static
 bool IMEStateManager::CanHandleWith(const nsPresContext* aPresContext) {
   return aPresContext && aPresContext->GetPresShell() &&
diff --git a/dom/events/IMEStateManager.h b/dom/events/IMEStateManager.h
index 028d5ebd4f66..6a4bcb9db935 100644
--- a/dom/events/IMEStateManager.h
+++ b/dom/events/IMEStateManager.h
@@ -164,6 +164,13 @@ class IMEStateManager {
       nsPresContext& aPresContext);
   MOZ_CAN_RUN_SCRIPT static nsresult OnRemoveContent(
       nsPresContext& aPresContext, dom::Element& aElement);
+  /**
+   * Called when the parent chain of the observing element of IMEContentObserver
+   * is changed.
+   */
+  MOZ_CAN_RUN_SCRIPT static void OnParentChainChangedOfObservingElement(
+      IMEContentObserver& aObserver);
+
   /**
    * OnChangeFocus() should be called when focused content is changed or
    * IME enabled state is changed.  If nobody has focus, set both aPresContext
diff --git a/dom/events/TextComposition.cpp b/dom/events/TextComposition.cpp
index 9f6a58d72f52..9c447ee16aef 100644
--- a/dom/events/TextComposition.cpp
+++ b/dom/events/TextComposition.cpp
@@ -1059,10 +1059,10 @@ TextComposition* TextCompositionArray::GetCompositionFor(
 TextComposition* TextCompositionArray::GetCompositionInContent(
     nsPresContext* aPresContext, nsIContent* aContent) {
   // There should be only one composition per content object.
-  for (index_type i = Length(); i > 0; --i) {
-    nsINode* node = ElementAt(i - 1)->GetEventTargetNode();
-    if (node && node->IsInclusiveDescendantOf(aContent)) {
-      return ElementAt(i - 1);
+  for (TextComposition* const composition : Reversed(*this)) {
+    nsINode* node = composition->GetEventTargetNode();
+    if (node && node->IsInclusiveFlatTreeDescendantOf(aContent)) {
+      return composition;
     }
   }
   return nullptr;
diff --git a/editor/libeditor/HTMLEditor.cpp b/editor/libeditor/HTMLEditor.cpp
index 65019d04c6b8..2b78da32ab43 100644
--- a/editor/libeditor/HTMLEditor.cpp
+++ b/editor/libeditor/HTMLEditor.cpp
@@ -825,8 +825,14 @@ nsresult HTMLEditor::FocusedElementOrDocumentBecomesNotEditable(
     aHTMLEditor->mHasFocus = false;
     aHTMLEditor->mIsInDesignMode = false;
 
-    const RefPtr<Element> focusedElement =
-        nsFocusManager::GetFocusedElementStatic();
+    RefPtr<Element> focusedElement = nsFocusManager::GetFocusedElementStatic();
+    if (focusedElement && !focusedElement->IsInComposedDoc()) {
+      // nsFocusManager may keep storing the focused element even after
+      // disconnected from the tree, but HTMLEditor cannot work with editable
+      // nodes not in a composed document.  Therefore, we should treat no
+      // focused element in the case.
+      focusedElement = nullptr;
+    }
     TextControlElement* const focusedTextControlElement =
         TextControlElement::FromNodeOrNull(focusedElement);
     if ((focusedElement && focusedElement->IsEditable() &&
@@ -854,14 +860,19 @@ nsresult HTMLEditor::FocusedElementOrDocumentBecomesNotEditable(
   // If the element becomes not editable without focus change, IMEStateManager
   // does not have a chance to disable IME.  Therefore, (even if we fail to
   // handle the emulated blur/focus above,) we should notify IMEStateManager of
-  // the editing state change.
-  RefPtr<Element> focusedElement = nsFocusManager::GetFocusedElementStatic();
-  RefPtr<nsPresContext> presContext =
-      focusedElement ? focusedElement->GetPresContext(
-                           Element::PresContextFor::eForComposedDoc)
-                     : aDocument.GetPresContext();
-  if (presContext) {
-    IMEStateManager::MaybeOnEditableStateDisabled(*presContext, focusedElement);
+  // the editing state change.  Note that if the window of the HTMLEditor has
+  // already lost focus, we don't need to do that and we should not touch the
+  // other windows.
+  if (aHTMLEditor->OurWindowHasFocus()) {
+    if (RefPtr<nsPresContext> presContext = aHTMLEditor->GetPresContext()) {
+      RefPtr<Element> focusedElement =
+          nsFocusManager::GetFocusedElementStatic();
+      MOZ_ASSERT_IF(focusedElement,
+                    focusedElement->GetPresContext(
+                        Element::PresContextFor::eForComposedDoc));
+      IMEStateManager::MaybeOnEditableStateDisabled(*presContext,
+                                                    focusedElement);
+    }
   }
 
   return rv;
@@ -878,10 +889,13 @@ nsresult HTMLEditor::OnBlur(const EventTarget* aEventTarget) {
                       ? "true"
                       : "false")
                : "N/A"));
+  const Element* eventTargetAsElement =
+      Element::FromEventTargetOrNull(aEventTarget);
 
   // If another element already has focus, we should not maintain the selection
   // because we may not have the rights doing it.
-  if (nsFocusManager::GetFocusedElementStatic()) {
+  const Element* focusedElement = nsFocusManager::GetFocusedElementStatic();
+  if (focusedElement && focusedElement != eventTargetAsElement) {
     // XXX If we had focus and new focused element is a text control, we may
     // need to notify focus of its TextEditor...
     mIsInDesignMode = false;
@@ -892,7 +906,8 @@ nsresult HTMLEditor::OnBlur(const EventTarget* aEventTarget) {
   // If we're in the designMode and blur occurs, the target must be the document
   // node.  If a blur event is fired and the target is an element, it must be
   // delayed blur event at initializing the `HTMLEditor`.
-  if (mIsInDesignMode && Element::FromEventTargetOrNull(aEventTarget)) {
+  if (mIsInDesignMode && eventTargetAsElement &&
+      eventTargetAsElement->IsInComposedDoc()) {
     return NS_OK;
   }
 
diff --git a/testing/web-platform/tests/editing/crashtests/remove_body_in_design_mode_during_focus_move.html b/testing/web-platform/tests/editing/crashtests/remove_body_in_design_mode_during_focus_move.html
new file mode 100644
index 000000000000..20559c6dc1ca
--- /dev/null
+++ b/testing/web-platform/tests/editing/crashtests/remove_body_in_design_mode_during_focus_move.html
@@ -0,0 +1,37 @@
+<!doctype html>
+<html>
+<head>
+<meta charset="utf-8">
+<script>
+"use strict";
+
+document.addEventListener("DOMContentLoaded", async () => {
+  const q = document.querySelector("q");
+  const data = document.querySelector("data");
+  const meter = document.querySelector("meter");
+  const button = document.querySelector("button");
+  button.addEventListener(
+    "focusout",
+    () => document.createElement("a").append(q)
+  );
+  const promiseButtonFocus = new Promise(resolve => {
+    button.addEventListener("button", resolve, {once: true});
+  });
+  button.focus();
+  promiseButtonFocus;
+  document.designMode = "on";
+  q.insertBefore(document.body, data);
+  meter.insertBefore(data, meter.childNodes[0]);
+}, {once: true});
+</script>
+</head>
+<body>
+<q>
+  <data>
+</q>
+<dialog>
+  <meter>
+</dialog>
+<button></button>
+</body>
+</html>
diff --git a/widget/tests/browser/browser.toml b/widget/tests/browser/browser.toml
index f22bc53b47b9..0224d3f6bd0e 100644
--- a/widget/tests/browser/browser.toml
+++ b/widget/tests/browser/browser.toml
@@ -21,6 +21,9 @@ skip-if = ["os != 'win'"]
 
 ["browser_test_fullscreen_size.js"]
 
+["browser_test_ime_state_after_body_removed_and_reconnected_in_designMode.js"]
+support-files = [ "../file_ime_state_test_helper.js" ]
+
 ["browser_test_ime_state_in_contenteditable_on_focus_move_in_remote_content.js"]
 support-files = [
   "file_ime_state_tests.html",
diff --git a/widget/tests/browser/browser_test_ime_state_after_body_removed_and_reconnected_in_designMode.js b/widget/tests/browser/browser_test_ime_state_after_body_removed_and_reconnected_in_designMode.js
new file mode 100644
index 000000000000..bd4a7fbcca23
--- /dev/null
+++ b/widget/tests/browser/browser_test_ime_state_after_body_removed_and_reconnected_in_designMode.js
@@ -0,0 +1,168 @@
+/* Any copyright is dedicated to the Public Domain.
+   http://creativecommons.org/publicdomain/zero/1.0/ */
+
+"use strict";
+
+/* import-globals-from ../file_ime_state_test_helper.js */
+
+Services.scriptloader.loadSubScript(
+  "chrome://mochitests/content/browser/widget/tests/browser/file_ime_state_test_helper.js",
+  this
+);
+
+add_task(async function test_replace_body_in_designMode() {
+  await BrowserTestUtils.withNewTab(
+    "data:text/html,<html><body><br></body></html>",
+    async function (browser) {
+      const tipWrapper = new TIPWrapper(window);
+      ok(
+        tipWrapper.isAvailable(),
+        "test_replace_body_in_designMode: TextInputProcessor should've been initialized"
+      );
+
+      function waitFor(aWaitingNotification) {
+        return new Promise(resolve => {
+          tipWrapper.onIMEFocusBlur = aNotification => {
+            if (aNotification != aWaitingNotification) {
+              return;
+            }
+            tipWrapper.onIMEFocusBlur = null;
+            resolve();
+          };
+        });
+      }
+      const waitForInitialFocus = waitFor("notify-focus");
+      await SpecialPowers.spawn(browser, [], () => {
+        content.document.designMode = "on";
+        content.document.body.focus();
+      });
+      info("test_replace_body_in_designMode: Waiting for initial IME focus...");
+      await waitForInitialFocus;
+      Assert.equal(
+        window.windowUtils.IMEStatus,
+        Ci.nsIDOMWindowUtils.IME_STATUS_ENABLED,
+        "test_replace_body_in_designMode: IME should be enabled when the document becomes editable"
+      );
+
+      tipWrapper.clearFocusBlurNotifications();
+      const waitForRefocusAfterBodyRemoval = waitFor("notify-focus");
+      await SpecialPowers.spawn(browser, [], () => {
+        content.wrappedJSObject.body = content.document.body;
+        content.document.body.remove();
+      });
+      info(
+        "test_replace_body_in_designMode: Waiting for IME refocus after the <body> is removed..."
+      );
+      await waitForRefocusAfterBodyRemoval;
+      Assert.equal(
+        window.windowUtils.IMEStatus,
+        Ci.nsIDOMWindowUtils.IME_STATUS_ENABLED,
+        "test_replace_body_in_designMode: IME should be enabled after the <body> is removed"
+      );
+
+      await SpecialPowers.spawn(browser, [], () => {
+        content.document.documentElement.appendChild(
+          content.wrappedJSObject.body
+        );
+      });
+      tipWrapper.clearFocusBlurNotifications();
+      await new Promise(resolve => {
+        window.requestAnimationFrame(() =>
+          window.requestAnimationFrame(resolve)
+        );
+      });
+      // FIXME: IME should be refocused when new <body> is inserted because
+      // HTMLEditor::GetRoot() returns the reconnected <body> now, but
+      // IMEContentObserver keeps observing the <html>.
+      Assert.equal(
+        tipWrapper.numberOfBlurNotifications +
+          tipWrapper.numberOfFocusNotifications,
+        0,
+        "test_replace_body_in_designMode: IME should not be refocused when the <body> is reconnected"
+      );
+      Assert.equal(
+        window.windowUtils.IMEStatus,
+        Ci.nsIDOMWindowUtils.IME_STATUS_ENABLED,
+        "test_replace_body_in_designMode: IME should be enabled after the <body> is reconnected"
+      );
+    }
+  );
+});
+
+add_task(async function test_replace_document_element_in_designMode() {
+  await BrowserTestUtils.withNewTab(
+    "data:text/html,<html><body><br></body></html>",
+    async function (browser) {
+      const tipWrapper = new TIPWrapper(window);
+      ok(
+        tipWrapper.isAvailable(),
+        "test_replace_document_element_in_designMode: TextInputProcessor should've been initialized"
+      );
+
+      function waitFor(aWaitingNotification) {
+        return new Promise(resolve => {
+          tipWrapper.onIMEFocusBlur = aNotification => {
+            if (aNotification != aWaitingNotification) {
+              return;
+            }
+            tipWrapper.onIMEFocusBlur = null;
+            resolve();
+          };
+        });
+      }
+      const waitForInitialFocus = waitFor("notify-focus");
+      await SpecialPowers.spawn(browser, [], () => {
+        content.document.designMode = "on";
+        content.document.body.focus();
+      });
+      info(
+        "test_replace_document_element_in_designMode: Waiting for initial IME focus..."
+      );
+      await waitForInitialFocus;
+      Assert.equal(
+        window.windowUtils.IMEStatus,
+        Ci.nsIDOMWindowUtils.IME_STATUS_ENABLED,
+        "test_replace_document_element_in_designMode: IME should be enabled when the document becomes editable"
+      );
+
+      tipWrapper.clearFocusBlurNotifications();
+      const waitForRefocusAfterDocumentElementRemoval = waitFor("notify-blur");
+      await SpecialPowers.spawn(browser, [], () => {
+        content.wrappedJSObject.documentElement =
+          content.document.documentElement;
+        content.document.documentElement.remove();
+      });
+      info(
+        "test_replace_document_element_in_designMode: Waiting for IME blur after the <html> is removed..."
+      );
+      await waitForRefocusAfterDocumentElementRemoval;
+      Assert.equal(
+        window.windowUtils.IMEStatus,
+        Ci.nsIDOMWindowUtils.IME_STATUS_ENABLED,
+        "test_replace_document_element_in_designMode: IME should be enabled after the <html> is removed"
+      );
+
+      await SpecialPowers.spawn(browser, [], () => {
+        content.document.appendChild(content.wrappedJSObject.documentElement);
+      });
+      tipWrapper.clearFocusBlurNotifications();
+      await new Promise(resolve => {
+        window.requestAnimationFrame(() =>
+          window.requestAnimationFrame(resolve)
+        );
+      });
+      // FIXME: IME should be focused when new root element is inserted.
+      Assert.equal(
+        tipWrapper.numberOfBlurNotifications +
+          tipWrapper.numberOfFocusNotifications,
+        0,
+        "test_replace_document_element_in_designMode: IME should not be refocused when the <html> is reconnected"
+      );
+      Assert.equal(
+        window.windowUtils.IMEStatus,
+        Ci.nsIDOMWindowUtils.IME_STATUS_ENABLED,
+        "test_replace_document_element_in_designMode: IME should be enabled after the <html> is reconnected"
+      );
+    }
+  );
+});