chore: Internal reactive store utils

Author: pan-kot

src/internal/index.ts

@@ -39,3 +39,4 @@ export { default as circleIndex } from './utils/circle-index';
 export { default as Portal, PortalProps } from './portal';
 export { useMergeRefs } from './use-merge-refs';
 export { useRandomId, useUniqueId } from './use-unique-id';
+export { ReactiveStore, ReadonlyReactiveStore, useReaction, useSelector } from './reactive-store';

src/internal/reactive-store/__tests__/reactive-store.test.tsx

@@ -0,0 +1,175 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import React, { useRef, useState } from 'react';
+import { render, screen } from '@testing-library/react';
+
+import { ReactiveStore, useReaction, useSelector } from '../index';
+import { act } from 'react-dom/test-utils';
+
+interface State {
+  name: string;
+  values: Record<string, number>;
+}
+
+describe('ReactiveStore', () => {
+  let store = new ReactiveStore<State>({ name: '', values: {} });
+
+  beforeEach(() => {
+    store = new ReactiveStore<State>({ name: 'Test', values: { A: 1, B: 2 } });
+  });
+
+  function Provider({
+    store: customStore,
+    update,
+  }: {
+    store?: ReactiveStore<State>;
+    update?: (state: State) => State;
+  }) {
+    const providerStore = customStore ?? store;
+
+    const renderCounter = useRef(0);
+    renderCounter.current += 1;
+
+    useReaction(
+      providerStore,
+      s => s.name,
+      (newName, prevName) => {
+        const div = document.querySelector('[data-testid="reaction-name"]')!;
+        div.textContent = `${prevName} -> ${newName}`;
+      }
+    );
+
+    return (
+      <div>
+        <div data-testid="provider">Provider ({renderCounter.current})</div>
+        <SubscriberName store={providerStore} />
+        <SubscriberItemsList store={providerStore} />
+        <StoreUpdater store={providerStore} update={update} />
+        <div data-testid="reaction-name"></div>
+      </div>
+    );
+  }
+
+  function StoreUpdater({ store, update }: { store: ReactiveStore<State>; update?: (state: State) => State }) {
+    useState(() => {
+      if (update) {
+        store.set(prev => update(prev));
+      }
+      return null;
+    });
+    return null;
+  }
+
+  function SubscriberName({ store }: { store: ReactiveStore<State> }) {
+    const value = useSelector(store, s => s.name);
+    const renderCounter = useRef(0);
+    renderCounter.current += 1;
+    return (
+      <div data-testid="subscriber-name">
+        Subscriber name ({renderCounter.current}) {value}
+      </div>
+    );
+  }
+
+  function SubscriberItemsList({ store }: { store: ReactiveStore<State> }) {
+    const items = useSelector(store, s => s.values);
+    const itemIds = Object.keys(items);
+    const renderCounter = useRef(0);
+    renderCounter.current += 1;
+    return (
+      <div>
+        <div data-testid="subscriber-items">
+          Subscriber items ({renderCounter.current}) {itemIds.join(', ')}
+        </div>
+        {itemIds.map(itemId => (
+          <div key={itemId}>
+            <SubscriberItem id={itemId} store={store} />
+          </div>
+        ))}
+      </div>
+    );
+  }
+
+  function SubscriberItem({ id, store }: { id: string; store: ReactiveStore<State> }) {
+    const value = useSelector(store, s => s.values[id]);
+    const renderCounter = useRef(0);
+    renderCounter.current += 1;
+    return (
+      <div data-testid={`subscriber-${id}`}>
+        Subscriber {id} ({renderCounter.current}) {value}
+      </div>
+    );
+  }
+
+  test('initializes state correctly', () => {
+    render(<Provider />);
+
+    expect(screen.getByTestId('provider').textContent).toBe('Provider (1)');
+    expect(screen.getByTestId('subscriber-name').textContent).toBe('Subscriber name (1) Test');
+    expect(screen.getByTestId('subscriber-items').textContent).toBe('Subscriber items (1) A, B');
+    expect(screen.getByTestId('subscriber-A').textContent).toBe('Subscriber A (1) 1');
+    expect(screen.getByTestId('subscriber-B').textContent).toBe('Subscriber B (1) 2');
+  });
+
+  test('handles updates correctly', () => {
+    render(<Provider />);
+
+    act(() => store.set(prev => ({ ...prev, name: 'Test', values: { ...prev.values, B: 3, C: 4 } })));
+
+    expect(screen.getByTestId('provider').textContent).toBe('Provider (1)');
+    expect(screen.getByTestId('subscriber-name').textContent).toBe('Subscriber name (1) Test');
+    expect(screen.getByTestId('subscriber-items').textContent).toBe('Subscriber items (2) A, B, C');
+    expect(screen.getByTestId('subscriber-A').textContent).toBe('Subscriber A (2) 1');
+    expect(screen.getByTestId('subscriber-B').textContent).toBe('Subscriber B (2) 3');
+    expect(screen.getByTestId('subscriber-C').textContent).toBe('Subscriber C (1) 4');
+
+    act(() => store.set(prev => ({ ...prev, name: 'Updated' })));
+
+    expect(screen.getByTestId('provider').textContent).toBe('Provider (1)');
+    expect(screen.getByTestId('subscriber-name').textContent).toBe('Subscriber name (2) Updated');
+    expect(screen.getByTestId('subscriber-items').textContent).toBe('Subscriber items (2) A, B, C');
+    expect(screen.getByTestId('subscriber-A').textContent).toBe('Subscriber A (2) 1');
+    expect(screen.getByTestId('subscriber-B').textContent).toBe('Subscriber B (2) 3');
+    expect(screen.getByTestId('subscriber-C').textContent).toBe('Subscriber C (1) 4');
+  });
+
+  test('reacts to updates with useReaction', () => {
+    render(<Provider />);
+
+    act(() => store.set(prev => ({ ...prev, name: 'Reaction test' })));
+
+    expect(screen.getByTestId('subscriber-name').textContent).toBe('Subscriber name (2) Reaction test');
+    expect(screen.getByTestId('reaction-name').textContent).toBe('Test -> Reaction test');
+  });
+
+  test('unsubscribes listeners on unmount', () => {
+    const { unmount } = render(<Provider />);
+
+    expect(store).toEqual(expect.objectContaining({ _listeners: expect.objectContaining({ length: 5 }) }));
+
+    unmount();
+
+    expect(store).toEqual(expect.objectContaining({ _listeners: expect.objectContaining({ length: 0 }) }));
+  });
+
+  test('synchronizes updates done between render and effect', () => {
+    render(<Provider update={state => ({ ...state, name: 'Test!' })} />);
+
+    expect(screen.getByTestId('subscriber-name').textContent).toBe('Subscriber name (2) Test!');
+  });
+
+  test('reacts to store replacement', () => {
+    const { rerender } = render(<Provider />);
+
+    expect(screen.getByTestId('provider').textContent).toBe('Provider (1)');
+    expect(screen.getByTestId('subscriber-name').textContent).toBe('Subscriber name (1) Test');
+    expect(screen.getByTestId('reaction-name').textContent).toBe('');
+
+    rerender(<Provider store={new ReactiveStore<State>({ name: 'Other test', values: {} })} />);
+
+    expect(screen.getByTestId('provider').textContent).toBe('Provider (2)');
+    expect(screen.getByTestId('subscriber-name').textContent).toBe('Subscriber name (3) Other test');
+    expect(screen.getByTestId('reaction-name').textContent).toBe('Test -> Other test');
+  });
+});

