From 8a51691a4de418b7bde82f67ceb6f1f050529d17 Mon Sep 17 00:00:00 2001
From: James Kerr <kerr@hey.com>
Date: Thu, 28 Mar 2024 15:55:27 -0700
Subject: [PATCH] Add Multiple Sort Fields

---
 .../src/nodes/default-accessors.ts            |  4 +-
 .../src/nodes/source-data-accessor.ts         | 42 +++++++++++++++----
 modules/react-arborist/src/utils.ts           |  5 +++
 3 files changed, 40 insertions(+), 11 deletions(-)

diff --git a/modules/react-arborist/src/nodes/default-accessors.ts b/modules/react-arborist/src/nodes/default-accessors.ts
index 239a610..1290958 100644
--- a/modules/react-arborist/src/nodes/default-accessors.ts
+++ b/modules/react-arborist/src/nodes/default-accessors.ts
@@ -23,7 +23,7 @@ export function createDefaultAccessors<T>(): SourceDataAccessors<T> {
         return true;
       }
     },
-    sortBy: (_d: T) => 0,
-    sortOrder: "asc",
+    sortBy: [],
+    sortOrder: [],
   };
 }
diff --git a/modules/react-arborist/src/nodes/source-data-accessor.ts b/modules/react-arborist/src/nodes/source-data-accessor.ts
index 58f1300..9620847 100644
--- a/modules/react-arborist/src/nodes/source-data-accessor.ts
+++ b/modules/react-arborist/src/nodes/source-data-accessor.ts
@@ -1,11 +1,15 @@
+import { toArray } from "../utils";
 import { createDefaultAccessors } from "./default-accessors";
 
+type GetSortField<T> = (d: T) => number | string | boolean;
+type SortOrder = "asc" | "desc";
+
 export type SourceDataAccessors<T> = {
   id: (d: T) => string;
   children: (d: T) => T[] | null;
   isLeaf: (d: T) => boolean;
-  sortBy: (d: T) => number | string | boolean;
-  sortOrder: "asc" | "desc";
+  sortBy: GetSortField<T> | GetSortField<T>[];
+  sortOrder: SortOrder | SortOrder[];
 };
 
 export class SourceDataAccessor<T> {
@@ -28,19 +32,39 @@ export class SourceDataAccessor<T> {
   }
 
   sort(array: T[]) {
-    return array.sort((a, b) => {
-      const first = this.access.sortBy(a);
-      const second = this.access.sortBy(b);
+    const orders = toArray(this.access.sortOrder);
+    const compares = toArray(this.access.sortBy);
 
-      if (this.asc) {
-        return first < second ? -1 : 1;
-      } else {
-        return first > second ? -1 : 1;
+    return array.sort((a, b) => {
+      for (let i = 0; i < compares.length; i++) {
+        const comparator = this.createComparator(
+          compares[i],
+          orders[i] || "asc",
+        );
+        const result = comparator(a, b);
+        if (result !== 0) return result;
       }
+      return 0;
     });
   }
 
   get asc() {
     return this.access.sortOrder === "asc";
   }
+
+  createComparator(getField: GetSortField<T>, sortOrder: SortOrder) {
+    return (a: T, b: T) => {
+      const first = getField(a);
+      const second = getField(b);
+
+      if (sortOrder === "asc") {
+        if (first < second) return -1;
+        if (first > second) return 1;
+      } else {
+        if (first < second) return 1;
+        if (first > second) return -1;
+      }
+      return 0;
+    };
+  }
 }
diff --git a/modules/react-arborist/src/utils.ts b/modules/react-arborist/src/utils.ts
index af07730..dee8c07 100644
--- a/modules/react-arborist/src/utils.ts
+++ b/modules/react-arborist/src/utils.ts
@@ -173,3 +173,8 @@ export function getInsertParentId(tree: TreeApi<any>) {
   if (focus.parent && !focus.parent.isRoot) return focus.parent.id;
   return null;
 }
+
+export function toArray<T>(itemOrArray: T | T[]): T[] {
+  if (Array.isArray(itemOrArray)) return itemOrArray;
+  else return [itemOrArray];
+}