fix: react useLoroStore

packages/react/package.json

@@ -48,9 +48,12 @@
         "react": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0"
     },
     "devDependencies": {
+        "@testing-library/react": "^16.3.0",
         "@types/node": "^20.10.5",
         "@types/react": "^18.2.45",
+        "jsdom": "^23.0.1",
         "react": "^18.2.0",
+        "react-dom": "^18.2.0",
         "typescript": "^5.3.3"
     }
 }

packages/react/src/hooks.tsx

@@ -81,34 +81,42 @@ export function useLoroStore<S extends SchemaType>(
     options: UseLoroStoreOptions<S>,
 ) {
     // Create a stable reference to the store
-    const storeRef = useRef<Mirror<S> | null>(null);
+    const storeRef = useRef<{ mirror: Mirror<S>; doc: LoroDoc } | null>(null);
 
     // Initialize the store and get initial state
-    const getStore = useCallback((): Mirror<S> => {
+    const getMirror = useCallback((): Mirror<S> => {
         let store = storeRef.current;
-        if (!store) {
-            store = new Mirror(options);
+        if (!store || store.doc !== options.doc) {
+            store = { mirror: new Mirror(options), doc: options.doc };
             storeRef.current = store;
         }
-        return store;
-    }, [options]);
+        return store.mirror;
+    }, [
+        options.doc,
+        options.debug,
+        options.initialState,
+        options.schema,
+        options.throwOnValidationError,
+        options.validateUpdates,
+    ]);
 
     // Get the current state
     const [state, setLocalState] = useState<InferType<S>>(() => {
-        return getStore().getState();
+        return getMirror().getState();
     });
 
     // Subscribe to state changes
     useEffect(() => {
-        const store = getStore();
+        const store = getMirror();
 
         // Update local state when the store changes
         const unsubscribe = store.subscribe((newState: InferType<S>) => {
             setLocalState(newState);
         });
 
+        setLocalState(getMirror().getState());
         return unsubscribe;
-    }, [getStore]);
+    }, [getMirror]);
 
     // Create a stable setState function
     type SetStateFn = {
@@ -120,15 +128,15 @@ export function useLoroStore<S extends SchemaType>(
     };
     const setState: SetStateFn = useCallback(
         (updater: unknown) => {
-            getStore().setState(updater as never);
+            getMirror().setState(updater as never);
         },
-        [getStore],
+        [getMirror],
     ) as unknown as SetStateFn;
 
     return {
         state,
         setState,
-        store: getStore(),
+        store: getMirror(),
     };
 }
 
@@ -160,6 +168,7 @@ export function useLoroValue<S extends SchemaType, R>(
 
     // Subscribe to changes
     useEffect(() => {
+        setValue(selector(store.getState()));
         const unsubscribe = store.subscribe((state: InferType<S>) => {
             const newValue = selector(state);
             setValue(newValue);
@@ -312,6 +321,7 @@ export function createLoroContext<S extends SchemaType>(schema: S) {
         const [state, setState] = useState(store.getState());
 
         useEffect(() => {
+            setState(store.getState());
             const unsubscribe = store.subscribe((newState: InferType<S>) => {
                 setState(newState);
             });

packages/react/tests/doc-switch.test.tsx

@@ -0,0 +1,114 @@
+// @vitest-environment jsdom
+import React from "react";
+import { describe, expect, it } from "vitest";
+import { render, screen, waitFor } from "@testing-library/react";
+import { LoroDoc } from "loro-crdt";
+import { Mirror, schema } from "loro-mirror";
+import { createLoroContext, useLoroStore } from "../src/index.js";
+
+const counterSchema = schema({
+    data: schema.LoroMap({
+        counter: schema.Number({ defaultValue: 0 }),
+    }),
+});
+
+function seedDoc(doc: LoroDoc, counter: number) {
+    const mirror = new Mirror({
+        doc,
+        schema: counterSchema,
+        initialState: { data: { counter: 0 } },
+    });
+    mirror.setState((s) => {
+        s.data.counter = counter;
+    });
+    mirror.dispose();
+}
+
+describe("React hooks - switching LoroDoc", () => {
+    it("useLoroStore updates state when switching to a different doc", async () => {
+        const docA = new LoroDoc();
+        const docB = new LoroDoc();
+
+        seedDoc(docA, 1);
+        seedDoc(docB, 2);
+
+        function CounterView({ doc }: { doc: LoroDoc }) {
+            const { state } = useLoroStore({
+                doc,
+                schema: counterSchema,
+                initialState: { data: { counter: 0 } },
+            });
+
+            return <div data-testid="counter">{state.data.counter}</div>;
+        }
+
+        const { rerender } = render(<CounterView doc={docA} />);
+        await waitFor(() =>
+            expect(screen.getByTestId("counter").textContent).toBe("1"),
+        );
+
+        rerender(<CounterView doc={docB} />);
+        await waitFor(() =>
+            expect(screen.getByTestId("counter").textContent).toBe("2"),
+        );
+    });
+
+    it("createLoroContext hooks reflect the new doc state after provider doc changes", async () => {
+        const docA = new LoroDoc();
+        const docB = new LoroDoc();
+
+        seedDoc(docA, 10);
+        seedDoc(docB, 20);
+
+        const { LoroProvider, useLoroState, useLoroSelector } =
+            createLoroContext(counterSchema);
+
+        function StateView() {
+            const [state] = useLoroState();
+            return <div data-testid="counter">{state.data.counter}</div>;
+        }
+
+        function SelectorView() {
+            const counter = useLoroSelector((s) => s.data.counter);
+            return <div data-testid="counter">{counter}</div>;
+        }
+
+        const stateView = render(
+            <LoroProvider doc={docA} initialState={{ data: { counter: 0 } }}>
+                <StateView />
+            </LoroProvider>,
+        );
+        await waitFor(() =>
+            expect(screen.getByTestId("counter").textContent).toBe("10"),
+        );
+
+        stateView.rerender(
+            <LoroProvider doc={docB} initialState={{ data: { counter: 0 } }}>
+                <StateView />
+            </LoroProvider>,
+        );
+        await waitFor(() =>
+            expect(screen.getByTestId("counter").textContent).toBe("20"),
+        );
+        stateView.unmount();
+
+        const selectorView = render(
+            <LoroProvider doc={docA} initialState={{ data: { counter: 0 } }}>
+                <SelectorView />
+            </LoroProvider>,
+        );
+        await waitFor(() =>
+            expect(screen.getByTestId("counter").textContent).toBe("10"),
+        );
+
+        selectorView.rerender(
+            <LoroProvider doc={docB} initialState={{ data: { counter: 0 } }}>
+                <SelectorView />
+            </LoroProvider>,
+        );
+        await waitFor(() =>
+            expect(screen.getByTestId("counter").textContent).toBe("20"),
+        );
+        selectorView.unmount();
+    });
+});