From 7ae5e336480f38daad9bfe0ccfca68a4ff7921bf Mon Sep 17 00:00:00 2001
From: Ivan Kopenkov <ivan@ennabl.com>
Date: Wed, 19 Oct 2022 12:00:08 +0300
Subject: [PATCH] fix half checked state when children are disabled

---
 docs/examples/basic.jsx  |  4 +-
 src/Tree.tsx             |  8 +++-
 src/utils/conductUtil.ts | 80 +++++++++++++++++++++++-----------------
 3 files changed, 54 insertions(+), 38 deletions(-)

diff --git a/docs/examples/basic.jsx b/docs/examples/basic.jsx
index ab986e32b..996160d51 100644
--- a/docs/examples/basic.jsx
+++ b/docs/examples/basic.jsx
@@ -33,7 +33,7 @@ const treeData = [
 
 class Demo extends React.Component {
   static defaultProps = {
-    keys: ['0-0-0-0'],
+    keys: ['0-0-0-0', '0-0-1-0'],
   };
 
   constructor(props) {
@@ -117,7 +117,7 @@ class Demo extends React.Component {
           <TreeNode title="parent 1" key="0-0">
             <TreeNode title={customLabel} key="0-0-0">
               <TreeNode title="leaf" key="0-0-0-0" style={{ background: 'rgba(255, 0, 0, 0.1)' }} />
-              <TreeNode title="leaf" key="0-0-0-1" />
+              <TreeNode title="leaf" key="0-0-0-1" disabled />
             </TreeNode>
             <TreeNode title="parent 1-1" key="0-0-1">
               <TreeNode title="parent 1-1-0" key="0-0-1-0" disableCheckbox />
diff --git a/src/Tree.tsx b/src/Tree.tsx
index a050fafec..941aad4d9 100644
--- a/src/Tree.tsx
+++ b/src/Tree.tsx
@@ -909,10 +909,14 @@ class Tree<TreeDataType extends DataNode | BasicDataNode = DataNode> extends Rea
         keyEntities,
       );
 
-      // If remove, we do it again to correction
-      if (!checked) {
+      const isHafCheckedAndNothingToCheckMore = treeNode.halfChecked && oriCheckedKeys.length === checkedKeys.length;
+
+      // If remove or halfChecked was triggered and nothing changed, 
+      // we do it again to correction
+      if (!checked || isHafCheckedAndNothingToCheckMore) {
         const keySet = new Set(checkedKeys);
         keySet.delete(key);
+        halfCheckedKeys = halfCheckedKeys.filter(item => item !== key);
         ({ checkedKeys, halfCheckedKeys } = conductCheck(
           Array.from(keySet),
           { checked: false, halfCheckedKeys },
diff --git a/src/utils/conductUtil.ts b/src/utils/conductUtil.ts
index c5d8d3f31..bbcac05e1 100644
--- a/src/utils/conductUtil.ts
+++ b/src/utils/conductUtil.ts
@@ -66,25 +66,29 @@ function fillConductCheck<TreeDataType extends BasicDataNode = DataNode>(
       }
 
       let allChecked = true;
-      let partialChecked = false;
-
-      (parent.children || [])
-        .filter(childEntity => !syntheticGetCheckDisabled(childEntity.node))
-        .forEach(({ key }) => {
-          const checked = checkedKeys.has(key);
-          if (allChecked && !checked) {
-            allChecked = false;
-          }
-          if (!partialChecked && (checked || halfCheckedKeys.has(key))) {
-            partialChecked = true;
-          }
-        });
+      let hasChecked = false;
+
+      (parent.children || []).forEach(({ key }) => {
+        const checked = checkedKeys.has(key);
+        if (checked) {
+          hasChecked = true;
+        } else if (halfCheckedKeys.has(key)) {
+          hasChecked = true;
+          allChecked = false;
+        } else {
+          allChecked = false;
+        }
+      });
 
       if (allChecked) {
         checkedKeys.add(parent.key);
-      }
-      if (partialChecked) {
+        halfCheckedKeys.delete(parent.key);
+      } else if (hasChecked) {
         halfCheckedKeys.add(parent.key);
+        checkedKeys.delete(parent.key);
+      } else {
+        halfCheckedKeys.delete(parent.key);
+        checkedKeys.delete(parent.key);
       }
 
       visitedKeys.add(parent.key);
@@ -93,7 +97,7 @@ function fillConductCheck<TreeDataType extends BasicDataNode = DataNode>(
 
   return {
     checkedKeys: Array.from(checkedKeys),
-    halfCheckedKeys: Array.from(removeFromCheckedKeys(halfCheckedKeys, checkedKeys)),
+    halfCheckedKeys: Array.from(halfCheckedKeys),
   };
 }
 
@@ -108,6 +112,9 @@ function cleanConductCheck<TreeDataType extends BasicDataNode = DataNode>(
   const checkedKeys = new Set<Key>(keys);
   let halfCheckedKeys = new Set<Key>(halfKeys);
 
+  console.log({keys,
+    halfKeys})
+
   // Remove checked keys from top to bottom
   for (let level = 0; level <= maxLevel; level += 1) {
     const entities = levelEntities.get(level) || new Set();
@@ -118,6 +125,7 @@ function cleanConductCheck<TreeDataType extends BasicDataNode = DataNode>(
         children
           .filter(childEntity => !syntheticGetCheckDisabled(childEntity.node))
           .forEach(childEntity => {
+            halfCheckedKeys.delete(childEntity.key);
             checkedKeys.delete(childEntity.key);
           });
       }
@@ -145,25 +153,29 @@ function cleanConductCheck<TreeDataType extends BasicDataNode = DataNode>(
       }
 
       let allChecked = true;
-      let partialChecked = false;
-
-      (parent.children || [])
-        .filter(childEntity => !syntheticGetCheckDisabled(childEntity.node))
-        .forEach(({ key }) => {
-          const checked = checkedKeys.has(key);
-          if (allChecked && !checked) {
-            allChecked = false;
-          }
-          if (!partialChecked && (checked || halfCheckedKeys.has(key))) {
-            partialChecked = true;
-          }
-        });
-
-      if (!allChecked) {
-        checkedKeys.delete(parent.key);
-      }
-      if (partialChecked) {
+      let hasChecked = false;
+
+      (parent.children || []).forEach(({ key }) => {
+        const checked = checkedKeys.has(key);
+        if (checked) {
+          hasChecked = true;
+        } else if (halfCheckedKeys.has(key)) {
+          hasChecked = true;
+          allChecked = false;
+        } else {
+          allChecked = false;
+        }
+      });
+
+      if (allChecked) {
+        checkedKeys.add(parent.key);
+        halfCheckedKeys.delete(parent.key);
+      } else if (hasChecked) {
         halfCheckedKeys.add(parent.key);
+        checkedKeys.delete(parent.key);
+      } else {
+        halfCheckedKeys.delete(parent.key);
+        checkedKeys.delete(parent.key);
       }
 
       visitedKeys.add(parent.key);
