1 files changed, 42 insertions, 11 deletions

diff --git a/components/script/dom/document.rs b/components/script/dom/document.rs
index 157829412ce..97ca60af115 100644
--- a/components/script/dom/document.rs
+++ b/components/script/dom/document.rs
@@ -204,6 +204,16 @@ pub enum IsHTMLDocument {
     NonHTMLDocument,
 }
 
+#[derive(JSTraceable, MallocSizeOf)]
+#[unrooted_must_root_lint::must_root]
+enum FocusTransaction {
+    /// No focus operation is in effect.
+    NotInTransaction,
+    /// A focus operation is in effect.
+    /// Contains the element that has most recently requested focus for itself.
+    InTransaction(Option<Dom<Element>>),
+}
+
 /// <https://dom.spec.whatwg.org/#document>
 #[dom_struct]
 pub struct Document {
@@ -243,8 +253,8 @@ pub struct Document {
     ready_state: Cell<DocumentReadyState>,
     /// Whether the DOMContentLoaded event has already been dispatched.
     domcontentloaded_dispatched: Cell<bool>,
-    /// The element that has most recently requested focus for itself.
-    possibly_focused: MutNullableDom<Element>,
+    /// The state of this document's focus transaction.
+    focus_transaction: DomRefCell<FocusTransaction>,
     /// The element that currently has the document focus context.
     focused: MutNullableDom<Element>,
     /// The script element that is currently executing.
@@ -1011,21 +1021,41 @@ impl Document {
 
     /// Initiate a new round of checking for elements requesting focus. The last element to call
     /// `request_focus` before `commit_focus_transaction` is called will receive focus.
-    pub fn begin_focus_transaction(&self) {
-        self.possibly_focused.set(None);
+    fn begin_focus_transaction(&self) {
+        *self.focus_transaction.borrow_mut() = FocusTransaction::InTransaction(Default::default());
     }
 
     /// Request that the given element receive focus once the current transaction is complete.
-    pub fn request_focus(&self, elem: &Element) {
-        if elem.is_focusable_area() {
-            self.possibly_focused.set(Some(elem))
+    /// If None is passed, then whatever element is currently focused will no longer be focused
+    /// once the transaction is complete.
+    pub(crate) fn request_focus(&self, elem: Option<&Element>, focus_type: FocusType) {
+        let implicit_transaction = matches!(
+            *self.focus_transaction.borrow(),
+            FocusTransaction::NotInTransaction
+        );
+        if implicit_transaction {
+            self.begin_focus_transaction();
+        }
+        if elem.map_or(true, |e| e.is_focusable_area()) {
+            *self.focus_transaction.borrow_mut() =
+                FocusTransaction::InTransaction(elem.map(Dom::from_ref));
+        }
+        if implicit_transaction {
+            self.commit_focus_transaction(focus_type);
         }
     }
 
     /// Reassign the focus context to the element that last requested focus during this
     /// transaction, or none if no elements requested it.
-    pub fn commit_focus_transaction(&self, focus_type: FocusType) {
-        if self.focused == self.possibly_focused.get().as_deref() {
+    fn commit_focus_transaction(&self, focus_type: FocusType) {
+        let possibly_focused = match *self.focus_transaction.borrow() {
+            FocusTransaction::NotInTransaction => unreachable!(),
+            FocusTransaction::InTransaction(ref elem) => {
+                elem.as_ref().map(|e| DomRoot::from_ref(&**e))
+            },
+        };
+        *self.focus_transaction.borrow_mut() = FocusTransaction::NotInTransaction;
+        if self.focused == possibly_focused.as_ref().map(|e| &**e) {
             return;
         }
         if let Some(ref elem) = self.focused.get() {
@@ -1040,7 +1070,7 @@ impl Document {
             }
         }
 
-        self.focused.set(self.possibly_focused.get().as_deref());
+        self.focused.set(possibly_focused.as_ref().map(|e| &**e));
 
         if let Some(ref elem) = self.focused.get() {
             elem.set_focus_state(true);
@@ -1140,6 +1170,7 @@ impl Document {
             }
 
             self.begin_focus_transaction();
+            self.request_focus(Some(&*el), FocusType::Element);
         }
 
         // https://w3c.github.io/uievents/#event-type-click
@@ -2980,7 +3011,7 @@ impl Document {
             stylesheet_list: MutNullableDom::new(None),
             ready_state: Cell::new(ready_state),
             domcontentloaded_dispatched: Cell::new(domcontentloaded_dispatched),
-            possibly_focused: Default::default(),
+            focus_transaction: DomRefCell::new(FocusTransaction::NotInTransaction),
             focused: Default::default(),
             current_script: Default::default(),
             pending_parsing_blocking_script: Default::default(),