import * as React from "react"; import { DataNode, FlattenNode, NodeElement, DataEntity, Key, EventDataNode } from "../interface"; import { toArray, getPosition, isTreeNode } from "../utils/common-util"; import { TreeNodeProps } from "../tree-node"; const getKey = (key: Key, pos: string) => { if (key !== null && key !== undefined) { return key; } return pos; }; /** * Convert `children` of Tree into `treeData` structure. */ const convertTreeToData = (rootNodes: React.ReactNode): DataNode[] => { const dig = (node: React.ReactNode): DataNode[] => { const treeNodes = toArray(node) as NodeElement[]; return treeNodes .map(treeNode => { //Filter invalid node if (!isTreeNode(treeNode)) { return null; } const { key } = treeNode; const { children, ...rest } = treeNode.props; const dataNode: DataNode = { key, ...rest }; const parsedChildren = dig(children); if (parsedChildren.length) { dataNode.children = parsedChildren; } return dataNode; }) .filter((dataNode: DataNode) => dataNode); }; return dig(rootNodes); }; /** * Flat nest tree data into flatten list, * this is used for virtual list render. * @param treeNodeList Origin data node list * @param expandedKeys * need expanded keys, provide `true` means all expanded */ const flattenTreeData = ( treeNodeList: DataNode[] = [], expandedKeys: Key[] | true = [] ): FlattenNode[] => { const expandedKeySet = new Set(expandedKeys === true ? [] : expandedKeys); const flattenList: FlattenNode[] = []; const dig = (list: DataNode[], parent: FlattenNode = null): FlattenNode[] => { return list.map((treeNode, index) => { const pos: string = getPosition(parent ? parent.pos : "0", index); const mergedKey = getKey(treeNode.key, pos); //Add FlattenDataNode into list const flattenNode: FlattenNode = { ...treeNode, parent, pos, children: null, data: treeNode, isStart: [...(parent ? parent.isStart : []), index === 0], isEnd: [...(parent ? parent.isEnd : []), index === list.length - 1] }; flattenList.push(flattenNode); //Loop treeNode children if (expandedKeys === true || expandedKeySet.has(mergedKey)) { flattenNode.children = dig(treeNode.children || [], flattenNode); } else { flattenNode.children = []; } return flattenNode; }); }; dig(treeNodeList); return flattenList; }; // Traverse all the data by `treeData` const traverseDataNodes = ( dataNodes: DataNode[], callback: (data: { node: DataNode; index: number; pos: string; key: Key; parentPos: string | number; level: number; }) => void ) => { const processNode = ( node: DataNode, index?: number, parent?: { node: DataNode; pos: string; level: number } ) => { const children = node ? node.children : dataNodes; const pos = node ? getPosition(parent.pos, index) : "0"; //Process node if is not root if (node) { const data = { node, index, pos, key: node.key !== null ? node.key : pos, parentPos: parent.node ? parent.pos : null, level: parent.level + 1 }; callback(data); } //Process children node if (children) { children.forEach((subNode, subIndex) => { processNode(subNode, subIndex, { node, pos, level: parent ? parent.level + 1 : -1 }); }); } }; processNode(null); }; interface Wrapper { posEntities: Record; keyEntities: Record; } //Convert `treeData` into entity records const convertDataToEntities = ( dataNodes: DataNode[], { initWrapper, processEntity, onProcessFinished }: { initWrapper?: (wrapper: Wrapper) => Wrapper; processEntity?: (entity: DataEntity, wrapper: Wrapper) => void; onProcessFinished?: (wrapper: Wrapper) => void; } = {} ) => { const posEntities = {}; const keyEntities = {}; let wrapper = { posEntities, keyEntities }; if (initWrapper) { wrapper = initWrapper(wrapper) || wrapper; } traverseDataNodes(dataNodes, item => { const { node, index, pos, key, parentPos, level } = item; const entity: DataEntity = { node, index, key, pos, level }; const mergedKey = getKey(key, pos); posEntities[pos] = entity; keyEntities[mergedKey] = entity; // Fill children entity.parent = posEntities[parentPos]; if (entity.parent) { entity.parent.children = entity.parent.children || []; entity.parent.children.push(entity); } if (processEntity) { processEntity(entity, wrapper); } }); if (onProcessFinished) { onProcessFinished(wrapper); } return wrapper; }; interface TreeNodeRequiredProps { expandedKeys: Key[]; selectedKeys: Key[]; loadedKeys: Key[]; loadingKeys: Key[]; dragOverNodeKey: Key; dropPosition: number; keyEntities: Record; } //Get TreeNode props with Tree props const getTreeNodeProps = ( key: Key, { expandedKeys, selectedKeys, loadedKeys, loadingKeys, dragOverNodeKey, dropPosition, keyEntities }: TreeNodeRequiredProps ) => { const entity = keyEntities[key]; const treeNodeProps = { eventKey: key, expanded: expandedKeys.indexOf(key) !== -1, selected: selectedKeys.indexOf(key) !== -1, loaded: loadedKeys.indexOf(key) !== -1, loading: loadingKeys.indexOf(key) !== -1, pos: String(entity ? entity.pos : ""), //Drag props dragOver: dragOverNodeKey === key && dropPosition === 0, dragOverGapTop: dragOverNodeKey === key && dropPosition === -1, dragOverGapBottom: dragOverNodeKey === key && dropPosition === 1 }; return treeNodeProps; }; //Convert Node Props to Event Data const convertNodePropsToEventData = (props: TreeNodeProps): EventDataNode => { const { data, expanded, selected, loaded, loading, dragOver, dragOverGapTop, dragOverGapBottom, pos, active } = props; const eventData = { ...data, expanded, selected, loaded, loading, dragOver, dragOverGapTop, dragOverGapBottom, pos, active }; if (!("props" in eventData)) { Object.defineProperty(eventData, "props", { get() { return props; } }); } return eventData; }; export { getKey, convertTreeToData, flattenTreeData, traverseDataNodes, convertDataToEntities, TreeNodeRequiredProps, getTreeNodeProps, convertNodePropsToEventData };