import * as React from 'react' ; import { Key , GetRowKey } from 'antd/es/table/interface' ; interface MapCache < RecordType > { data ? : RecordType [ ] ; childrenColumnName ? : string ; kvMap ? : Map < Key , RecordType > ; getRowKey ? : Function ; } export function findAllChildrenKeys < RecordType > ( data , getRowKey , childrenColumnName ) { const keys = [ ] ; function dig ( list ) { if ( ! Array . isArray ( list ) ) { return ; } ( list || [ ] ) . forEach ( ( item , index ) => { keys . push ( getRowKey ( item , index ) ) ; dig ( ( item as any ) [ childrenColumnName ] ) ; } ) ; } dig ( data ) ; return keys ; } export default function useLazyKVMap < RecordType > ( data , childrenColumnName , getRowKey ) { const mapCacheRef = React . useRef < MapCache < RecordType > > ( { } ) ; function getRecordByKey ( key ) { if ( ! mapCacheRef . current || mapCacheRef . current . data !== data || mapCacheRef . current . childrenColumnName !== childrenColumnName || mapCacheRef . current . getRowKey !== getRowKey ) { const kvMap = new Map < Key , RecordType > ( ) ; function dig ( records ) { records . forEach ( ( record , index ) => { const rowKey = getRowKey ( record , index ) ; kvMap . set ( rowKey , record ) ; } ) ; } dig ( data ) ; mapCacheRef . current = { data , childrenColumnName , kvMap , getRowKey , } ; } return mapCacheRef . current . kvMap ! . get ( key ) ! ; } return [ getRecordByKey ] ; }