Save focus from shadow DOM

src/Modal.tsx

@@ -281,9 +281,15 @@ const Modal: React.ForwardRefExoticComponent<
     useImperativeHandle(ref, () => modal, [modal]);
 
     if (canUseDOM && !prevShow && show) {
-      lastFocusRef.current = activeElement(
-        ownerWindow?.document,
-      ) as HTMLElement | null;
+      let last = activeElement(ownerWindow?.document) as HTMLElement | null;
+      while (
+        last &&
+        last.shadowRoot instanceof ShadowRoot &&
+        last.shadowRoot.activeElement instanceof HTMLElement
+      ) {
+        last = last.shadowRoot.activeElement;
+      }
+      lastFocusRef.current = last;
     }
 
     // TODO: I think this needs to be in an effect

test/ModalSpec.tsx

@@ -329,5 +329,48 @@ describe('<Modal>', () => {
         expect(document.activeElement!.classList.contains('modal')).toBe(true);
       });
     });
+
+    it('should return focus to previously focused element in shadow DOM when modal closes', async () => {
+      class CustomElementImplementation extends HTMLElement {
+        constructor() {
+          super();
+          const shadowRoot = this.attachShadow({ mode: 'open' });
+          const button = document.createElement('button');
+          shadowRoot.appendChild(button);
+        }
+      }
+      const CustomElementTag = `focus-test-${Math.floor(
+        Math.random() * 9999999,
+      )}`;
+      window.customElements.define(
+        CustomElementTag,
+        CustomElementImplementation,
+      );
+      function ButtonAndModal(props: { show?: boolean }) {
+        return (
+          <>
+            <CustomElementTag />
+            <Modal {...props}>
+              <input autoFocus />
+            </Modal>
+          </>
+        );
+      }
+
+      const { rerender, container } = render(<ButtonAndModal />);
+      const element = container.querySelector(CustomElementTag);
+      const button = element!.shadowRoot!.querySelector('button');
+      button!.focus();
+
+      expect(document.activeElement).toBe(element);
+      expect(element?.shadowRoot?.activeElement).toBe(button);
+
+      rerender(<ButtonAndModal show />);
+      expect(document.activeElement).not.toBe(button);
+
+      rerender(<ButtonAndModal />);
+      expect(document.activeElement).toBe(element);
+      expect(element?.shadowRoot?.activeElement).toBe(button);
+    });
   });
 });