Mirror of https://github.com/roostorg/coop
github.com/roostorg/coop
1import partition from 'lodash/partition';
2
3export class TreeNode<T> {
4 public key: string;
5 public value: T;
6 public parent: TreeNode<T> | undefined;
7 public children: TreeNode<T>[];
8
9 constructor(key: string, value: T, parent?: TreeNode<T>) {
10 this.key = key;
11 this.value = value;
12 this.parent = parent;
13 this.children = [];
14 }
15
16 get isLeaf() {
17 return this.children.length === 0;
18 }
19
20 get hasChildren() {
21 return !this.isLeaf;
22 }
23}
24
25export class Tree<T extends any> {
26 public root: TreeNode<T>;
27
28 constructor(key: string, value: T) {
29 this.root = new TreeNode(key, value);
30 }
31
32 *preOrderTraversal(node = this.root): Generator<TreeNode<T>, void, unknown> {
33 yield node;
34 if (node.children.length) {
35 for (const child of node.children) {
36 yield* this.preOrderTraversal(child);
37 }
38 }
39 }
40
41 *postOrderTraversal(node = this.root): Generator<TreeNode<T>, void, unknown> {
42 if (node.children.length) {
43 for (const child of node.children) {
44 yield* this.postOrderTraversal(child);
45 }
46 }
47 yield node;
48 }
49
50 *parentTraversal(node: TreeNode<T>): Generator<TreeNode<T>, void, unknown> {
51 let current: TreeNode<T> | undefined = node;
52 while (current) {
53 yield current;
54 current = current.parent;
55 }
56 }
57
58 getPathToNode(node: TreeNode<T>) {
59 return Array.from(this.parentTraversal(node)).reverse();
60 }
61
62 insert(parentNodeKey: string, key: string, value: T) {
63 for (const node of this.preOrderTraversal()) {
64 if (node.key === parentNodeKey) {
65 node.children.push(new TreeNode(key, value, node));
66 return true;
67 }
68 }
69 return false;
70 }
71
72 remove(key: string) {
73 for (const node of this.preOrderTraversal()) {
74 const filtered = node.children.filter((c) => c.key !== key);
75 if (filtered.length !== node.children.length) {
76 node.children = filtered;
77 return true;
78 }
79 }
80 return false;
81 }
82
83 find(key: string) {
84 for (const node of this.preOrderTraversal()) {
85 if (node.key === key) return node;
86 }
87 return undefined;
88 }
89
90 size() {
91 let size = 0;
92
93 for (const _ of this.preOrderTraversal()) {
94 size++;
95 }
96 return size;
97 }
98
99 replaceNode(oldNodeKey: string, newNodeKey: string, newNodeValue: T) {
100 for (const node of this.preOrderTraversal()) {
101 const oldNodeIndex = node.children.findIndex((c) => c.key === oldNodeKey);
102 if (oldNodeIndex >= 0) {
103 const newNode = new TreeNode(newNodeKey, newNodeValue, node);
104 newNode.children = [...(node.children[oldNodeIndex]?.children ?? [])];
105 node.children.splice(oldNodeIndex, 1, newNode);
106 return;
107 }
108 }
109 }
110
111 /**
112 * Returns an array constructed from iterating over all the nodes
113 * in the tree and extracting a particular value from each node,
114 * as defined the by extract function.
115 */
116 getValues<U>(extract: (node: T) => U) {
117 const values: U[] = [];
118 for (const node of this.preOrderTraversal()) {
119 values.push(extract(node.value));
120 }
121 return values;
122 }
123
124 /**
125 * Function that determines whether all nodes in the
126 * tree are unique, or whether there are duplicate nodes.
127 * @param isEqual - a function that determines equality
128 * between two objects of type T
129 */
130 duplicateValues(isEqual: (v1: T, v2: T) => boolean) {
131 const nodes: TreeNode<T>[] = [];
132 const duplicates: T[] = [];
133 for (const node of this.preOrderTraversal()) {
134 if (
135 nodes.some((previousNode) => isEqual(node.value, previousNode.value))
136 ) {
137 duplicates.push(node.value);
138 }
139 nodes.push(node);
140 }
141 return duplicates;
142 }
143
144 filterTree(filterPredicate: (nodeValue: T) => boolean): Tree<T> | null {
145 const filterRecursive = (currentNode: TreeNode<T>): TreeNode<T> | null => {
146 const filteredChildren = currentNode.children
147 .map((child) => filterRecursive(child))
148 .filter((child): child is TreeNode<T> => child !== null);
149
150 if (filterPredicate(currentNode.value) || filteredChildren.length > 0) {
151 const newNode = new TreeNode(currentNode.key, currentNode.value);
152 newNode.children = filteredChildren;
153 for (const child of filteredChildren) {
154 child.parent = newNode;
155 }
156 return newNode;
157 }
158
159 return null;
160 };
161
162 const filteredRoot = filterRecursive(this.root);
163
164 if (!filteredRoot) {
165 return null;
166 }
167
168 // Manually reconstruct the tree from the filteredRoot
169 const newTree = new Tree<T>(filteredRoot.key, filteredRoot.value);
170 newTree.root.children = filteredRoot.children;
171 // Ensure all child nodes have their parent property correctly set
172 for (const child of newTree.root.children) {
173 child.parent = newTree.root;
174 }
175
176 return newTree;
177 }
178}
179
180export function sortTreeNodeList(nodeList: readonly any[], agg: any[]): any[] {
181 if (nodeList.length === 0) {
182 return agg;
183 }
184 const [nextLevelNodes, lowerNodes] = partition(
185 nodeList,
186 (node) =>
187 node.parentId == null || agg.some((it) => it.id === node.parentId),
188 );
189 return sortTreeNodeList(lowerNodes, [...agg, ...nextLevelNodes]);
190}
191
192/**
193 * Constructs a tree from a flattened list of nodes, where each node has an id
194 * and an optional parentId. An example input is:
195 * [
196 * { id: '1', parentId: null },
197 * { id: '2', parentId: null },
198 * { id: '3', parentId: '1' },
199 * { id: '4', parentId: '2' },
200 * ]
201 * And the output would be a tree with the following structure:
202 * root
203 * / \
204 * 1 2
205 * / \
206 * 3 4
207 */
208export function treeFromList<T>(
209 nodeList: readonly any[],
210 root: T,
211 constructNode: (values: any) => T,
212): Tree<T> {
213 const tree = new Tree('root', root);
214 const sortedList = sortTreeNodeList(nodeList, []);
215 for (const policy of sortedList) {
216 const parentKey =
217 sortedList.find((it) => it.id === policy.parentId)?.name ?? 'root';
218 tree.insert(parentKey, policy.name, constructNode(policy));
219 }
220 return tree;
221}
222
223/**
224 * Constructs a multilevel list from a flattened list of nodes, where each node
225 * has an id and an optional parentId. An example input is:
226 * [
227 * { id: '1', parentId: null },
228 * { id: '2', parentId: null },
229 * { id: '3', parentId: '1' },
230 * { id: '4', parentId: '2' },
231 * { id: '5', parentId: '2' },
232 * ]
233 * And the output would be a list with the following structure:
234 * [
235 * { id: '1', parentId: null, children: [{ id: '3', parentId: '1' }] },
236 * { id: '2', parentId: null, children: [{ id: '4', parentId: '2' }, { id: '5', parentId: '2' }] },
237 * ]
238 */
239export function multilevelListFromFlatList<
240 T extends { id: string; parentId?: string | null | undefined },
241>(
242 nodeList: readonly T[],
243): (T & {
244 children?: T[];
245})[] {
246 const map = new Map<string, T & { children?: T[] }>();
247 const roots: (T & { children?: T[] })[] = [];
248
249 // Initialize map with all nodes
250 nodeList.forEach((node) => {
251 map.set(node.id, { ...node } as T & { children?: T[] });
252 });
253
254 // Build the tree
255 nodeList.forEach((node) => {
256 if (node.parentId == null) {
257 roots.push(map.get(node.id)!);
258 } else {
259 const parent = map.get(node.parentId);
260 if (parent) {
261 if (!parent.children) {
262 parent.children = [];
263 }
264 parent.children.push(map.get(node.id)!);
265 }
266 }
267 });
268
269 return roots;
270}