feat: supplement maxCount logic for complicated cases

examples/mutiple-with-maxCount.tsx

@@ -20,6 +20,20 @@ export default () => {
           key: '1-2',
           value: '1-2',
           title: '1-2',
+          disabled: true,
+          children: [
+            {
+              key: '1-2-1',
+              value: '1-2-1',
+              title: '1-2-1',
+              disabled: true,
+            },
+            {
+              key: '1-2-2',
+              value: '1-2-2',
+              title: '1-2-2',
+            },
+          ],
         },
         {
           key: '1-3',
@@ -63,21 +77,18 @@ export default () => {
         maxCount={3}
         treeData={treeData}
       />
-
       <h2>checkable with maxCount</h2>
       <TreeSelect
         style={{ width: 300 }}
         multiple
         treeCheckable
         // showCheckedStrategy="SHOW_ALL"
-        showCheckedStrategy="SHOW_PARENT"
-        // showCheckedStrategy="SHOW_CHILD"
+        // showCheckedStrategy="SHOW_PARENT"
         maxCount={4}
         treeData={treeData}
         onChange={onChange}
         value={value}
       />
-
       <h2>checkable with maxCount and treeCheckStrictly</h2>
       <TreeSelect
         style={{ width: 300 }}

src/OptionList.tsx

@@ -9,9 +9,11 @@ import useMemo from 'rc-util/lib/hooks/useMemo';
 import * as React from 'react';
 import LegacyContext from './LegacyContext';
 import TreeSelectContext from './TreeSelectContext';
-import type { DataNode, Key, SafeKey } from './interface';
+import type { DataNode, FieldNames, Key, SafeKey } from './interface';
 import { getAllKeys, isCheckDisabled } from './utils/valueUtil';
 import { useEvent } from 'rc-util';
+import { formatStrategyValues } from './utils/strategyUtil';
+import { conductCheck } from 'rc-tree/lib/utils/conductUtil';
 
 const HIDDEN_STYLE = {
   width: 0,
@@ -47,8 +49,9 @@ const OptionList: React.ForwardRefRenderFunction<ReviseRefOptionListProps> = (_,
     treeExpandAction,
     treeTitleRender,
     onPopupScroll,
-    displayValues,
     isOverMaxCount,
+    maxCount,
+    showCheckedStrategy,
   } = React.useContext(TreeSelectContext);
 
   const {
@@ -80,11 +83,6 @@ const OptionList: React.ForwardRefRenderFunction<ReviseRefOptionListProps> = (_,
     (prev, next) => next[0] && prev[1] !== next[1],
   );
 
-  const memoRawValues = React.useMemo(
-    () => (displayValues || []).map(v => v.value),
-    [displayValues],
-  );
-
   // ========================== Values ==========================
   const mergedCheckedKeys = React.useMemo(() => {
     if (!checkable) {
@@ -163,8 +161,69 @@ const OptionList: React.ForwardRefRenderFunction<ReviseRefOptionListProps> = (_,
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [searchValue]);
 
+  const disabledCacheRef = React.useRef(new Map<Key, boolean>());
+  const lastCheckedKeysRef = React.useRef<Key[]>([]);
+  const lastMaxCountRef = React.useRef<number>(null);
+
+  const resetCache = React.useCallback(() => {
+    disabledCacheRef.current.clear();
+    lastCheckedKeysRef.current = [...checkedKeys];
+    lastMaxCountRef.current = maxCount;
+  }, [checkedKeys, maxCount]);
+
+  React.useEffect(() => {
+    resetCache();
+  }, [checkedKeys, maxCount]);
+
+  const getSelectableKeys = (targetNode: DataNode, names: FieldNames): Key[] => {
+    const keys = [targetNode[names.value]];
+    if (!Array.isArray(targetNode.children)) {
+      return keys;
+    }
+
+    return targetNode.children.reduce((acc, child) => {
+      if (!child.disabled) {
+        acc.push(...getSelectableKeys(child, names));
+      }
+      return acc;
+    }, keys);
+  };
+
   const nodeDisabled = useEvent((node: DataNode) => {
-    return isOverMaxCount && !memoRawValues.includes(node[fieldNames.value]);
+    const nodeValue = node[fieldNames.value];
+
+    if (checkedKeys.includes(nodeValue)) {
+      return false;
+    }
+
+    if (isOverMaxCount) {
+      return true;
+    }
+
+    const cacheKey = `${nodeValue}-${checkedKeys.join(',')}-${maxCount}`;
+
+    // check cache
+    if (disabledCacheRef.current.has(cacheKey)) {
+      return disabledCacheRef.current.get(cacheKey);
+    }
+
+    // calculate disabled state
+    const selectableNodeKeys = getSelectableKeys(node, fieldNames);
+    const simulatedCheckedKeys = [...checkedKeys, ...selectableNodeKeys];
+    const { checkedKeys: conductedKeys } = conductCheck(simulatedCheckedKeys, true, keyEntities);
+    const simulatedDisplayValues = formatStrategyValues(
+      conductedKeys as SafeKey[],
+      showCheckedStrategy,
+      keyEntities,
+      fieldNames,
+    );
+
+    const isDisabled = simulatedDisplayValues.length > maxCount;
+
+    // update cache
+    disabledCacheRef.current.set(cacheKey, isDisabled);
+
+    return isDisabled;
   });
 
   // ========================== Get First Selectable Node ==========================

src/TreeSelect.tsx

@@ -422,13 +422,6 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
         mergedFieldNames,
       );
 
-      // if multiple and maxCount is set, check if exceed maxCount
-      if (mergedMultiple && maxCount !== undefined) {
-        if (formattedKeyList.length > maxCount) {
-          return;
-        }
-      }
-
       const labeledValues = convert2LabelValues(newRawValues);
       setInternalValue(labeledValues);
 
@@ -623,8 +616,9 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
       treeExpandAction,
       treeTitleRender,
       onPopupScroll,
-      displayValues: cachedDisplayValues,
       isOverMaxCount,
+      maxCount,
+      showCheckedStrategy: mergedShowCheckedStrategy,
     };
   }, [
     virtual,
@@ -639,8 +633,8 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
     treeTitleRender,
     onPopupScroll,
     maxCount,
-    cachedDisplayValues,
     mergedMultiple,
+    mergedShowCheckedStrategy,
   ]);
 
   // ======================= Legacy Context =======================

src/TreeSelectContext.ts

@@ -1,6 +1,7 @@
 import * as React from 'react';
 import type { ExpandAction } from 'rc-tree/lib/Tree';
-import type { DataNode, FieldNames, Key, LabeledValueType } from './interface';
+import type { DataNode, FieldNames, Key } from './interface';
+import { CheckedStrategy } from './utils/strategyUtil';
 
 export interface TreeSelectContextProps {
   virtual?: boolean;
@@ -14,8 +15,9 @@ export interface TreeSelectContextProps {
   treeExpandAction?: ExpandAction;
   treeTitleRender?: (node: any) => React.ReactNode;
   onPopupScroll?: React.UIEventHandler<HTMLDivElement>;
-  displayValues?: LabeledValueType[];
   isOverMaxCount?: boolean;
+  maxCount?: number;
+  showCheckedStrategy?: CheckedStrategy;
 }
 
 const TreeSelectContext = React.createContext<TreeSelectContextProps>(null as any);

tests/Select.maxCount.spec.tsx

@@ -372,3 +372,132 @@ describe('TreeSelect.maxCount with treeCheckStrictly', () => {
     expect(handleChange).toHaveBeenCalledTimes(4);
   });
 });
+
+describe('TreeSelect.maxCount with complex scenarios', () => {
+  const complexTreeData = [
+    {
+      key: 'asia',
+      value: 'asia',
+      title: 'Asia',
+      children: [
+        {
+          key: 'china',
+          value: 'china',
+          title: 'China',
+          children: [
+            { key: 'beijing', value: 'beijing', title: 'Beijing' },
+            { key: 'shanghai', value: 'shanghai', title: 'Shanghai' },
+            { key: 'guangzhou', value: 'guangzhou', title: 'Guangzhou' },
+          ],
+        },
+        {
+          key: 'japan',
+          value: 'japan',
+          title: 'Japan',
+          children: [
+            { key: 'tokyo', value: 'tokyo', title: 'Tokyo' },
+            { key: 'osaka', value: 'osaka', title: 'Osaka' },
+          ],
+        },
+      ],
+    },
+    {
+      key: 'europe',
+      value: 'europe',
+      title: 'Europe',
+      children: [
+        {
+          key: 'uk',
+          value: 'uk',
+          title: 'United Kingdom',
+          children: [
+            { key: 'london', value: 'london', title: 'London' },
+            { key: 'manchester', value: 'manchester', title: 'Manchester' },
+          ],
+        },
+        {
+          key: 'france',
+          value: 'france',
+          title: 'France',
+          disabled: true,
+          children: [
+            { key: 'paris', value: 'paris', title: 'Paris' },
+            { key: 'lyon', value: 'lyon', title: 'Lyon' },
+          ],
+        },
+      ],
+    },
+  ];
+
+  it('should handle complex tree structure with maxCount correctly', () => {
+    const handleChange = jest.fn();
+    const { getByRole } = render(
+      <TreeSelect
+        treeData={complexTreeData}
+        treeCheckable
+        treeDefaultExpandAll
+        multiple
+        maxCount={3}
+        onChange={handleChange}
+        open
+      />,
+    );
+
+    const container = getByRole('tree');
+
+    // 选择一个顶层节点
+    const asiaNode = within(container).getByText('Asia');
+    fireEvent.click(asiaNode);
+    expect(handleChange).not.toHaveBeenCalled(); // 不应该触发，因为会超过 maxCount
+
+    // 选择叶子节点
+    const beijingNode = within(container).getByText('Beijing');
+    const shanghaiNode = within(container).getByText('Shanghai');
+    const tokyoNode = within(container).getByText('Tokyo');
+    const londonNode = within(container).getByText('London');
+
+    fireEvent.click(beijingNode);
+    fireEvent.click(shanghaiNode);
+    fireEvent.click(tokyoNode);
+    expect(handleChange).toHaveBeenCalledTimes(3);
+
+    // 尝试选择第四个节点，应该被阻止
+    fireEvent.click(londonNode);
+    expect(handleChange).toHaveBeenCalledTimes(3);
+
+    // 验证禁用状态
+    expect(londonNode.closest('div')).toHaveClass('rc-tree-select-tree-treenode-disabled');
+  });
+
+  it('should handle maxCount with mixed selection strategies', () => {
+    const handleChange = jest.fn();
+
+    const { getByRole } = render(
+      <TreeSelect
+        treeData={complexTreeData}
+        treeCheckable
+        treeDefaultExpandAll
+        multiple
+        maxCount={3}
+        onChange={handleChange}
+        defaultValue={['uk']}
+        open
+      />,
+    );
+
+    const container = getByRole('tree');
+
+    const tokyoNode = within(container).getByText('Tokyo');
+    fireEvent.click(tokyoNode);
+
+    // because UK node will show two children, so it will trigger one change
+    expect(handleChange).toHaveBeenCalledTimes(1);
+
+    const beijingNode = within(container).getByText('Beijing');
+    fireEvent.click(beijingNode);
+
+    // should not trigger change
+    expect(handleChange).toHaveBeenCalledTimes(1);
+    expect(beijingNode.closest('div')).toHaveClass('rc-tree-select-tree-treenode-disabled');
+  });
+});