[@mantine/core] Tree: Add checked logic

packages/@mantine/core/src/components/Tree/Tree.tsx

@@ -40,6 +40,9 @@ export interface RenderTreeNodePayload {
   /** Node data from the `data` prop of `Tree` */
   node: TreeNodeData;
 
+  /** Tree controller instance, return value of `useTree` hook */
+  tree: TreeController;
+
   /** Props to spread into the root node element */
   elementProps: {
     className: string;

packages/@mantine/core/src/components/Tree/TreeNode.tsx

@@ -174,6 +174,7 @@ export function TreeNode({
           node,
           level,
           selected,
+          tree: controller,
           expanded: controller.expandedState[node.value] || false,
           hasChildren: Array.isArray(node.children) && node.children.length > 0,
           elementProps,

packages/@mantine/core/src/components/Tree/get-all-checked-nodes/get-all-checked-nodes.test.ts

@@ -0,0 +1,68 @@
+import { getAllCheckedNodes } from './get-all-checked-nodes';
+
+describe('@mantine/core/Tree/get-all-checked-nodes', () => {
+  it('returns all checked nodes for a flat tree', () => {
+    const data = [
+      { value: '1', label: '1' },
+      { value: '2', label: '2' },
+      { value: '3', label: '3' },
+    ];
+
+    const { result } = getAllCheckedNodes(data, ['1', '2']);
+
+    expect(result).toStrictEqual([
+      { checked: true, indeterminate: false, value: '1', hasChildren: false },
+      { checked: true, indeterminate: false, value: '2', hasChildren: false },
+    ]);
+  });
+
+  it('returns all checked nodes for a nested tree with one level of nesting', () => {
+    const data = [
+      {
+        value: '1',
+        label: '1',
+        children: [
+          { value: '2', label: '2' },
+          { value: '3', label: '3' },
+        ],
+      },
+    ];
+
+    const { result } = getAllCheckedNodes(data, ['2']);
+
+    expect(result).toStrictEqual([
+      { checked: true, indeterminate: false, value: '2', hasChildren: false },
+      { checked: false, indeterminate: true, value: '1', hasChildren: true },
+    ]);
+  });
+
+  it('returns all checked nodes for a nested tree with multiple levels of nesting', () => {
+    const data = [
+      {
+        value: '1',
+        label: '1',
+        children: [
+          {
+            value: '2',
+            label: '2',
+            children: [
+              { value: '3', label: '3' },
+              { value: '4', label: '4' },
+              { value: '5', label: '5' },
+            ],
+          },
+          { value: '6', label: '6' },
+        ],
+      },
+    ];
+
+    const { result } = getAllCheckedNodes(data, ['3', '4']);
+
+    expect(result).toStrictEqual([
+      { checked: true, indeterminate: false, value: '3', hasChildren: false },
+      { checked: true, indeterminate: false, value: '4', hasChildren: false },
+      { checked: false, indeterminate: true, value: '2', hasChildren: true },
+      { checked: false, indeterminate: true, value: '1', hasChildren: true },
+    ]);
+  });
+});

packages/@mantine/core/src/components/Tree/get-all-checked-nodes/get-all-checked-nodes.ts

@@ -0,0 +1,42 @@
+import type { TreeNodeData } from '../Tree';
+
+export interface CheckedNodeStatus {
+  checked: boolean;
+  indeterminate: boolean;
+  hasChildren: boolean;
+  value: string;
+}
+
+export function getAllCheckedNodes(
+  data: TreeNodeData[],
+  checkedState: string[],
+  acc: CheckedNodeStatus[] = []
+) {
+  const currentTreeChecked: CheckedNodeStatus[] = [];
+
+  for (const node of data) {
+    if (Array.isArray(node.children) && node.children.length > 0) {
+      const innerChecked = getAllCheckedNodes(node.children, checkedState, acc);
+      if (innerChecked.currentTreeChecked.length === node.children.length) {
+        const item = { checked: true, indeterminate: false, value: node.value, hasChildren: true };
+        currentTreeChecked.push(item);
+        acc.push(item);
+      } else if (innerChecked.currentTreeChecked.length > 0) {
+        const item = { checked: false, indeterminate: true, value: node.value, hasChildren: true };
+        currentTreeChecked.push(item);
+        acc.push(item);
+      }
+    } else if (checkedState.includes(node.value)) {
+      const item: CheckedNodeStatus = {
+        checked: true,
+        indeterminate: false,
+        value: node.value,
+        hasChildren: false,
+      };
+      currentTreeChecked.push(item);
+      acc.push(item);
+    }
+  }
+
+  return { result: acc, currentTreeChecked };
+}

packages/@mantine/core/src/components/Tree/index.ts

@@ -9,3 +9,4 @@ export type {
   RenderTreeNodePayload,
 } from './Tree';
 export type { UseTreeInput, UseTreeReturnType } from './use-tree';
+export type { CheckedNodeStatus } from './get-all-checked-nodes/get-all-checked-nodes';

packages/@mantine/core/src/components/Tree/is-node-checked/is-node-checked.test.ts

@@ -0,0 +1,56 @@
+import { isNodeChecked } from './is-node-checked';
+
+describe('@mantine/core/Tree/is-node-checked', () => {
+  it('detects checked node in flat tree', () => {
+    const data = [
+      { value: '1', label: '1' },
+      { value: '2', label: '2' },
+      { value: '3', label: '3' },
+    ];
+
+    expect(isNodeChecked('1', data, ['1', '2'])).toBe(true);
+    expect(isNodeChecked('2', data, ['1', '2'])).toBe(true);
+    expect(isNodeChecked('3', data, ['1', '2'])).toBe(false);
+  });
+
+  it('detects checked node in nested tree with one level of nesting', () => {
+    const data = [
+      {
+        value: '1',
+        label: '1',
+        children: [
+          { value: '2', label: '2' },
+          { value: '3', label: '3' },
+        ],
+      },
+    ];
+
+    expect(isNodeChecked('1', data, ['2'])).toBe(false);
+    expect(isNodeChecked('2', data, ['2'])).toBe(true);
+    expect(isNodeChecked('1', data, ['2', '3'])).toBe(true);
+  });
+
+  it('detects checked node in nested tree with multiple levels of nesting', () => {
+    const data = [
+      {
+        value: '1',
+        label: '1',
+        children: [
+          {
+            value: '2',
+            label: '2',
+            children: [
+              { value: '3', label: '3' },
+              { value: '4', label: '4' },
+              { value: '5', label: '5' },
+            ],
+          },
+          { value: '6', label: '6' },
+        ],
+      },
+    ];
+
+    expect(isNodeChecked('2', data, ['3', '4'])).toBe(false);
+    expect(isNodeChecked('2', data, ['3', '4', '5'])).toBe(true);
+  });
+});

packages/@mantine/core/src/components/Tree/is-node-checked/is-node-checked.ts

@@ -0,0 +1,19 @@
+import { getAllCheckedNodes } from '../get-all-checked-nodes/get-all-checked-nodes';
+import type { TreeNodeData } from '../Tree';
+
+export function isNodeChecked(
+  value: string,
+  data: TreeNodeData[],
+  checkedState: string[]
+): boolean {
+  if (checkedState.length === 0) {
+    return false;
+  }
+
+  if (checkedState.includes(value)) {
+    return true;
+  }
+
+  const checkedNodes = getAllCheckedNodes(data, checkedState).result;
+  return checkedNodes.some((node) => node.value === value && node.checked);
+}

packages/@mantine/core/src/components/Tree/is-node-indeterminate/is-node-indeterminate.test.ts

@@ -0,0 +1,54 @@
+import { isNodeIndeterminate } from './is-node-indeterminate';
+
+describe('@mantine/core/Tree/is-node-indeterminate', () => {
+  it('detects indeterminate node in flat tree', () => {
+    const data = [
+      { value: '1', label: '1' },
+      { value: '2', label: '2' },
+      { value: '3', label: '3' },
+    ];
+
+    expect(isNodeIndeterminate('1', data, ['1', '2'])).toBe(false);
+  });
+
+  it('detects indeterminate node in nested tree with one level of nesting', () => {
+    const data = [
+      {
+        value: '1',
+        label: '1',
+        children: [
+          { value: '2', label: '2' },
+          { value: '3', label: '3' },
+        ],
+      },
+    ];
+
+    expect(isNodeIndeterminate('1', data, ['2'])).toBe(true);
+    expect(isNodeIndeterminate('2', data, ['2'])).toBe(false);
+    expect(isNodeIndeterminate('1', data, ['2', '3'])).toBe(false);
+  });
+
+  it('detects indeterminate node in nested tree with multiple levels of nesting', () => {
+    const data = [
+      {
+        value: '1',
+        label: '1',
+        children: [
+          {
+            value: '2',
+            label: '2',
+            children: [
+              { value: '3', label: '3' },
+              { value: '4', label: '4' },
+              { value: '5', label: '5' },
+            ],
+          },
+          { value: '6', label: '6' },
+        ],
+      },
+    ];
+
+    expect(isNodeIndeterminate('2', data, ['3', '4'])).toBe(true);
+    expect(isNodeIndeterminate('2', data, ['3', '4', '5'])).toBe(false);
+  });
+});

packages/@mantine/core/src/components/Tree/is-node-indeterminate/is-node-indeterminate.ts

@@ -0,0 +1,15 @@
+import { getAllCheckedNodes } from '../get-all-checked-nodes/get-all-checked-nodes';
+import type { TreeNodeData } from '../Tree';
+
+export function isNodeIndeterminate(
+  value: string,
+  data: TreeNodeData[],
+  checkedState: string[]
+): boolean {
+  if (checkedState.length === 0) {
+    return false;
+  }
+
+  const checkedNodes = getAllCheckedNodes(data, checkedState).result;
+  return checkedNodes.some((node) => node.value === value && node.indeterminate);
+}