use global stateful component everywhere

defineGlobalComponent.tsx

import {
  Dispatch,
  memo,
  SetStateAction,
  useEffect,
  ComponentType,
  useMemo,
  useState,
  useCallback,
} from "react";
import { createRoot } from "react-dom/client";
import create from "zustand";

function defineGlobalComponent<T = {}>(
  getComponent: (api: {
    // why typeof here not working?
    useState: <S>(
      initialState: S | (() => S)
    ) => [S, Dispatch<SetStateAction<S>>];
  }) => ComponentType<T>
): ComponentType<T> {
  const store = create<any>(() => ({}));
  return memo((props) => {
    let stateIndex = 0;
    let maxStateIndex: number;
    const useState = useCallback((s: any) => {
      if (maxStateIndex !== undefined) stateIndex = stateIndex % maxStateIndex;
      const id = stateIndex;
      if (!(id in store.getState())) {
        store.setState({ [id]: typeof s === "function" ? s() : s });
      }
      stateIndex++;
      return [store((s) => s[id]), (v: any) => store.setState({ [id]: v })];
    }, []);
    useEffect(() => {
      // record useState calls on the end
      maxStateIndex = stateIndex;
    }, []);
    const Component: any = useMemo(() => getComponent({ useState: useState as any }), []);
    return <Component {...props} />;
  }) as any;
}

const GlobalCounter = defineGlobalComponent(({ useState }) => () => {
  const [count, setCount] = useState(0);
  const [input, setInput] = useState("");
  return (
    <div>
      <button onClick={() => setCount(count + 1)}>{count}</button>
      <input value={input} onChange={(e) => setInput(e.target.value)}></input>
      <h3>Another Counter</h3>
      <GlobalAnother uniqId={`inner`} />
    </div>
  );
});

const GlobalAnother = defineGlobalComponent(({ useState }) => {
  return ({ uniqId }: { uniqId: string }) => {
    const [count, setCount] = useState(2);
    const doubled = useMemo(() => count * 2, [count]);
    return (
      <>
        <button onClick={() => setCount(count + 2)}>
          {uniqId}: {count}
        </button>
        <div>doubled is {doubled}</div>
      </>
    );
  };
});

function App() {
  const [outerInput, setOuterInput] = useState("");

  return (
    <>
      <h3>Outer Input</h3>
      <input
        value={outerInput}
        onChange={(e) => setOuterInput(e.target.value)}
      />
      <h3>Shared Counter and Input</h3>
      <GlobalCounter />
      <GlobalCounter />
      <GlobalCounter />
      <h3>Another Counter</h3>
      <GlobalAnother uniqId={`${outerInput}1`} />
      <GlobalAnother uniqId={`${outerInput}2`} />
    </>
  );
}

createRoot(document.getElementById("root")!).render(<App />);