src/internal/reactive-store/index.ts

@@ -0,0 +1,121 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { useEffect, useRef, useState } from 'react';
+
+type Selector<S, R> = (state: S) => R;
+type Listener<S> = (state: S, prevState: S) => void;
+
+export interface ReadonlyReactiveStore<S> {
+  get(): S;
+  subscribe<R>(selector: Selector<S, R>, listener: Listener<S>): () => void;
+  unsubscribe(listener: Listener<S>): void;
+}
+
+/**
+ * A pub/sub state management util that registers listeners by selectors.
+ * It comes with React utils that subscribe to state changes and trigger effects or React state updates.
+ *
+ * For simple states, a store can be defined as `ReactiveStore<StateType>`. For more complex states,
+ * it is recommended to create a custom class extending ReactiveStore and providing custom setters,
+ * for example:
+ *   class TableStore extends ReactiveStore<TableState> {
+ *     setVisibleColumns(visibleColumns) {
+ *       this.set((prev) => ({ ...prev, visibleColumns }));
+ *     }
+ *     // ...
+ *   }
+ *
+ * The store instance is usually created once when the component is mounted, which can be achieved with React's
+ * useRef or useMemo utils. To make the store aware of component's properties it is enough to assign them on
+ * every render, unless a state recomputation is required (in which case a useEffect is needed).
+ *   const store = useRef(new TableStore()).current;
+ *   store.totalColumns = props.totalColumns;
+ *
+ * As long as every selector un-subscribes on un-mount (which is the case when `useSelector()` helper is used),
+ * there is no need to do any cleanup on the store itself.
+ */
+export class ReactiveStore<S> implements ReadonlyReactiveStore<S> {
+  private _state: S;
+  private _listeners: [Selector<S, unknown>, Listener<S>][] = [];
+
+  constructor(state: S) {
+    this._state = state;
+  }
+
+  public get(): S {
+    return this._state;
+  }
+
+  public set(cb: (state: S) => S): void {
+    const prevState = this._state;
+    const newState = cb(prevState);
+
+    this._state = newState;
+
+    for (const [selector, listener] of this._listeners) {
+      if (selector(prevState) !== selector(newState)) {
+        listener(newState, prevState);
+      }
+    }
+  }
+
+  public subscribe<R>(selector: Selector<S, R>, listener: Listener<S>): () => void {
+    this._listeners.push([selector, listener]);
+    return () => this.unsubscribe(listener);
+  }
+
+  public unsubscribe(listener: Listener<S>): void {
+    this._listeners = this._listeners.filter(([, storedListener]) => storedListener !== listener);
+  }
+}
+
+/**
+ * Triggers an effect every time the selected store state changes.
+ */
+export function useReaction<S, R>(
+  store: ReadonlyReactiveStore<S>,
+  selector: Selector<S, R>,
+  effect: Listener<R>
+): void {
+  const prevStore = useRef(store);
+  useEffect(
+    () => {
+      if (prevStore.current !== store) {
+        effect(selector(store.get()), selector(prevStore.current.get()));
+        prevStore.current = store;
+      }
+      const unsubscribe = store.subscribe(selector, (next, prev) => effect(selector(next), selector(prev)));
+      return unsubscribe;
+    },
+    // Ignoring selector and effect as they are expected to stay constant.
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+    [store]
+  );
+}
+
+/**
+ * Creates React state that updates every time the selected store state changes.
+ */
+export function useSelector<S, R>(store: ReadonlyReactiveStore<S>, selector: Selector<S, R>): R {
+  const [state, setState] = useState<R>(selector(store.get()));
+
+  // We create subscription synchronously during the first render cycle to ensure the store updates that
+  // happen after the first render but before the first effect are not lost.
+  const unsubscribeRef = useRef(store.subscribe(selector, newState => setState(selector(newState))));
+  // When the component is un-mounted or the store reference changes, the old subscription is cancelled
+  // (and the new subscription is created for the new store instance).
+  const prevStore = useRef(store);
+  useEffect(() => {
+    if (prevStore.current !== store) {
+      setState(selector(store.get()));
+      unsubscribeRef.current = store.subscribe(selector, newState => setState(selector(newState)));
+      prevStore.current = store;
+    }
+    return () => unsubscribeRef.current();
+    // Ignoring selector as it is expected to stay constant.
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [store]);
+
+  return state;
+}