Merge pull request #82 from gadget-inc/fix-snapshot-references-in-maps

Ensure we properly snapshot reference types in maps/arrays

Author: thegedge

spec/getSnapshot.spec.ts

@@ -1,6 +1,6 @@
 import type { IsExact } from "conditional-type-checks";
 import { assert } from "conditional-type-checks";
-import type { SnapshotOut } from "../src";
+import type { IReferenceType, SnapshotOut } from "../src";
 import { types } from "../src";
 import { getSnapshot } from "../src/api";
 import { TestClassModel } from "./fixtures/TestClassModel";
@@ -99,4 +99,35 @@ describe("getSnapshot", () => {
     verifySnapshot(snapshot);
     assert<IsExact<typeof snapshot.map.test_key, SnapshotOut<typeof NamedThing>>>(true);
   });
+
+  test("snapshots a reference type correctly in a map / array", () => {
+    const RootType = types.model({
+      things: types.map(NamedThing),
+      refMap: types.map(types.reference(NamedThing)),
+      refArray: types.array(types.reference(NamedThing)),
+    });
+
+    const instance = RootType.createReadOnly({
+      things: {
+        a: { key: "a", name: "test a" },
+        b: { key: "b", name: "test b" },
+      },
+      refMap: { a: "a", b: "b" },
+      refArray: ["a", "b"],
+    });
+
+    const snapshot = getSnapshot(instance);
+    assert<IsExact<typeof snapshot.things.test_key, SnapshotOut<typeof NamedThing>>>(true);
+    assert<IsExact<typeof snapshot.refMap.test_key, SnapshotOut<IReferenceType<typeof NamedThing>>>>(true);
+    expect(snapshot.refMap.a).toEqual("a");
+    expect(snapshot.refMap.b).toEqual("b");
+    expect(snapshot.refArray[0]).toEqual("a");
+    expect(snapshot.refArray[1]).toEqual("b");
+
+    const reconstructedInstance = RootType.createReadOnly(snapshot);
+    expect(reconstructedInstance.refMap.get("a")?.name).toEqual("test a");
+    expect(reconstructedInstance.refMap.get("b")?.name).toEqual("test b");
+    expect(reconstructedInstance.refArray[0]?.name).toEqual("test a");
+    expect(reconstructedInstance.refArray[1]?.name).toEqual("test b");
+  });
 });

src/snapshot.ts

@@ -1,8 +1,8 @@
 import type { IStateTreeNode as MSTStateTreeNode } from "mobx-state-tree";
 import { getSnapshot as mstGetSnapshot, isStateTreeNode as mstIsStateTreeNode } from "mobx-state-tree";
 import { getType, isModelType, isReferenceType, isStateTreeNode } from "./api";
-import { QuickArray } from "./array";
-import { QuickMap } from "./map";
+import { ArrayType, QuickArray } from "./array";
+import { MapType, QuickMap } from "./map";
 import { $identifier } from "./symbols";
 import type { IAnyType, IStateTreeNode, SnapshotOut } from "./types";
 
@@ -16,11 +16,21 @@ export function getSnapshot<T extends IAnyType>(value: IStateTreeNode<T>): Snaps
 
 const snapshot = (value: any): unknown => {
   if (value instanceof QuickArray) {
-    return Array.from(value.map((v) => snapshot(v)));
+    const type = getType(value) as ArrayType<IAnyType>;
+    const childrenAreReferences = isReferenceType(type.childrenType);
+
+    return Array.from(value.map((v) => (childrenAreReferences ? v[$identifier] : snapshot(v))));
   }
 
   if (value instanceof QuickMap) {
-    return Object.fromEntries(Array.from(value.entries()).map(([k, v]) => [k, snapshot(v)]));
+    const type = getType(value) as MapType<IAnyType>;
+    const childrenAreReferences = isReferenceType(type.childrenType);
+
+    return Object.fromEntries(
+      Array.from(value.entries()).map(([k, v]) => {
+        return [k, childrenAreReferences ? v[$identifier] : snapshot(v)];
+      })
+    );
   }
 
   if (value instanceof Date) {