From cf1e0bee8179d42053621e3fbea2fd7917e42ccc Mon Sep 17 00:00:00 2001
From: peterkogo <peter@xyflow.com>
Date: Thu, 17 Apr 2025 11:23:13 +0200
Subject: [PATCH] introduce generics (wip)

---
 .../container/SvelteFlow/SvelteFlow.svelte    | 11 +++---
 .../lib/container/SvelteFlow/Wrapper.svelte   |  5 ++-
 .../src/lib/container/SvelteFlow/types.ts     | 25 ++++++------
 packages/svelte/src/lib/store/index.ts        | 15 +++----
 .../src/lib/store/initial-store.svelte.ts     | 29 +++++++-------
 packages/svelte/src/lib/store/types.ts        | 39 ++++++++++++-------
 .../svelte/src/lib/store/visibleElements.ts   | 15 ++++---
 packages/svelte/src/lib/types/general.ts      | 10 ++++-
 8 files changed, 88 insertions(+), 61 deletions(-)

diff --git a/packages/svelte/src/lib/container/SvelteFlow/SvelteFlow.svelte b/packages/svelte/src/lib/container/SvelteFlow/SvelteFlow.svelte
index e501526a..4db22601 100644
--- a/packages/svelte/src/lib/container/SvelteFlow/SvelteFlow.svelte
+++ b/packages/svelte/src/lib/container/SvelteFlow/SvelteFlow.svelte
@@ -1,4 +1,5 @@
-<script lang="ts">
+<script lang="ts" generics="NodeType extends Node = Node, EdgeType extends Edge = Edge">
+  import type { Edge, Node } from '$lib/types';
   import { getContext, setContext, onDestroy } from 'svelte';
   import type { HTMLAttributes } from 'svelte/elements';
   import { ConnectionLineType, PanOnScrollMode } from '@xyflow/system';
@@ -67,10 +68,10 @@
     edges = $bindable([]),
     viewport = $bindable(undefined),
     ...props
-  }: SvelteFlowProps & HTMLAttributes<HTMLDivElement> = $props();
+  }: SvelteFlowProps<NodeType, EdgeType> & HTMLAttributes<HTMLDivElement> = $props();
 
   // svelte-ignore non_reactive_update
