122 lines
3.8 KiB
TypeScript
122 lines
3.8 KiB
TypeScript
import type { TreeDataNode } from "antd";
|
|
import React from "react";
|
|
import { create } from "zustand";
|
|
import type { NodeContext } from "../components/Graph";
|
|
|
|
export interface TreeStore {
|
|
nodesFlatById: Map<React.Key, TreeDataNode>,
|
|
parentIdByChildId: Map<React.Key, string>;
|
|
rootNodes: TreeDataNode[];
|
|
tree: TreeDataNode[];
|
|
add: (childNode: NodeContext, parentNodeId: string | undefined) => void;
|
|
remove: (nodeId: string) => void;
|
|
reset: () => void;
|
|
}
|
|
|
|
export const useGraphLayersTreeStore = create<TreeStore>()((set) => ({
|
|
nodesFlatById: new Map<React.Key, TreeDataNode>(),
|
|
parentIdByChildId: new Map<React.Key, string>(),
|
|
rootNodes: [],
|
|
tree: [],
|
|
add: (childNodeContext, parentNodeId) => set((state) => {
|
|
const childNode = nodeContextToTreeNode(childNodeContext);
|
|
if (parentNodeId) {
|
|
const parentNode = state.nodesFlatById.get(parentNodeId);
|
|
if (parentNode) {
|
|
parentNode.children = parentNode.children ? [...parentNode.children, childNode] : [childNode];
|
|
const parentIdByChildId = new Map(state.parentIdByChildId);
|
|
parentIdByChildId.set(childNode.key, parentNodeId);
|
|
const nodesFlatById = new Map(state.nodesFlatById);
|
|
nodesFlatById.set(childNode.key, childNode);
|
|
|
|
const newState = {
|
|
nodesFlatById: nodesFlatById,
|
|
parentIdByChildId: parentIdByChildId,
|
|
rootNodes: [...state.rootNodes],
|
|
tree: createTree([...state.rootNodes], nodesFlatById)
|
|
}
|
|
|
|
return newState;
|
|
} else {
|
|
throw Error(`There is no parent node with id: ${parentNodeId}`)
|
|
}
|
|
} else {
|
|
const nodesFlatById = new Map(state.nodesFlatById);
|
|
nodesFlatById.set(childNode.key, childNode);
|
|
const newRootNodes = [...state.rootNodes, childNode];
|
|
const newState = {
|
|
nodesFlatById: nodesFlatById,
|
|
parentIdByChildId: state.parentIdByChildId,
|
|
rootNodes: newRootNodes,
|
|
tree: createTree(newRootNodes, nodesFlatById)
|
|
}
|
|
return newState;
|
|
}
|
|
}),
|
|
remove: (nodeId) => set((state) => {
|
|
const node = state.nodesFlatById.get(nodeId);
|
|
|
|
if (!node) {
|
|
return state;
|
|
}
|
|
|
|
const nodesFlatById = new Map(state.nodesFlatById);
|
|
nodesFlatById.delete(nodeId);
|
|
const parentNodeId = state.parentIdByChildId.get(nodeId);
|
|
|
|
if (parentNodeId) {
|
|
const parentNode = state.nodesFlatById.get(parentNodeId);
|
|
|
|
if (parentNode) {
|
|
parentNode.children = parentNode.children
|
|
? [...parentNode.children?.filter(n => n.key !== nodeId)]
|
|
: parentNode.children;
|
|
const parentIdByChildId = new Map(state.parentIdByChildId);
|
|
parentIdByChildId.delete(nodeId);
|
|
|
|
return {
|
|
rootNodes: [...state.rootNodes],
|
|
nodesFlatById: nodesFlatById,
|
|
parentIdByChildId: parentIdByChildId,
|
|
tree: createTree([...state.rootNodes], nodesFlatById)
|
|
}
|
|
}
|
|
}
|
|
|
|
return {
|
|
rootNodes: [...state.rootNodes.filter(n => n.key !== nodeId)],
|
|
nodesFlatById: nodesFlatById,
|
|
parentIdByChildId: state.parentIdByChildId,
|
|
tree: createTree([...state.rootNodes], nodesFlatById)
|
|
}
|
|
}),
|
|
reset: () => set({
|
|
nodesFlatById: new Map<React.Key, TreeDataNode>(),
|
|
parentIdByChildId: new Map<React.Key, string>(),
|
|
rootNodes: [],
|
|
tree: [],
|
|
}),
|
|
}));
|
|
|
|
export function createTree(nodes: TreeDataNode[], nodesFlatById: Map<React.Key, TreeDataNode>): TreeDataNode[] {
|
|
const n = [...nodes];
|
|
const result = [];
|
|
|
|
for (const node of n) {
|
|
const stateNode = nodesFlatById.get(node.key);
|
|
|
|
if (stateNode) {
|
|
stateNode.children = createTree(stateNode?.children ?? [], nodesFlatById);
|
|
result.push(stateNode);
|
|
}
|
|
}
|
|
|
|
return result;
|
|
}
|
|
|
|
export function nodeContextToTreeNode(nodeContext: NodeContext): TreeDataNode {
|
|
return {
|
|
key: nodeContext.nodeId,
|
|
title: nodeContext.nodeName,
|
|
} as TreeDataNode;
|
|
} |