From 66111b9afa8546103d4cbf0b597c77c889d0a7cd Mon Sep 17 00:00:00 2001
From: Arman Ozak <arman@armanozak.com>
Date: Tue, 16 Jun 2020 14:21:28 +0300
Subject: [PATCH] feat: improve tree util types

---
 .../packages/core/src/lib/utils/tree-utils.ts | 44 ++++++++++++-------
 1 file changed, 27 insertions(+), 17 deletions(-)

diff --git a/npm/ng-packs/packages/core/src/lib/utils/tree-utils.ts b/npm/ng-packs/packages/core/src/lib/utils/tree-utils.ts
index 86844a4375..607a30ec52 100644
--- a/npm/ng-packs/packages/core/src/lib/utils/tree-utils.ts
+++ b/npm/ng-packs/packages/core/src/lib/utils/tree-utils.ts
@@ -1,20 +1,24 @@
-export class TreeNode<T extends any> {
+export class TreeNodeFactory<T extends object> {
   children: TreeNode<T>[] = [];
   isLeaf = true;
 
   constructor(props: T) {
     Object.assign(this, props);
   }
+
+  static create<T extends object>(props: T) {
+    return new TreeNodeFactory<T>(props) as TreeNode<T>;
+  }
 }
 
-export function createTreeFromList<T extends object, R extends BranchOrLeaf<T>>(
+export function createTreeFromList<T extends object, R extends unknown>(
   list: T[],
-  keySelector: (item: T) => NodeKey,
-  parentKeySelector: (item: T) => NodeKey,
-  valueMapper = (item: T) => new TreeNode(item) as R,
+  keySelector: (item: T) => number | string | Symbol,
+  parentKeySelector: typeof keySelector,
+  valueMapper: (item: T) => R,
 ) {
   const map = createMapFromList(list, keySelector, valueMapper);
-  const tree: ReturnType<typeof valueMapper>[] = [];
+  const tree: NodeValue<T, typeof valueMapper>[] = [];
 
   list.forEach(row => {
     const id = keySelector(row);
@@ -23,8 +27,8 @@ export function createTreeFromList<T extends object, R extends BranchOrLeaf<T>>(
 
     if (parentId) {
       const parent = map.get(parentId);
-      parent.children.push(node);
-      parent.isLeaf = false;
+      (parent as any).children.push(node);
+      (parent as any).isLeaf = false;
     } else {
       tree.push(node);
     }
@@ -33,19 +37,25 @@ export function createTreeFromList<T extends object, R extends BranchOrLeaf<T>>(
   return tree;
 }
 
-export function createMapFromList<T extends object, R extends BranchOrLeaf<T>>(
+export function createMapFromList<T extends object, R extends unknown>(
   list: T[],
-  keySelector: (item: T) => NodeKey,
-  valueMapper = (item: T) => new TreeNode(item) as R,
+  keySelector: (item: T) => number | string | Symbol,
+  valueMapper: (item: T) => R,
 ) {
-  const map = new Map<ReturnType<typeof keySelector>, ReturnType<typeof valueMapper>>();
+  type Key = ReturnType<typeof keySelector>;
+  type Value = NodeValue<T, typeof valueMapper>;
+  const map = new Map<Key, Value>();
   list.forEach(row => map.set(keySelector(row), valueMapper(row)));
   return map;
 }
 
-type NodeKey = number | string | Symbol;
-
-interface BranchOrLeaf<T> {
-  children: BranchOrLeaf<T>[];
+export type TreeNode<T extends object> = {
+  [K in keyof T]: T[K];
+} & {
+  children: TreeNode<T>[];
   isLeaf: boolean;
-}
+};
+
+type NodeValue<T extends object, F extends (...args: any) => any> = F extends undefined
+  ? TreeNode<T>
+  : ReturnType<F>;