-  let store = createStore({
+  let store = createStore<NodeType, EdgeType>({
     props,
     width,
     height,
@@ -95,7 +96,7 @@
   });
 
   // Set store for provider context
-  const providerContext = getContext<ProviderContext>(key);
+  const providerContext = getContext<ProviderContext<NodeType, EdgeType>>(key);
   if (providerContext && providerContext.setStore) {
     providerContext.setStore(store);
   }
@@ -106,7 +107,7 @@
     getStore() {
       return store;
     }
-  } satisfies StoreContext);
+  } satisfies StoreContext<NodeType, EdgeType>);
 
   onDestroy(() => {
     store.reset();
diff --git a/packages/svelte/src/lib/container/SvelteFlow/Wrapper.svelte b/packages/svelte/src/lib/container/SvelteFlow/Wrapper.svelte
index f1f78eb1..610dce84 100644
--- a/packages/svelte/src/lib/container/SvelteFlow/Wrapper.svelte
+++ b/packages/svelte/src/lib/container/SvelteFlow/Wrapper.svelte
@@ -1,8 +1,9 @@
-<script lang="ts">
+<script lang="ts" generics="NodeType extends Node = Node, EdgeType extends Edge = Edge">
   import type { HTMLAttributes } from 'svelte/elements';
   import type { Snippet } from 'svelte';
   import { type SvelteFlowRestProps } from '$lib/store/types';
   import { toPxString } from '$lib/utils';
+  import type { Node, Edge, NodeTypes } from '$lib/types';
 
   let {
     width,
@@ -21,7 +22,7 @@
     clientWidth?: number;
     clientHeight?: number;
     children?: Snippet;
-    rest: SvelteFlowRestProps & HTMLAttributes<HTMLDivElement>;
+    rest: SvelteFlowRestProps<NodeType, EdgeType> & HTMLAttributes<HTMLDivElement>;
   } = $props();
 
   // Unfortunately we have to destructure the props here this way,
diff --git a/packages/svelte/src/lib/container/SvelteFlow/types.ts b/packages/svelte/src/lib/container/SvelteFlow/types.ts
index 117b2d9c..118e6969 100644
--- a/packages/svelte/src/lib/container/SvelteFlow/types.ts
+++ b/packages/svelte/src/lib/container/SvelteFlow/types.ts
@@ -42,7 +42,10 @@ import type {
 import type { Component } from 'svelte';
 import type { EdgeEvents, NodeEvents, NodeSelectionEvents, PaneEvents } from '$lib/types/events';
 
-export type SvelteFlowProps = NodeEvents &
+export type SvelteFlowProps<
+  NodeType extends Node = Node,
+  EdgeType extends Edge = Edge
+> = NodeEvents &
   NodeSelectionEvents &
   EdgeEvents &
   PaneEvents & {
@@ -63,7 +66,7 @@ export type SvelteFlowProps = NodeEvents &
      *  }
      * ]);
      */
-    nodes?: Node[];
+    nodes?: NodeType[];
     /** An array of edges to render in a controlled flow.
      * @example
      * const edges = $state.raw([
@@ -74,7 +77,7 @@ export type SvelteFlowProps = NodeEvents &
      *  }
      * ]);
      */
-    edges?: Edge[];
+    edges?: EdgeType[];
     /** Custom node types to be available in a flow.
      *
      * Svelte Flow matches a node's type to a component in the nodeTypes object.
@@ -130,7 +133,7 @@ export type SvelteFlowProps = NodeEvents &
      *  nodes: [{id: 'node-1'}, {id: 'node-2'}], // nodes to fit
      * };
      */
-    fitViewOptions?: FitViewOptions;
+    fitViewOptions?: FitViewOptions<NodeType>;
     /** Defines nodes relative position to its coordinates
      * @example
      * [0, 0] // default, top left
@@ -387,11 +390,11 @@ export type SvelteFlowProps = NodeEvents &
      *  console.log('deleted edges:', edges);
      * }}
      */
-    ondelete?: OnDelete;
+    ondelete?: OnDelete<NodeType, EdgeType>;
     /** This handler gets called before the user deletes nodes or edges and provides a way to abort the deletion by returning false. */
-    onbeforedelete?: OnBeforeDelete;
+    onbeforedelete?: OnBeforeDelete<NodeType, EdgeType>;
     /** This handler gets called when a new edge is created. You can use it to modify the newly created edge. */
-    onbeforeconnect?: OnBeforeConnect;
+    onbeforeconnect?: OnBeforeConnect<EdgeType>;
     /** This event gets fired when a connection successfully completes and an edge is created. */
     onconnect?: OnConnect;
     /** When a user starts to drag a connection line, this event gets fired. */
@@ -399,13 +402,13 @@ export type SvelteFlowProps = NodeEvents &
     /** When a user stops dragging a connection line, this event gets fired. */
     onconnectend?: OnConnectEnd;
     /** This event gets fired when after an edge was reconnected*/
-    onreconnect?: OnReconnect;
+    onreconnect?: OnReconnect<EdgeType>;
     /** This event gets fired when a user starts to reconnect an edge */
-    onreconnectstart?: OnRecoonnectStart;
+    onreconnectstart?: OnRecoonnectStart<EdgeType>;
     /** This event gets fired when a user stops reconnecting an edge */
-    onreconnectend?: OnReconnectEnd;
+    onreconnectend?: OnReconnectEnd<EdgeType>;
     /** This handler gets called when an edge is reconnected. You can use it to modify the edge before the update is applied. */
-    onbeforereconnect?: OnBeforeReconnect;
+    onbeforereconnect?: OnBeforeReconnect<EdgeType>;
     /** A connection is started by clicking on a handle */
     onclickconnectstart?: OnConnectStart;
     /** A connection is finished by clicking on a handle */
diff --git a/packages/svelte/src/lib/store/index.ts b/packages/svelte/src/lib/store/index.ts
index bbdcfee4..4312b2b3 100644
--- a/packages/svelte/src/lib/store/index.ts
+++ b/packages/svelte/src/lib/store/index.ts
@@ -25,8 +25,10 @@ export const key = Symbol();
 
 export { useStore } from '../hooks/useStore';
 
-export function createStore(signals: StoreSignals): SvelteFlowStore {
-  const store = getInitialStore(signals);
+export function createStore<NodeType extends Node = Node, EdgeType extends Edge = Edge>(
+  signals: StoreSignals<NodeType, EdgeType>
+): SvelteFlowStore<NodeType, EdgeType> {
+  const store = getInitialStore<NodeType, EdgeType>(signals);
 
   function setNodeTypes(nodeTypes: NodeTypes) {
     store.nodeTypes = {
@@ -42,9 +44,8 @@ export function createStore(signals: StoreSignals): SvelteFlowStore {
     };
   }
 
-  function addEdge(edgeParams: Edge | Connection) {
-    console.log(edgeParams);
-    store.edges = addEdgeUtil(edgeParams, store.edges);
+  function addEdge(edgeParams: EdgeType | Connection) {
+    store.edges = addEdgeUtil<EdgeType>(edgeParams, store.edges);
   }
 
   const updateNodePositions: UpdateNodePositions = (nodeDragItems, dragging = false) => {
@@ -76,7 +77,7 @@ export function createStore(signals: StoreSignals): SvelteFlowStore {
       store.resolveFitView();
     }
 
-    const newNodes = new Map<string, Node>();
+    const newNodes = new Map<string, NodeType>();
     for (const change of changes) {
       const userNode = store.nodeLookup.get(change.id)?.internals.userNode;
 
@@ -389,7 +390,7 @@ export function createStore(signals: StoreSignals): SvelteFlowStore {
     updateConnection,
     cancelConnection,
     reset
-  } satisfies SvelteFlowStoreActions);
+  } satisfies SvelteFlowStoreActions<NodeType, EdgeType>);
 
   return storeWithActions;
 }
diff --git a/packages/svelte/src/lib/store/initial-store.svelte.ts b/packages/svelte/src/lib/store/initial-store.svelte.ts
index ae0f9c1e..8fdc18e8 100644
--- a/packages/svelte/src/lib/store/initial-store.svelte.ts
+++ b/packages/svelte/src/lib/store/initial-store.svelte.ts
@@ -56,6 +56,7 @@ import type {
   OnBeforeDelete,
   IsValidConnection,
   Edge,
+  Node,
   EdgeLayouted,
   InternalNode,
   OnBeforeReconnect
@@ -79,7 +80,9 @@ export const initialEdgeTypes = {
   step: StepEdgeInternal
 };
 
-export const getInitialStore = (signals: StoreSignals) => {
+export function getInitialStore<NodeType extends Node = Node, EdgeType extends Edge = Edge>(
+  signals: StoreSignals<NodeType, EdgeType>
+) {
   // We use a class here, because Svelte adds getters & setter for us.
   // Inline classes have some performance implications but we just call it once (max twice).
   class SvelteFlowStore {
@@ -115,7 +118,7 @@ export const getInitialStore = (signals: StoreSignals) => {
     });
     viewportInitialized: boolean = $derived(this.panZoom !== null);
 
-    _edges: Edge[] = $derived.by(() => {
+    _edges: EdgeType[] = $derived.by(() => {
       updateConnectionLookup(this.connectionLookup, this.edgeLookup, signals.edges);
       return signals.edges;
     });
@@ -135,10 +138,10 @@ export const getInitialStore = (signals: StoreSignals) => {
       signals.edges = edges;
     }
 
-    nodeLookup: NodeLookup = new Map();
-    parentLookup: ParentLookup = new Map();
+    nodeLookup: NodeLookup<InternalNode<NodeType>> = new Map();
+    parentLookup: ParentLookup<InternalNode<NodeType>> = new Map();
     connectionLookup: ConnectionLookup = new Map();
-    edgeLookup: EdgeLookup = new Map();
+    edgeLookup: EdgeLookup<EdgeType> = new Map();
 
     _prevVisibleEdges = new Map<string, EdgeLayouted>();
     visible = $derived.by(() => {
@@ -184,7 +187,7 @@ export const getInitialStore = (signals: StoreSignals) => {
         });
       } else {
         visibleNodes = this.nodeLookup;
-        visibleEdges = getLayoutedEdges(options as EdgeLayoutAllOptions);
+        visibleEdges = getLayoutedEdges(options as EdgeLayoutAllOptions<NodeType, EdgeType>);
       }
 
       return {
@@ -286,18 +289,18 @@ export const getInitialStore = (signals: StoreSignals) => {
     onlyRenderVisibleElements: boolean = $derived(signals.props.onlyRenderVisibleElements ?? false);
     onerror: OnError = $derived(signals.props.onflowerror ?? devWarn);
 
-    ondelete?: OnDelete = $derived(signals.props.ondelete);
-    onbeforedelete?: OnBeforeDelete = $derived(signals.props.onbeforedelete);
+    ondelete?: OnDelete<NodeType, EdgeType> = $derived(signals.props.ondelete);
+    onbeforedelete?: OnBeforeDelete<NodeType, EdgeType> = $derived(signals.props.onbeforedelete);
 
     onbeforeconnect?: OnBeforeConnect = $derived(signals.props.onbeforeconnect);
     onconnect?: OnConnect = $derived(signals.props.onconnect);
     onconnectstart?: OnConnectStart = $derived(signals.props.onconnectstart);
     onconnectend?: OnConnectEnd = $derived(signals.props.onconnectend);
 
-    onbeforereconnect?: OnBeforeReconnect = $derived(signals.props.onbeforereconnect);
-    onreconnect?: OnReconnect = $derived(signals.props.onreconnect);
-    onreconnectstart?: OnRecoonnectStart = $derived(signals.props.onreconnectstart);
-    onreconnectend?: OnReconnectEnd = $derived(signals.props.onreconnectend);
+    onbeforereconnect?: OnBeforeReconnect<EdgeType> = $derived(signals.props.onbeforereconnect);
+    onreconnect?: OnReconnect<EdgeType> = $derived(signals.props.onreconnect);
+    onreconnectstart?: OnRecoonnectStart<EdgeType> = $derived(signals.props.onreconnectstart);
+    onreconnectend?: OnReconnectEnd<EdgeType> = $derived(signals.props.onreconnectend);
 
     clickConnect?: boolean = $derived(signals.props.clickConnect ?? true);
     onclickconnectstart?: OnConnectStart = $derived(signals.props.onclickconnectstart);
@@ -374,7 +377,7 @@ export const getInitialStore = (signals: StoreSignals) => {
     }
   }
   return new SvelteFlowStore();
-};
+}
 
 // Only way to check if an object is a proxy
 // is to see if is failes to perform a structured clone
diff --git a/packages/svelte/src/lib/store/types.ts b/packages/svelte/src/lib/store/types.ts
index 4f8c9b07..d6efb023 100644
--- a/packages/svelte/src/lib/store/types.ts
+++ b/packages/svelte/src/lib/store/types.ts
@@ -13,10 +13,10 @@ import type { getInitialStore } from './initial-store.svelte';
 import type { Node, Edge, NodeTypes, EdgeTypes, FitViewOptions } from '$lib/types';
 import type { SvelteFlowProps } from '$lib/container/SvelteFlow';
 
-export type SvelteFlowStoreActions = {
+export type SvelteFlowStoreActions<NodeType extends Node = Node, EdgeType extends Edge = Edge> = {
   setNodeTypes: (nodeTypes: NodeTypes) => void;
   setEdgeTypes: (edgeTypes: EdgeTypes) => void;
-  addEdge: (edge: Edge | Connection) => void;
+  addEdge: (edge: EdgeType | Connection) => void;
   zoomIn: (options?: ViewportHelperFunctionOptions) => Promise<boolean>;
   zoomOut: (options?: ViewportHelperFunctionOptions) => Promise<boolean>;
   setMinZoom: (minZoom: number) => void;
@@ -26,7 +26,7 @@ export type SvelteFlowStoreActions = {
   fitView: (options?: FitViewOptions) => Promise<boolean>;
   updateNodePositions: UpdateNodePositions;
   updateNodeInternals: (updates: Map<string, InternalNodeUpdate>) => void;
-  unselectNodesAndEdges: (params?: { nodes?: Node[]; edges?: Edge[] }) => void;
+  unselectNodesAndEdges: (params?: { nodes?: NodeType[]; edges?: EdgeType[] }) => void;
   addSelectedNodes: (ids: string[]) => void;
   addSelectedEdges: (ids: string[]) => void;
   handleNodeSelection: (id: string, unselect?: boolean, nodeRef?: HTMLDivElement | null) => void;
@@ -38,8 +38,8 @@ export type SvelteFlowStoreActions = {
   reset(): void;
 };
 
-export type SvelteFlowRestProps = Omit<
-  SvelteFlowProps,
+export type SvelteFlowRestProps<NodeType extends Node = Node, EdgeType extends Edge = Edge> = Omit<
+  SvelteFlowProps<NodeType, EdgeType>,
   | 'width'
   | 'height'
   | 'class'
@@ -89,24 +89,33 @@ export type SvelteFlowRestProps = Omit<
   | 'viewport'
 >;
 
-export type StoreSignals = {
-  props: SvelteFlowRestProps;
+export type StoreSignals<NodeType extends Node = Node, EdgeType extends Edge = Edge> = {
+  props: SvelteFlowRestProps<NodeType, EdgeType>;
   width?: number;
   height?: number;
-  nodes: Node[];
-  edges: Edge[];
+  nodes: NodeType[];
+  edges: EdgeType[];
   viewport?: Viewport;
 };
 
-export type SvelteFlowStoreState = ReturnType<typeof getInitialStore>;
+export type SvelteFlowStoreState<
+  NodeType extends Node = Node,
+  EdgeType extends Edge = Edge
+> = ReturnType<typeof getInitialStore<NodeType, EdgeType>>;
 
-export type SvelteFlowStore = SvelteFlowStoreState & SvelteFlowStoreActions;
+export type SvelteFlowStore<
+  NodeType extends Node = Node,
+  EdgeType extends Edge = Edge
+> = SvelteFlowStoreState<NodeType, EdgeType> & SvelteFlowStoreActions<NodeType, EdgeType>;
 
-export type StoreContext = {
-  getStore: () => SvelteFlowStore;
+export type StoreContext<NodeType extends Node = Node, EdgeType extends Edge = Edge> = {
+  getStore: () => SvelteFlowStore<NodeType, EdgeType>;
   provider: boolean;
 };
 
-export type ProviderContext = StoreContext & {
-  setStore: (store: SvelteFlowStore) => void;
+export type ProviderContext<
+  NodeType extends Node = Node,
+  EdgeType extends Edge = Edge
+> = StoreContext & {
+  setStore: (store: SvelteFlowStore<NodeType, EdgeType>) => void;
 };
diff --git a/packages/svelte/src/lib/store/visibleElements.ts b/packages/svelte/src/lib/store/visibleElements.ts
index b8dc4462..373d5ac1 100644
--- a/packages/svelte/src/lib/store/visibleElements.ts
+++ b/packages/svelte/src/lib/store/visibleElements.ts
@@ -1,4 +1,4 @@
-import type { DefaultEdgeOptions, Edge, EdgeLayouted, InternalNode } from '$lib/types';
+import type { DefaultEdgeOptions, Node, Edge, EdgeLayouted, InternalNode } from '$lib/types';
 import {
   ConnectionMode,
   getEdgePosition,
@@ -25,17 +25,18 @@ export function getVisibleNodes(
   return visibleNodes;
 }
 
-export interface EdgeLayoutBaseOptions {
-  edges: Edge[];
+export interface EdgeLayoutBaseOptions<NodeType extends Node = Node, EdgeType extends Edge = Edge> {
+  edges: EdgeType[];
   defaultEdgeOptions: DefaultEdgeOptions;
   elevateEdgesOnSelect: boolean;
   previousEdges: Map<string, EdgeLayouted>;
-  nodeLookup: NodeLookup;
+  nodeLookup: NodeLookup<InternalNode<NodeType>>;
   connectionMode: ConnectionMode;
   onerror: OnError;
 }
 
-export interface EdgeLayoutAllOptions extends EdgeLayoutBaseOptions {
+export interface EdgeLayoutAllOptions<NodeType extends Node = Node, EdgeType extends Edge = Edge>
+  extends EdgeLayoutBaseOptions<NodeType, EdgeType> {
   onlyRenderVisible: never;
   visibleNodes: never;
   transform: never;
@@ -51,7 +52,9 @@ export interface EdgeLayoutOnlyVisibleOptions extends EdgeLayoutBaseOptions {
   onlyRenderVisible: true;
 }
 
-export type EdgeLayoutOptions = EdgeLayoutAllOptions | EdgeLayoutOnlyVisibleOptions;
+export type EdgeLayoutOptions<NodeType extends Node = Node, EdgeType extends Edge = Edge> =
+  | EdgeLayoutAllOptions<NodeType, EdgeType>
+  | EdgeLayoutOnlyVisibleOptions;
 
 export function getLayoutedEdges(options: EdgeLayoutOptions): Map<string, EdgeLayouted> {
   const {
diff --git a/packages/svelte/src/lib/types/general.ts b/packages/svelte/src/lib/types/general.ts
index 27129228..76a7afe5 100644
--- a/packages/svelte/src/lib/types/general.ts
+++ b/packages/svelte/src/lib/types/general.ts
@@ -23,8 +23,14 @@ export type ConnectionData = {
 
 export type FitViewOptions<NodeType extends Node = Node> = FitViewOptionsBase<NodeType>;
 
-export type OnDelete = (params: { nodes: Node[]; edges: Edge[] }) => void;
-export type OnBeforeConnect = (connection: Connection) => Edge | Connection | void;
+export type OnDelete<NodeType extends Node = Node, EdgeType extends Edge = Edge> = (params: {
+  nodes: NodeType[];
+  edges: EdgeType[];
+}) => void;
+
+export type OnBeforeConnect<EdgeType extends Edge = Edge> = (
+  connection: Connection
+) => EdgeType | Connection | void;
 export type OnBeforeReconnect<EdgeType extends Edge = Edge> = (
   newEdge: EdgeType,
   oldEdge: EdgeType