Update selection restore handling for modals

src/useDialogFocus.ts

···
     document.addEventListener('keydown', onKey);
 
     return () => {
+      element.removeEventListener('mousedown', onClick);
+      document.body.removeEventListener('focusin', onFocus);
+      document.removeEventListener('keydown', onKey);
+
       const active = document.activeElement as HTMLElement;
       if (!active || contains(element, active)) {
         restoreSelection(selection);
       }
-
-      element.removeEventListener('mousedown', onClick);
-      document.body.removeEventListener('focusin', onFocus);
-      document.removeEventListener('keydown', onKey);
     };
   }, [ref.current, disabled]);
 }

src/useMenuFocus.ts

···
     document.addEventListener('keydown', onKey);
 
     return () => {
+      document.body.removeEventListener('focusin', onFocus);
+      document.removeEventListener('keydown', onKey);
+
       const active = document.activeElement as HTMLElement;
       if (!active || contains(element, active)) {
         restoreSelection(selection);
       }
-
-      document.body.removeEventListener('focusin', onFocus);
-      document.removeEventListener('keydown', onKey);
     };
   }, [ref.current, disabled]);
 }

src/useModalFocus.ts

···
   const hasPriority = usePriority(ref, disabled);
 
   useLayoutEffect(() => {
-    const { current: element } = ref;
-    if (!element || disabled) return;
+    if (disabled) return;
 
     let selection: RestoreSelection | null = null;
-    if (!document.activeElement || !contains(element, document.activeElement)) {
-      const newTarget = getAutofocusTarget(element);
-      selection = snapshotSelection(element);
-      newTarget.focus();
+    if (
+      !document.activeElement ||
+      !contains(ref.current, document.activeElement)
+    ) {
+      const newTarget = ref.current ? getAutofocusTarget(ref.current) : null;
+      selection = snapshotSelection();
+      if (newTarget) newTarget.focus();
     }
 
     function onBlur(event: FocusEvent) {
+      const { current: element } = ref;
       if (!hasPriority.current || !element || event.defaultPrevented) return;
 
       if (
···
     }
 
     function onKeyDown(event: KeyboardEvent) {
+      const { current: element } = ref;
       if (!hasPriority.current || !element || event.defaultPrevented) return;
 
       if (event.code === 'Tab') {
···
     document.addEventListener('keydown', onKeyDown);
 
     return () => {
-      restoreSelection(selection);
       document.body.removeEventListener('focusout', onBlur);
       document.removeEventListener('keydown', onKeyDown);
+      restoreSelection(selection);
     };
-  }, [ref.current, hasPriority, disabled]);
+  }, [ref, hasPriority, disabled]);
 }