@lcap/asl/src/types/utils/traverse.ts

NetEase Application Specific Language

/**
 * 用于遍历的节点接口
 */
 export interface INode {
    children?: Array<INode>,
    [prop: string]: any,
}
/**
 * 用于遍历的节点信息
 */
export interface INodeInfo<T extends INode> {
    node: T,
    parent?: T,
    stack?: Array<T>,
    nodePath?: string,
    jsonPath?: string,
    index?: number,
    key?: string | number,
}

export interface TraverseOptions {
    mode?: 'onlyChildren' | 'onlyArray' | 'anyObject',
    depthFirst?: boolean,
    excludedKeySet?: Set<string>,
}

function handleNext<T extends INode>(
    func: (next: INodeInfo<T>) => any,
    current: INodeInfo<T>,
    options?: TraverseOptions,
) {
    if (options.mode === 'onlyChildren') {
        const list = current.node && current.node.children;
        if (!list)
            return;

        for (let index = 0; index < list.length; index++) {
            const child = list[index];
            if (!child)
                continue;

            const nodeInfo: INodeInfo<T> = {
                node: child as T,
                parent: current.node,
                stack: current.stack.concat([current.node]),
                nodePath: `${current.nodePath}/${index}`,
                jsonPath: `${current.jsonPath}.children[${index}]`.replace(/^\./, ''),
                index,
                key: index,
            };
            const result = func(nodeInfo);
            // if (result !== undefined)
            //     return result;
        }
    } else if (options.mode === 'onlyArray') {
        for (const key in current.node) {
            if (!current.node.hasOwnProperty(key) || options.excludedKeySet.has(key))
                continue;
            const value = current.node[key];
            if (!Array.isArray(value))
                continue;

            const list = value;
            for (let index = 0; index < list.length; index++) {
                const child = list[index];
                if (!child)
                    continue;

                const nodeInfo: INodeInfo<T> = {
                    node: child as T,
                    parent: current.node,
                    stack: current.stack.concat([current.node]),
                    nodePath: `${current.nodePath}/${index}`,
                    jsonPath: `${current.jsonPath}.${key}[${index}]`.replace(/^\./, ''),
                    index,
                    key: index,
                };
                const result = func(nodeInfo);
                // if (result !== undefined)
                //     return result;
            }
        }
    } else if (options.mode === 'anyObject') {
        for (const key in current.node) {
            if (!current.node.hasOwnProperty(key) || options.excludedKeySet.has(key))
                continue;
            const value = current.node[key];
            if (Array.isArray(value)) {
                const list = value;
                for (let index = 0; index < list.length; index++) {
                    const child = list[index];
                    if (!child)
                        continue;

                    const nodeInfo: INodeInfo<T> = {
                        node: child as T,
                        parent: current.node,
                        stack: current.stack.concat([current.node]),
                        nodePath: `${current.nodePath}/${index}`,
                        jsonPath: `${current.jsonPath}.${key}[${index}]`.replace(/^\./, ''),
                        index,
                        key: index,
                    };
                    const result = func(nodeInfo);
                    // if (result !== undefined)
                    //     return result;
                }
            } else if (typeof value === 'object' && value !== null) {
                const nodeInfo: INodeInfo<T> = {
                    node: value as T,
                    parent: current.node,
                    stack: current.stack.concat([current.node]),
                    nodePath: `${current.nodePath}/${key}`,
                    jsonPath: `${current.jsonPath}['${key}']`,
                    index: undefined,
                    key,
                };
                const result = func(nodeInfo);
                // if (result !== undefined)
                //     return result;
            }
        }
    }
}

/**
 * 遍历树
 * @param func
 * @param init 初始点信息
 * @param options.mode 'onlyChildren' 表示只向下遍历 children 的元素（默认）；'onlyArray' 表示只向下遍历 Array 的元素；'anyObject' 表示只要属性为 Object 类型都会遍历
 * @param options.depthFirst 默认广度优先搜索，是否深度优先搜索
 * @example traverse((current) => console.log(current.node.name), { init: rootNode });
 * @example traverse((current) => console.log(current.jsonPath), { init: ast }, { mode: 'anyObject' });
 */
export function traverse<T extends INode>(
    func: (current: INodeInfo<T>) => any,
    init: INodeInfo<T>,
    options?: TraverseOptions,
) {
    init = Object.assign({ parent: undefined, stack: [], nodePath: '', jsonPath: '', index: undefined }, init);
    options = Object.assign({ mode: 'onlyChildren', depthFirst: false, excludedKeySet: new Set() }, options);

    if (options.depthFirst) {
        func(init);
        handleNext<T>((next) => traverse(func, next, options), init, options);
    } else {
        let queue: Array<INodeInfo<T>> = [];
        queue = queue.concat(init);
        let current: INodeInfo<T>;
        while ((current = queue.shift())) {
            const result = func(current);
            if (result !== undefined)
                return result;

            handleNext<T>((next) => queue.push(next), current, options);
        }
    }
}

export default traverse;