react antd tree 前端实现搜索展开并加亮关键词

大概效果如图：
搜索时自动展开并加亮。antd也有大概实现该功能。但其展开逻辑有点小问题。当搜索的是子类时未进行递归寻找父级id。导致没有实现展开功能，只实现了加亮。

import React, { Component } from 'react';
import { connect } from 'react-redux';
import { Card, Input, Spin, Tree, Icon, Modal, message } from '@whalecloud/fdx';

const Search = Input.Search;
const { TreeNode } = Tree;

// 数组去重
const unique = arr => {
  return Array.from(new Set(arr))
}

// 搜索父级ID。此处没有使用上。当树的数据每一项数据中没返回父级id时，就需要使用该方法获取父级id
const getParentKey = (key, tree) => {
  let parentKey;
  for (let i = 0; i < tree.length; i++) {
    const node = tree[i];
    if (node.children) {
      if (node.children.some(item => item.treeId === key)) {
        parentKey = node.treeId;
      } else if (getParentKey(key, node.children)) {
        parentKey = getParentKey(key, node.children);
      }
    }
  }
  return parentKey;
};

class Catalog extends Component {
  state = {
    treeList: [], // 树的数据。此处树id为treeId的。树名称为treeName。树的父级id为parentid
    treeLoading: false,
    expandedKeys: [], // 记录树展开的id
    searchValue: '', // 搜索值
    autoExpandParent: true,
  }

  onExpand = (expandedKeys) => {
    this.setState({
      expandedKeys,
      autoExpandParent: false,
    });
  }

  onSearch = value => {
    let expandedKeys = [];
    // 搜索时，寻找需展开的节点
    const getExpandKeys = arr => {
      arr.forEach(item => {
        if (item.treeName.indexOf(value) > -1) {
          expandedKeys.push(item.parentid); // 此处每一项的数据就有父级id，所以直接把父级id塞进展开数组即可
          // 倘若每一项的数据中不包含父级id。则需要执行下面的方法
          // expandedKeys.push(getParentKey(item.treeId,this.state.treeList));
        }
        if(item.children && item.children.length > 0) {
          getExpandKeys(item.children);
        }
      })
    }
    // 此处想实现的效果是搜索为空时，不展开树。若想搜索为空时，全部树都展开，则去掉这层if包装即可
    if(value) { 
      getExpandKeys(this.state.treeList);
      expandedKeys = unique(expandedKeys);
    }
    this.setState({
      expandedKeys,
      searchValue: value,
      autoExpandParent: true,
    });
  }

  render() {
    const { treeLoading, treeList, searchValue, expandedKeys, autoExpandParent } = this.state;

    const loop = data => data.map((item) => {
      const index = item.treeName.indexOf(searchValue);
      const beforeStr = item.treeName.substr(0, index);
      const afterStr = item.treeName.substr(index + searchValue.length);
      const treeName = index > -1 ? (
        <span>
          {beforeStr}
          <span style={{ color: '#f50' }}>{searchValue}</span>
          {afterStr}
        </span>
      ) : <span>{item.treeName}</span>;
      if (item.children) {
        return (
          <TreeNode key={item.treeId} value={item.treeId} title={treeName} dataRef={item}>
            {loop(item.children)}
          </TreeNode>
        );
      }
      return <TreeNode key={item.treeId} title={treeName} value={item.treeId} dataRef={item} />;
    });

    return (
      <div>
          <Search
            placeholder="请输入"
            onSearch={this.onSearch}
            enterButton
          />
          
          <Spin spinning={treeLoading}>
            <Tree
              onExpand={this.onExpand}
              expandedKeys={expandedKeys}
              autoExpandParent={autoExpandParent}
            >
              {loop(this.state.treeList)}
            </Tree>
          </Spin>
      </div>
    )
  }
}
export default Catalog;