feat: supplement maxCount logic for complicated cases

aojunhao123 · aojunhao123 · commit c8b2c9242b3c · 2024-12-06T16:36:07.000+08:00
diff --git a/examples/mutiple-with-maxCount.tsx b/examples/mutiple-with-maxCount.tsx
@@ -77,7 +77,6 @@ export default () => {
         maxCount={3}
         treeData={treeData}
       />
-
       <h2>checkable with maxCount</h2>
       <TreeSelect
         style={{ width: 300 }}
@@ -90,7 +89,6 @@ export default () => {
         onChange={onChange}
         value={value}
       />
-
       <h2>checkable with maxCount and treeCheckStrictly</h2>
       <TreeSelect
         style={{ width: 300 }}
diff --git a/src/OptionList.tsx b/src/OptionList.tsx
@@ -9,7 +9,7 @@ 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';
@@ -49,7 +49,6 @@ const OptionList: React.ForwardRefRenderFunction<ReviseRefOptionListProps> = (_,
     treeExpandAction,
     treeTitleRender,
     onPopupScroll,
-    displayValues,
     isOverMaxCount,
     maxCount,
     showCheckedStrategy,
@@ -84,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) {
@@ -167,58 +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, fieldNames: FieldNames): Key[] => {
+    const keys = [targetNode[fieldNames.value]];
+    if (!Array.isArray(targetNode.children)) {
+      return keys;
+    }
+
+    return targetNode.children.reduce((acc, child) => {
+      if (!child.disabled) {
+        acc.push(...getSelectableKeys(child, fieldNames));
+      }
+      return acc;
+    }, keys);
+  };
+
   const nodeDisabled = useEvent((node: DataNode) => {
-    // Always enable selected nodes
-    if (checkedKeys.includes(node[fieldNames.value])) {
+    const nodeValue = node[fieldNames.value];
+
+    if (checkedKeys.includes(nodeValue)) {
       return false;
     }
 
-    // Get all selectable keys under current node considering conduction rules
-    const getSelectableKeys = (nodes: DataNode[]) => {
-      const keys: Key[] = [];
-      const traverse = (n: DataNode) => {
-        if (!n.disabled) {
-          keys.push(n[fieldNames.value]);
-          // Only traverse children if node is not disabled
-          if (Array.isArray(n.children)) {
-            n.children.forEach(traverse);
-          }
-        }
-      };
-      nodes.forEach(traverse);
-      return keys;
-    };
+    if (isOverMaxCount) {
+      return true;
+    }
 
-    const selectableNodeValues = getSelectableKeys([node]);
+    const cacheKey = `${nodeValue}-${checkedKeys.join(',')}-${maxCount}`;
 
-    // Simulate checked state after selecting current node
-    const simulatedCheckedKeys = [...checkedKeys, ...selectableNodeValues];
+    // 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);
-
-    // Calculate display keys based on strategy
-    const simulatedDisplayKeys = formatStrategyValues(
+    const simulatedDisplayValues = formatStrategyValues(
       conductedKeys as SafeKey[],
       showCheckedStrategy,
       keyEntities,
       fieldNames,
     );
 
-    const currentDisplayKeys = formatStrategyValues(
-      checkedKeys as SafeKey[],
-      showCheckedStrategy,
-      keyEntities,
-      fieldNames,
-    );
+    const isDisabled = simulatedDisplayValues.length > maxCount;
 
-    const newDisplayValuesCount = simulatedDisplayKeys.length - currentDisplayKeys.length;
-
-    // Check if selecting this node would exceed maxCount
-    if (isOverMaxCount || memoRawValues.length + newDisplayValuesCount > maxCount) {
-      return true;
-    }
+    // update cache
+    disabledCacheRef.current.set(cacheKey, isDisabled);
 
-    return false;
+    return isDisabled;
   });
 
   // ========================== Get First Selectable Node ==========================
diff --git a/src/TreeSelect.tsx b/src/TreeSelect.tsx
@@ -422,6 +422,8 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
         mergedFieldNames,
       );
 
+      console.log('triggerChange');
+
       const labeledValues = convert2LabelValues(newRawValues);
       setInternalValue(labeledValues);
 
@@ -616,7 +618,6 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
       treeExpandAction,
       treeTitleRender,
       onPopupScroll,
-      displayValues: cachedDisplayValues,
       isOverMaxCount,
       maxCount,
       showCheckedStrategy: mergedShowCheckedStrategy,
@@ -634,7 +635,6 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
     treeTitleRender,
     onPopupScroll,
     maxCount,
-    cachedDisplayValues,
     mergedMultiple,
     mergedShowCheckedStrategy,
   ]);
diff --git a/src/TreeSelectContext.ts b/src/TreeSelectContext.ts
@@ -15,7 +15,6 @@ export interface TreeSelectContextProps {
   treeExpandAction?: ExpandAction;
   treeTitleRender?: (node: any) => React.ReactNode;
   onPopupScroll?: React.UIEventHandler<HTMLDivElement>;
-  displayValues?: LabeledValueType[];
   isOverMaxCount?: boolean;
   maxCount?: number;
   showCheckedStrategy?: CheckedStrategy;
diff --git a/tests/Select.maxCount.spec.tsx b/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');
+  });
+});
