From d46319889290e48532950f132acab325a5e89a9a Mon Sep 17 00:00:00 2001
From: Segun Adebayo <joseshegs@gmail.com>
Date: Tue, 24 Sep 2024 14:43:17 +0100
Subject: [PATCH] refactor: improve support for shadow root

---
 packages/utilities/dom-query/src/env.ts       |  5 ++---
 packages/utilities/dom-query/src/scope.ts     |  4 ++--
 packages/utilities/focus-visible/src/index.ts | 22 +++++++++----------
 3 files changed, 14 insertions(+), 17 deletions(-)

diff --git a/packages/utilities/dom-query/src/env.ts b/packages/utilities/dom-query/src/env.ts
index 5d2066ce7e..15295c2fdf 100644
--- a/packages/utilities/dom-query/src/env.ts
+++ b/packages/utilities/dom-query/src/env.ts
@@ -17,9 +17,8 @@ export function getWindow(el: Node | ShadowRoot | Document | null | undefined) {
   return window
 }
 
-export function getActiveElement(el: HTMLElement): HTMLElement | null {
-  const doc = getDocument(el)
-  let activeElement = doc.activeElement as HTMLElement | null
+export function getActiveElement(rootNode: Document | ShadowRoot): HTMLElement | null {
+  let activeElement = rootNode.activeElement as HTMLElement | null
 
   while (activeElement?.shadowRoot) {
     const el = activeElement.shadowRoot.activeElement as HTMLElement | null
diff --git a/packages/utilities/dom-query/src/scope.ts b/packages/utilities/dom-query/src/scope.ts
index 80b148ec6e..852c671499 100644
--- a/packages/utilities/dom-query/src/scope.ts
+++ b/packages/utilities/dom-query/src/scope.ts
@@ -1,4 +1,4 @@
-import { getDocument } from "./env"
+import { getActiveElement, getDocument } from "./env"
 
 export interface ScopeContext {
   getRootNode?(): Document | ShadowRoot | Node
@@ -9,7 +9,7 @@ export function createScope<T>(methods: T) {
     getRootNode: (ctx: ScopeContext) => (ctx.getRootNode?.() ?? document) as Document | ShadowRoot,
     getDoc: (ctx: ScopeContext) => getDocument(dom.getRootNode(ctx)),
     getWin: (ctx: ScopeContext) => dom.getDoc(ctx).defaultView ?? window,
-    getActiveElement: (ctx: ScopeContext) => dom.getRootNode(ctx).activeElement,
+    getActiveElement: (ctx: ScopeContext) => getActiveElement(dom.getRootNode(ctx)),
     isActiveElement: (ctx: ScopeContext, elem: HTMLElement | null) => elem === dom.getActiveElement(ctx),
     getById: <T extends Element = HTMLElement>(ctx: ScopeContext, id: string) =>
       dom.getRootNode(ctx).getElementById(id) as T | null,
diff --git a/packages/utilities/focus-visible/src/index.ts b/packages/utilities/focus-visible/src/index.ts
index 864f972cce..00ac4f4799 100644
--- a/packages/utilities/focus-visible/src/index.ts
+++ b/packages/utilities/focus-visible/src/index.ts
@@ -2,7 +2,7 @@
  * Credit: Huge props to the team at Adobe for inspiring this implementation.
  * https://github.com/adobe/react-spectrum/blob/main/packages/%40react-aria/interactions/src/useFocusVisible.ts
  */
-import { getDocument, getWindow, isMac } from "@zag-js/dom-query"
+import { getDocument, getEventTarget, getWindow, isMac } from "@zag-js/dom-query"
 
 function isVirtualClick(event: MouseEvent | PointerEvent): boolean {
   if ((event as any).mozInputSource === 0 && event.isTrusted) return true
@@ -23,23 +23,19 @@ function isValidKey(e: KeyboardEvent) {
 const nonTextInputTypes = new Set(["checkbox", "radio", "range", "color", "file", "image", "button", "submit", "reset"])
 
 function isKeyboardFocusEvent(isTextInput: boolean, modality: Modality, e: HandlerEvent) {
-  const IHTMLInputElement =
-    typeof window !== "undefined" ? getWindow(e?.target as Element).HTMLInputElement : HTMLInputElement
-  const IHTMLTextAreaElement =
-    typeof window !== "undefined" ? getWindow(e?.target as Element).HTMLTextAreaElement : HTMLTextAreaElement
-  const IHTMLElement = typeof window !== "undefined" ? getWindow(e?.target as Element).HTMLElement : HTMLElement
-  const IKeyboardEvent = typeof window !== "undefined" ? getWindow(e?.target as Element).KeyboardEvent : KeyboardEvent
+  const target = e ? getEventTarget<Element>(e) : null
+  const win = getWindow(target)
 
   isTextInput =
     isTextInput ||
-    (e?.target instanceof IHTMLInputElement && !nonTextInputTypes.has(e?.target?.type)) ||
-    e?.target instanceof IHTMLTextAreaElement ||
-    (e?.target instanceof IHTMLElement && e?.target.isContentEditable)
+    (target instanceof win.HTMLInputElement && !nonTextInputTypes.has(target?.type)) ||
+    target instanceof win.HTMLTextAreaElement ||
+    (target instanceof win.HTMLElement && target.isContentEditable)
 
   return !(
     isTextInput &&
     modality === "keyboard" &&
-    e instanceof IKeyboardEvent &&
+    e instanceof win.KeyboardEvent &&
     !Reflect.has(FOCUS_VISIBLE_INPUT_KEYS, e.key)
   )
 }
@@ -108,7 +104,9 @@ function handleFocusEvent(e: FocusEvent) {
   // Firefox fires two extra focus events when the user first clicks into an iframe:
   // first on the window, then on the document. We ignore these events so they don't
   // cause keyboard focus rings to appear.
-  if (e.target === getWindow(e.target as Element) || e.target === getDocument(e.target as Element)) {
+  const target = getEventTarget(e)
+
+  if (target === getWindow(target as Element) || target === getDocument(target as Element)) {
     return
   }