From e50562feaa4ede9fc7cd4a6a15baef7f72954452 Mon Sep 17 00:00:00 2001
From: James Kerr <kerr@hey.com>
Date: Sun, 23 Jul 2023 19:44:00 -0700
Subject: [PATCH] Add a treeApi.setSelection method.

---
 .../src/components/provider.tsx               |  1 -
 .../react-arborist/src/interfaces/tree-api.ts | 24 ++++++++++++++-----
 .../src/state/selection-slice.ts              | 15 +++++++++---
 .../react-arborist/src/types/tree-props.ts    |  4 ++--
 packages/showcase/pages/gmail.tsx             | 16 +++++++++++--
 5 files changed, 46 insertions(+), 14 deletions(-)

diff --git a/packages/react-arborist/src/components/provider.tsx b/packages/react-arborist/src/components/provider.tsx
index a112f75..46c23a0 100644
--- a/packages/react-arborist/src/components/provider.tsx
+++ b/packages/react-arborist/src/components/provider.tsx
@@ -14,7 +14,6 @@ import {
   TreeApiContext,
 } from "../context";
 import { TreeApi } from "../interfaces/tree-api";
-import { IdObj } from "../types/utils";
 import { initialState } from "../state/initial";
 import { rootReducer, RootState } from "../state/root-reducer";
 import { HTML5Backend } from "react-dnd-html5-backend";
diff --git a/packages/react-arborist/src/interfaces/tree-api.ts b/packages/react-arborist/src/interfaces/tree-api.ts
index e9b6f29..b442a89 100644
--- a/packages/react-arborist/src/interfaces/tree-api.ts
+++ b/packages/react-arborist/src/interfaces/tree-api.ts
@@ -370,21 +370,33 @@ export class TreeApi<T> {
   }
 
   deselectAll() {
-    this.dispatch(selection.clear());
-    this.dispatch(selection.anchor(null));
-    this.dispatch(selection.mostRecent(null));
+    this.setSelection({ ids: [], anchor: null, mostRecent: null });
     safeRun(this.props.onSelect, this.selectedNodes);
   }
 
   selectAll() {
-    this.dispatch(selection.set(new Set(Object.keys(this.idToIndex))));
+    this.setSelection({
+      ids: Object.keys(this.idToIndex),
+      anchor: this.firstNode,
+      mostRecent: this.lastNode,
+    });
     this.dispatch(focus(this.lastNode?.id));
-    this.dispatch(selection.anchor(this.firstNode));
-    this.dispatch(selection.mostRecent(this.lastNode));
     if (this.focusedNode) safeRun(this.props.onFocus, this.focusedNode);
     safeRun(this.props.onSelect, this.selectedNodes);
   }
 
+  setSelection(args: {
+    ids: (IdObj | string)[] | null;
+    anchor: IdObj | string | null;
+    mostRecent: IdObj | string | null;
+  }) {
+    const ids = new Set(args.ids?.map(identify));
+    const anchor = identifyNull(args.anchor);
+    const mostRecent = identifyNull(args.mostRecent);
+    this.dispatch(selection.set({ ids, anchor, mostRecent }));
+    safeRun(this.props.onSelect, this.selectedNodes);
+  }
+
   /* Drag and Drop */
 
   get cursorParentId() {
diff --git a/packages/react-arborist/src/state/selection-slice.ts b/packages/react-arborist/src/state/selection-slice.ts
index c0c4947..62e2493 100644
--- a/packages/react-arborist/src/state/selection-slice.ts
+++ b/packages/react-arborist/src/state/selection-slice.ts
@@ -28,9 +28,13 @@ export const actions = {
     ids: (Array.isArray(id) ? id : [id]).map(identify),
   }),
 
-  set: (ids: Set<string>) => ({
+  set: (args: {
+    ids: Set<string>;
+    anchor: string | null;
+    mostRecent: string | null;
+  }) => ({
     type: "SELECTION_SET" as const,
-    ids,
+    ...args,
   }),
 
   mostRecent: (id: string | null | IdObj) => ({
@@ -64,7 +68,12 @@ export function reducer(
       action.ids.forEach((id) => ids.delete(id));
       return { ...state, ids: new Set(ids) };
     case "SELECTION_SET":
-      return { ...state, ids: new Set(action.ids) };
+      return {
+        ...state,
+        ids: action.ids,
+        mostRecent: action.mostRecent,
+        anchor: action.anchor,
+      };
     case "SELECTION_MOST_RECENT":
       return { ...state, mostRecent: action.id };
     case "SELECTION_ANCHOR":
diff --git a/packages/react-arborist/src/types/tree-props.ts b/packages/react-arborist/src/types/tree-props.ts
index 2bae985..de8fff4 100644
--- a/packages/react-arborist/src/types/tree-props.ts
+++ b/packages/react-arborist/src/types/tree-props.ts
@@ -5,7 +5,7 @@ import { ElementType, MouseEventHandler } from "react";
 import { ListOnScrollProps } from "react-window";
 import { NodeApi } from "../interfaces/node-api";
 import { OpenMap } from "../state/open-slice";
-import { useDragDropManager } from "react-dnd"
+import { useDragDropManager } from "react-dnd";
 
 export interface TreeProps<T> {
   /* Data Options */
@@ -76,5 +76,5 @@ export interface TreeProps<T> {
   dndRootElement?: globalThis.Node | null;
   onClick?: MouseEventHandler;
   onContextMenu?: MouseEventHandler;
-  dndManager?: ReturnType<typeof useDragDropManager>
+  dndManager?: ReturnType<typeof useDragDropManager>;
 }
diff --git a/packages/showcase/pages/gmail.tsx b/packages/showcase/pages/gmail.tsx
index 866262c..f9ef5f7 100644
--- a/packages/showcase/pages/gmail.tsx
+++ b/packages/showcase/pages/gmail.tsx
@@ -1,16 +1,27 @@
 import clsx from "clsx";
-import { CursorProps, NodeApi, NodeRendererProps, Tree } from "react-arborist";
+import {
+  CursorProps,
+  NodeApi,
+  NodeRendererProps,
+  Tree,
+  TreeApi,
+} from "react-arborist";
 import { gmailData, GmailItem } from "../data/gmail";
 import * as icons from "react-icons/md";
 import styles from "../styles/Gmail.module.css";
 import { FillFlexParent } from "../components/fill-flex-parent";
 import { SiGmail } from "react-icons/si";
 import { BsTree } from "react-icons/bs";
-import { useState } from "react";
+import { useEffect, useRef, useState } from "react";
 import Link from "next/link";
 
 export default function GmailSidebar() {
   const [term, setTerm] = useState("");
+  const globalTree = (tree: TreeApi<any> | null) => {
+    // @ts-ignore
+    window.tree = tree;
+  };
+
   return (
     <div className={styles.page}>
       <div className={styles.mainContent}>
@@ -28,6 +39,7 @@ export default function GmailSidebar() {
             {({ width, height }) => {
               return (
                 <Tree
+                  ref={globalTree}
                   initialData={gmailData}
                   width={width}
                   height={height}