From ebd13e18455ddb18918f8e000613930d6c257bac Mon Sep 17 00:00:00 2001
From: Anton Korzunov <thekashey@gmail.com>
Date: Mon, 12 Feb 2024 21:50:06 +1100
Subject: [PATCH] fix: refactor use-focus-state to cause one redraw

---
 src/use-focus-state.js | 82 ++++++++++++++++++++++++++++++++++++------
 stories/control.js     | 17 +++++++--
 2 files changed, 85 insertions(+), 14 deletions(-)

diff --git a/src/use-focus-state.js b/src/use-focus-state.js
index 0e335cc..795fd41 100644
--- a/src/use-focus-state.js
+++ b/src/use-focus-state.js
@@ -5,28 +5,88 @@ import { createNanoEvents } from './nano-events';
 
 const mainbus = createNanoEvents();
 
+let subscribeCounter = 0;
+
+const onFocusIn = event => mainbus.emit('assign', event.target);
+const onFocusOut = event => mainbus.emit('reset', event.target);
+
+/**
+ * attaches focusin/focusout listener-singlenton to the document
+ * it will emit "reset" event on blur/focusout and cause "set" on focus/focusin
+ */
+const useDocumentFocusSubscribe = () => {
+  useEffect(() => {
+    if (!subscribeCounter) {
+      document.addEventListener('focusin', onFocusIn);
+      document.addEventListener('focusout', onFocusOut);
+    }
+    subscribeCounter += 1;
+    return () => {
+      subscribeCounter -= 1;
+      if (!subscribeCounter) {
+        document.removeEventListener('focusin', onFocusIn);
+        document.removeEventListener('focusout', onFocusOut);
+      }
+    };
+  }, []);
+};
+
+const getFocusState = (target, current) => {
+  if (target === current) {
+    return 'self';
+  }
+  if (current.contains(target)) {
+    return 'within';
+  }
+  return 'within-boundary';
+};
+
 export const useFocusState = () => {
-  const [marker] = useState({});
   const [active, setActive] = useState(false);
+  const [state, setState] = useState('');
   const ref = useRef(null);
+  const focusState = useRef({});
 
-  const onFocus = useCallback(() => {
-    mainbus.emit('focus', marker);
+  // initial focus
+  useEffect(() => {
+    if (ref.current) {
+      setActive(
+        ref.current === document.activeElement
+          || ref.current.contains(document.activeElement),
+      );
+      setState(getFocusState(document.activeElement, ref.current));
+    }
   }, []);
 
-  useEffect(() => {
-    setActive(
-      ref.current === document.activeElement
-        || ref.current.contains(document.activeElement),
-    );
+  const onFocus = useCallback((e) => {
+    // element caught focus. Store, but do not set value yes
+    focusState.current = {
+      focused: true,
+      state: getFocusState(e.target, e.currentTarget),
+    };
   }, []);
 
-  useEffect(() => mainbus.on('focus', (event) => {
-    setActive(event === marker);
-  }), []);
+
+  useDocumentFocusSubscribe();
+  useEffect(() => {
+    const fout = mainbus.on('reset', () => {
+      // focus is going somewhere
+      focusState.current = {};
+    });
+    const fin = mainbus.on('assign', () => {
+      // focus event propagation is ended
+      setActive(focusState.current.focused || false);
+      setState(focusState.current.state || '');
+    });
+    return () => {
+      fout();
+      fin();
+    };
+  }, []);
 
   return {
     active,
+    state,
     onFocus,
     ref,
   };
diff --git a/stories/control.js b/stories/control.js
index f954b07..0045045 100644
--- a/stories/control.js
+++ b/stories/control.js
@@ -31,7 +31,12 @@ const ControlTrap = () => {
 
 const FocusButton = ({ children }) => {
   const { active, onFocus, ref } = useFocusState();
-  return <button tabIndex={active ? undefined : -1} onFocus={onFocus} ref={ref}>{children}</button>;
+  return (
+    <button tabIndex={active ? undefined : -1} onFocus={onFocus} ref={ref}>
+      {active ? '*' : ''}
+      {children}
+    </button>
+  );
 };
 const RowingFocusTrap = () => {
   const { autofocus, focusNext, focusPrev } = useFocusScope();
@@ -56,7 +61,7 @@ const RowingFocusTrap = () => {
       onKeyDown={onKey}
       onFocus={onFocus}
       ref={ref}
-      style={{ border: active ? '3px solid green' : '3px solid grey' }}
+      style={{ border: active ? '3px solid red' : '3px solid #EEE' }}
     >
       <button>zero</button>
       <FocusButton>Button1</FocusButton>
@@ -75,7 +80,13 @@ const ControlledFocusButton = ({ children, onFocus: reportFocused, isActive }) =
     }
   }, [active]);
 
-  return <button tabIndex={isActive ? undefined : -1} onFocus={onFocus} ref={ref}>{children}</button>;
+  return (
+    <button tabIndex={isActive ? undefined : -1} onFocus={onFocus} ref={ref}>
+      {active ? '*' : ''}
+      {children}
+      {isActive ? '*' : ''}
+    </button>
+  );
 };
 const ConstantRowingFocusTrap = () => {
   const { focusNext, focusPrev } = useFocusScope();