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,17 @@ 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,10 @@ 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';
 
 const HIDDEN_STYLE = {
   width: 0,
@@ -47,8 +48,9 @@ const OptionList: React.ForwardRefRenderFunction<ReviseRefOptionListProps> = (_,
     treeExpandAction,
     treeTitleRender,
     onPopupScroll,
-    displayValues,
-    isOverMaxCount,
+    leftMaxCount,
+    leafCountOnly,
+    valueEntities,
   } = React.useContext(TreeSelectContext);
 
   const {
@@ -80,11 +82,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 +160,60 @@ const OptionList: React.ForwardRefRenderFunction<ReviseRefOptionListProps> = (_,
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [searchValue]);
 
+  // ========================= Disabled =========================
+  const disabledCacheRef = React.useRef<Map<string, boolean>>(new Map());
+
+  // Clear cache if `leftMaxCount` changed
+  React.useEffect(() => {
+    if (leftMaxCount) {
+      disabledCacheRef.current.clear();
+    }
+  }, [leftMaxCount]);
+
+  function getDisabledWithCache(node: DataNode) {
+    const value = node[fieldNames.value];
+    if (!disabledCacheRef.current.has(value)) {
+      const entity = valueEntities.get(value);
+      const isLeaf = (entity.children || []).length === 0;
+
+      if (!isLeaf) {
+        const checkableChildren = entity.children.filter(
+          childTreeNode =>
+            !childTreeNode.node.disabled &&
+            !childTreeNode.node.disableCheckbox &&
+            !checkedKeys.includes(childTreeNode.node[fieldNames.value]),
+        );
+
+        const checkableChildrenCount = checkableChildren.length;
+        disabledCacheRef.current.set(value, checkableChildrenCount > leftMaxCount);
+      } else {
+        disabledCacheRef.current.set(value, false);
+      }
+    }
+    return disabledCacheRef.current.get(value);
+  }
+
   const nodeDisabled = useEvent((node: DataNode) => {
-    return isOverMaxCount && !memoRawValues.includes(node[fieldNames.value]);
+    const nodeValue = node[fieldNames.value];
+
+    if (checkedKeys.includes(nodeValue)) {
+      return false;
+    }
+
+    if (leftMaxCount === null) {
+      return false;
+    }
+
+    if (leftMaxCount <= 0) {
+      return true;
+    }
+
+    // This is a low performance calculation
+    if (leafCountOnly && leftMaxCount) {
+      return getDisabledWithCache(node);
+    }
+
+    return false;
   });
 
   // ========================== Get First Selectable Node ==========================

src/TreeSelect.tsx

@@ -408,6 +408,17 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
 
   const [cachedDisplayValues] = useCache(displayValues);
 
+  // ========================== MaxCount ==========================
+  const mergedMaxCount = React.useMemo(() => {
+    if (
+      mergedMultiple &&
+      (mergedShowCheckedStrategy === 'SHOW_CHILD' || treeCheckStrictly || !treeCheckable)
+    ) {
+      return maxCount;
+    }
+    return null;
+  }, [maxCount, mergedMultiple, treeCheckStrictly, mergedShowCheckedStrategy, treeCheckable]);
+
   // =========================== Change ===========================
   const triggerChange = useRefFunc(
     (
@@ -422,11 +433,9 @@ 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;
-        }
+      // Not allow pass with `maxCount`
+      if (mergedMaxCount && formattedKeyList.length > mergedMaxCount) {
+        return;
       }
 
       const labeledValues = convert2LabelValues(newRawValues);
@@ -607,9 +616,6 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
   });
 
   // ========================== Context ===========================
-  const isOverMaxCount =
-    mergedMultiple && maxCount !== undefined && cachedDisplayValues?.length >= maxCount;
-
   const treeSelectContext = React.useMemo<TreeSelectContextProps>(() => {
     return {
       virtual,
@@ -623,8 +629,10 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
       treeExpandAction,
       treeTitleRender,
       onPopupScroll,
-      displayValues: cachedDisplayValues,
-      isOverMaxCount,
+      leftMaxCount: maxCount ? maxCount - cachedDisplayValues.length : null,
+      leafCountOnly:
+        mergedShowCheckedStrategy === 'SHOW_CHILD' && !treeCheckStrictly && !!treeCheckable,
+      valueEntities,
     };
   }, [
     virtual,
@@ -639,8 +647,11 @@ const TreeSelect = React.forwardRef<BaseSelectRef, TreeSelectProps>((props, ref)
     treeTitleRender,
     onPopupScroll,
     maxCount,
-    cachedDisplayValues,
-    mergedMultiple,
+    cachedDisplayValues.length,
+    mergedShowCheckedStrategy,
+    treeCheckStrictly,
+    treeCheckable,
+    valueEntities,
   ]);
 
   // ======================= 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 type useDataEntities from './hooks/useDataEntities';
 
 export interface TreeSelectContextProps {
   virtual?: boolean;
@@ -14,8 +15,12 @@ export interface TreeSelectContextProps {
   treeExpandAction?: ExpandAction;
   treeTitleRender?: (node: any) => React.ReactNode;
   onPopupScroll?: React.UIEventHandler<HTMLDivElement>;
-  displayValues?: LabeledValueType[];
-  isOverMaxCount?: boolean;
+
+  // For `maxCount` usage
+  leftMaxCount: number | null;
+  /** When `true`, only take leaf node as count, or take all as count with `maxCount` limitation */
+  leafCountOnly: boolean;
+  valueEntities: ReturnType<typeof useDataEntities>['valueEntities'];
 }
 
 const TreeSelectContext = React.createContext<TreeSelectContextProps>(null as any);

src/utils/warningPropsUtil.ts

@@ -10,6 +10,8 @@ function warningProps(props: TreeSelectProps & { searchPlaceholder?: string }) {
     labelInValue,
     value,
     multiple,
+    showCheckedStrategy,
+    maxCount,
   } = props;
 
   warning(!searchPlaceholder, '`searchPlaceholder` has been removed.');
@@ -20,7 +22,7 @@ function warningProps(props: TreeSelectProps & { searchPlaceholder?: string }) {
 
   if (labelInValue || treeCheckStrictly) {
     warning(
-      toArray(value).every((val) => val && typeof val === 'object' && 'value' in val),
+      toArray(value).every(val => val && typeof val === 'object' && 'value' in val),
       'Invalid prop `value` supplied to `TreeSelect`. You should use { label: string, value: string | number } or [{ label: string, value: string | number }] instead.',
     );
   }
@@ -33,6 +35,13 @@ function warningProps(props: TreeSelectProps & { searchPlaceholder?: string }) {
   } else {
     warning(!Array.isArray(value), '`value` should not be array when `TreeSelect` is single mode.');
   }
+
+  if (maxCount && (showCheckedStrategy === 'SHOW_ALL' || showCheckedStrategy === 'SHOW_PARENT')) {
+    warning(
+      false,
+      '`maxCount` not work with `showCheckedStrategy=SHOW_ALL` or `showCheckedStrategy=SHOW_PARENT`.',
+    );
+  }
 }
 
 export default warningProps;