feat: serializeGraph accepts an identity function to prevent data duplication

Author: JesusTheHun

src/types.ts

@@ -1,15 +1,15 @@
 export type EdgeWeight = number;
 
-export type Edge<Node = unknown, Props = unknown> = {
-  source: Node;
-  target: Node;
+export type Edge<NodeIdentity = unknown, Props = unknown> = {
+  source: NodeIdentity;
+  target: NodeIdentity;
   weight?: EdgeWeight;
   props: Props;
 };
 
-export type Serialized<Node = unknown, LinkProps = unknown> = {
+export type Serialized<Node = unknown, LinkProps = unknown, NodeIdentity = Node> = {
   nodes: Node[];
-  links: Edge<Node, LinkProps>[];
+  links: Edge<NodeIdentity, LinkProps>[];
 };
 
 export type SerializedInput<Node = unknown, LinkProps = unknown> = {

src/utils/serializeGraph.spec.ts

@@ -1,4 +1,4 @@
-import { describe, expectTypeOf, it } from 'vitest';
+import { describe, expect, expectTypeOf, it } from 'vitest';
 import { Graph } from '../Graph.js';
 import { checkSerialized } from '../test-utils.js';
 import { Serialized } from '../types.js';
@@ -13,6 +13,44 @@ describe('serializeGraph', () => {
     checkSerialized(serialized);
   });
 
+  it('should use the node identity for link serialization', function () {
+    const nodeA = { id: 1, title: 'a' };
+    const nodeB = { id: 2, title: 'b' };
+
+    const graph = new Graph<{ id: number; title: string }, { type: string }>();
+    graph.addEdge(nodeA, nodeB, { props: { type: 'foo' } });
+
+    const serialized = serializeGraph(graph, (n) => n.id);
+
+    expect(serialized).toStrictEqual({
+      nodes: [nodeA, nodeB],
+      links: [{ source: 1, target: 2, props: { type: 'foo' } }],
+    });
+  });
+
+  it('should reuse the same identity when the node is met multiple times', function () {
+    const nodeA = { id: 1, title: 'a' };
+    const nodeB = { id: 2, title: 'b' };
+    const nodeC = { id: 3, title: 'c' };
+
+    const graph = new Graph<{ id: number; title: string }>();
+    graph.addEdge(nodeA, nodeC);
+    graph.addEdge(nodeB, nodeC);
+
+    // we use an object as identity
+    const serialized = serializeGraph(graph, (n) => ({ id: n.id }));
+
+    const nodeIdentityC1 = serialized.links.find(
+      (l) => l.source.id === nodeA.id && l.target.id === nodeC.id,
+    )?.target;
+    const nodeIdentityC2 = serialized.links.find(
+      (l) => l.source.id === nodeB.id && l.target.id === nodeC.id,
+    )?.target;
+
+    expect(nodeIdentityC1).toBeDefined();
+    expect(nodeIdentityC1).toBe(nodeIdentityC2);
+  });
+
   it.skip('should return a serialized input with type inferred from the graph', function () {
     const nodeA = { title: 'a' };
     const nodeB = { title: 'b' };

src/utils/serializeGraph.ts

@@ -2,7 +2,7 @@
  * Serializes the graph.
  */
 import { Graph } from '../Graph.js';
-import { Edge, Serialized } from '../types.js';
+import { NoInfer, Edge, Serialized } from '../types.js';
 
 type SerializeGraphOptions<IncludeDefaultWeight extends boolean = false> = {
   /**
@@ -17,30 +17,57 @@ type SerializeGraphOptions<IncludeDefaultWeight extends boolean = false> = {
 
 /**
  * Serialize the graph data set : nodes, edges, edges weight & properties.
- * @param graph
- * @param opts
+ *
+ * Optionally, you can pass a function that returns a unique value for a given node.
+ * When provided, the function will be used to avoid data duplication in the serialized object.
  */
-export function serializeGraph<Node, LinkProps, IncludeDefaultWeight extends boolean>(
+export function serializeGraph<
+  Node,
+  LinkProps,
+  IncludeDefaultWeight extends boolean,
+  NodeIdentity = Node,
+>(
   graph: Graph<Node, LinkProps>,
-  opts: SerializeGraphOptions<IncludeDefaultWeight> = {},
-): Serialized<Node, LinkProps> {
-  const { includeDefaultWeight = false } = opts;
+  ...args:
+    | [
+        identityFn: (node: NoInfer<Node>) => NodeIdentity,
+        SerializeGraphOptions<IncludeDefaultWeight>?,
+      ]
+    | [SerializeGraphOptions<IncludeDefaultWeight>?]
+): Serialized<Node, LinkProps, NodeIdentity> {
+  const identityFn = typeof args[0] === 'function' ? args[0] : undefined;
+  const opts = typeof args[0] === 'function' ? args[1] : args[0];
 
-  const serialized: Serialized<Node, LinkProps> = {
+  const { includeDefaultWeight = false } = opts ?? {};
+
+  const serialized: Serialized<Node, LinkProps, NodeIdentity> = {
     nodes: Array.from(graph.nodes),
     links: [],
   };
 
+  const nodeIdentityMap = new Map<Node, NodeIdentity>();
+
   serialized.nodes.forEach((node) => {
     const source = node;
     graph.adjacent(source)?.forEach((target) => {
       const edgeWeight = graph.getEdgeWeight(source, target);
       const edgeProps = graph.getEdgeProperties(source, target);
 
+      if (identityFn && !nodeIdentityMap.has(source)) {
+        nodeIdentityMap.set(source, identityFn(source));
+      }
+
+      if (identityFn && !nodeIdentityMap.has(target)) {
+        nodeIdentityMap.set(target, identityFn(target));
+      }
+
+      const sourceIdentity = nodeIdentityMap.get(source) ?? source;
+      const targetIdentity = nodeIdentityMap.get(target) ?? target;
+
       const link = {
-        source: source,
-        target: target,
-      } as Edge<Node, LinkProps>;
+        source: sourceIdentity,
+        target: targetIdentity,
+      } as Edge<NodeIdentity, LinkProps>;
 
       if (edgeWeight != 1 || includeDefaultWeight) {
         link.weight = edgeWeight;