From aae37a541651c9337e549fe9ee54597e77764de1 Mon Sep 17 00:00:00 2001
From: Braks <78412429+bcakmakoglu@users.noreply.github.com>
Date: Sat, 9 Apr 2022 20:30:42 +0200
Subject: [PATCH] feat(flow): Allow add and set actions to have a setter input

---
 package/src/store/actions.ts | 25 ++++++++++++++++++-------
 package/src/types/store.ts   | 26 +++++++++++++++++++-------
 2 files changed, 37 insertions(+), 14 deletions(-)

diff --git a/package/src/store/actions.ts b/package/src/store/actions.ts
index 93076db0..35de54ed 100644
--- a/package/src/store/actions.ts
+++ b/package/src/store/actions.ts
@@ -225,13 +225,15 @@ export default (state: State, getters: ComputedGetters): Actions => {
   const setNodes: Actions['setNodes'] = (nodes, extent?: CoordinateExtent) => {
     if (!state.initialized && !nodes.length) return
     if (!state.nodes) state.nodes = []
-    state.nodes = createGraphNodes(nodes, getters.getNode.value, state.nodes, extent ?? state.nodeExtent)
+    const curr = nodes instanceof Function ? nodes(state.nodes) : nodes
+    state.nodes = createGraphNodes(curr, getters.getNode.value, state.nodes, extent ?? state.nodeExtent)
   }
 
   const setEdges: Actions['setEdges'] = (edges) => {
     if (!state.initialized && !edges.length) return
+    const curr = edges instanceof Function ? edges(state.edges) : edges
 
-    state.edges = edges.reduce<GraphEdge[]>((res, edge) => {
+    state.edges = curr.reduce<GraphEdge[]>((res, edge) => {
       const sourceNode = getters.getNode.value(edge.source)!
       const targetNode = getters.getNode.value(edge.target)!
 
@@ -257,16 +259,22 @@ export default (state: State, getters: ComputedGetters): Actions => {
 
   const setElements: Actions['setElements'] = (elements, extent) => {
     if ((!state.initialized && !elements.length) || !elements) return
-    setNodes(elements.filter(isNode), extent)
-    setEdges(elements.filter(isEdge))
+    const curr = elements instanceof Function ? elements([...state.nodes, ...state.edges]) : elements
+
+    setNodes(curr.filter(isNode), extent)
+    setEdges(curr.filter(isEdge))
   }
 
   const addNodes: Actions['addNodes'] = (nodes, extent) => {
-    state.nodes.push(...createGraphNodes(nodes, getters.getNode.value, state.nodes, extent ?? state.nodeExtent))
+    const curr = nodes instanceof Function ? nodes(state.nodes) : nodes
+
+    state.nodes.push(...createGraphNodes(curr, getters.getNode.value, state.nodes, extent ?? state.nodeExtent))
   }
 
   const addEdges: Actions['addEdges'] = (params) => {
-    params.forEach((param) => {
+    const curr = params instanceof Function ? params(state.edges) : params
+
+    curr.forEach((param) => {
       const edge = addEdge(param, state.edges)
       if (edge) {
         const sourceNode = getters.getNode.value(edge.source)!
@@ -294,11 +302,14 @@ export default (state: State, getters: ComputedGetters): Actions => {
 
   const applyEdgeChanges: Actions['applyEdgeChanges'] = (changes) => applyEdges(changes, state.edges)
 
-  const setState: Actions['setState'] = (opts) => {
+  const setState: Actions['setState'] = (options) => {
     const skip = ['modelValue', 'nodes', 'edges', 'maxZoom', 'minZoom', 'translateExtent']
+    const opts = options instanceof Function ? options(state) : options
+
     if (typeof opts.modelValue !== 'undefined') setElements(opts.modelValue, opts.nodeExtent ?? state.nodeExtent)
     if (typeof opts.nodes !== 'undefined') setNodes(opts.nodes, opts.nodeExtent ?? state.nodeExtent)
     if (typeof opts.edges !== 'undefined') setEdges(opts.edges)
+
     Object.keys(opts).forEach((o) => {
       const option = opts[o as keyof typeof opts]
       if (!skip.includes(o) && isDef(option)) (<any>state)[o] = option
diff --git a/package/src/types/store.ts b/package/src/types/store.ts
index 0ae78186..8dc8cc53 100644
--- a/package/src/types/store.ts
+++ b/package/src/types/store.ts
@@ -95,16 +95,21 @@ export interface State<NodeData = ElementData, EdgeData = ElementData>
 }
 
 export interface Actions<NodeData = ElementData, EdgeData = ElementData> {
-  /** @deprecated use setNodes / setEdges instead */
-  setElements: (elements: Elements<NodeData, EdgeData>, extent?: CoordinateExtent) => void
+  /** parses elements (nodes + edges) and re-sets the state */
+  setElements: (
+    elements: Elements<NodeData, EdgeData> | ((elements: FlowElements<NodeData, EdgeData>) => Elements<NodeData>),
+    extent?: CoordinateExtent,
+  ) => void
   /** parses nodes and re-sets the state */
-  setNodes: (nodes: Node<NodeData>[], extent?: CoordinateExtent) => void
+  setNodes: (nodes: Node<NodeData>[] | ((nodes: GraphNode<NodeData>[]) => Node<NodeData>[]), extent?: CoordinateExtent) => void
   /** parses edges and re-sets the state */
-  setEdges: (edges: Edge<EdgeData>[]) => void
+  setEdges: (edges: Edge<EdgeData>[] | ((edges: GraphEdge<EdgeData>[]) => Edge<EdgeData>[])) => void
   /** parses nodes and adds to state */
-  addNodes: <NA = NodeData>(nodes: Node<NA>[], extent?: CoordinateExtent) => void
+  addNodes: <NA = NodeData>(nodes: Node<NA>[] | ((nodes: GraphNode<NA>[]) => Node<NA>[]), extent?: CoordinateExtent) => void
   /** parses edges and adds to state */
-  addEdges: <EA = EdgeData>(edgesOrConnections: (Edge<EA> | Connection)[]) => void
+  addEdges: <EA = EdgeData>(
+    edgesOrConnections: (Edge<EA> | Connection)[] | ((edges: GraphEdge<EA>[]) => (Edge<EA> | Connection)[]),
+  ) => void
   /** updates an edge */
   updateEdge: <EU = EdgeData>(oldEdge: GraphEdge<EU>, newConnection: Connection) => GraphEdge | false
   applyEdgeChanges: <ED = EdgeData>(changes: EdgeChange[]) => GraphEdge<ED>[]
@@ -117,7 +122,14 @@ export interface Actions<NodeData = ElementData, EdgeData = ElementData> {
   setTranslateExtent: (translateExtent: CoordinateExtent) => void
   resetSelectedElements: () => void
   setInteractive: (isInteractive: boolean) => void
-  setState: (state: Partial<FlowOptions<NodeData, EdgeData> & Omit<State, 'nodes' | 'edges' | 'modelValue'>>) => void
+  setState: (
+    state:
+      | Partial<FlowOptions<NodeData, EdgeData> & Omit<State, 'nodes' | 'edges' | 'modelValue'>>
+      | ((
+          state: State<NodeData, EdgeData>,
+        ) => Partial<FlowOptions<NodeData, EdgeData> & Omit<State, 'nodes' | 'edges' | 'modelValue'>>),
+  ) => void
+
   updateNodePosition: ({ id, diff, dragging }: { id?: string; diff?: XYPosition; dragging?: boolean }) => void
   updateNodeDimensions: (
     updates: {