diff --git a/src/TreeSelect.tsx b/src/TreeSelect.tsx
index 96324dd9..ae2fef13 100644
--- a/src/TreeSelect.tsx
+++ b/src/TreeSelect.tsx
@@ -236,7 +236,7 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
 
   const mergedId = useId(id);
   const treeConduction = treeCheckable && !treeCheckStrictly;
-  const mergedCheckable: boolean = !!(treeCheckable || treeCheckStrictly);
+  const mergedCheckable = treeCheckable || treeCheckStrictly;
   const mergedLabelInValue = treeCheckStrictly || labelInValue;
   const mergedMultiple = mergedCheckable || multiple;
 
diff --git a/src/hooks/useDataEntities.ts b/src/hooks/useDataEntities.ts
index 6e77011b..d015033c 100644
--- a/src/hooks/useDataEntities.ts
+++ b/src/hooks/useDataEntities.ts
@@ -6,7 +6,10 @@ import warning from 'rc-util/lib/warning';
 import { isNil } from '../utils/valueUtil';
 
 export default (treeData: any, fieldNames: FieldNames) =>
-  React.useMemo<{ keyEntities: Record<string, DataEntity> }>(() => {
+  React.useMemo<{
+    valueEntities: Map<RawValueType, DataEntity>;
+    keyEntities: Record<string, DataEntity>;
+  }>(() => {
     const collection = convertDataToEntities(treeData, {
       fieldNames,
       initWrapper: wrapper => ({
@@ -31,7 +34,5 @@ export default (treeData: any, fieldNames: FieldNames) =>
       },
     });
 
-    return collection;
-  }, [treeData, fieldNames]) as ReturnType<typeof convertDataToEntities> & {
-    valueEntities: Map<RawValueType, DataEntity>;
-  };
+    return collection as any;
+  }, [treeData, fieldNames]);
diff --git a/src/utils/legacyUtil.tsx b/src/utils/legacyUtil.tsx
index 328fd4f4..c300f5d8 100644
--- a/src/utils/legacyUtil.tsx
+++ b/src/utils/legacyUtil.tsx
@@ -1,13 +1,7 @@
 import * as React from 'react';
 import toArray from 'rc-util/lib/Children/toArray';
 import warning from 'rc-util/lib/warning';
-import type {
-  DataNode,
-  LegacyDataNode,
-  ChangeEventExtra,
-  RawValueType,
-  LegacyCheckedNode,
-} from '../interface';
+import type { DataNode, ChangeEventExtra, RawValueType, LegacyCheckedNode } from '../interface';
 import TreeNode from '../TreeNode';
 import type { DefaultOptionType, FieldNames } from '../TreeSelect';
 
@@ -39,9 +33,9 @@ export function convertChildrenToData(nodes: React.ReactNode): DataNode[] {
     .filter(data => data);
 }
 
-export function fillLegacyProps(dataNode: DataNode): LegacyDataNode {
+export function fillLegacyProps(dataNode: DataNode): any {
   if (!dataNode) {
-    return dataNode as LegacyDataNode;
+    return dataNode;
   }
 
   const cloneNode = { ...dataNode };
@@ -58,7 +52,7 @@ export function fillLegacyProps(dataNode: DataNode): LegacyDataNode {
     });
   }
 
-  return cloneNode as LegacyDataNode;
+  return cloneNode;
 }
 
 export function fillAdditionalInfo(
diff --git a/tests/Select.checkable.spec.js b/tests/Select.checkable.spec.js
index d5d55e90..eaaf859a 100644
--- a/tests/Select.checkable.spec.js
+++ b/tests/Select.checkable.spec.js
@@ -753,4 +753,21 @@ describe('TreeSelect.checkable', () => {
       expect.objectContaining({ checked: ['parent', 'child'] }),
     );
   });
+
+  it('customize checkable node', () => {
+    const wrapper = mount(
+      <TreeSelect
+        open
+        treeCheckable={<span className="little" />}
+        treeData={[
+          {
+            label: 'parent',
+            value: 'parent',
+          },
+        ]}
+      />,
+    );
+
+    expect(wrapper.exists('.little')).toBeTruthy();
+  });
 });