1  import emptyStack from '@iter-tools/imm-stack';
  2  
  3  const { isArray } = Array;
  4  const { freeze: freezeObject, isFrozen } = Object;
  5  const { isFinite } = Number;
  6  
  7  export const buildModule = (NODE_SIZE = 8) => {
  8    const sumNodes = (nodes) => {
  9      return nodes.map(getSum).reduce((a, b) => a + b, 0);
 10    };
 11  
 12    const treeFrom = (...trees) => [sumNodes(trees), trees];
 13  
 14    const from = (...values) => {
 15      let tree = [];
 16      for (const value of values) {
 17        tree = push(tree, value);
 18      }
 19      return tree;
 20    };
 21  
 22    const fromValues = (values) => {
 23      let tree = [];
 24      for (const value of values) {
 25        tree = push(tree, value);
 26      }
 27      return tree;
 28    };
 29  
 30    const findBalancePoint = (tree) => {
 31      const values = isFinite(tree[0]) ? tree[1] : tree;
 32      let leftSum = 0;
 33      let rightSum = getSum(tree);
 34      let balance = leftSum / rightSum;
 35  
 36      if (!values.length) return null;
 37      if (values.length === 1) return 1;
 38  
 39      for (let i = 1; i < values.length; i++) {
 40        const sum = getSum(values[i - 1]);
 41        const lastBalance = balance;
 42        leftSum += sum;
 43        rightSum -= sum;
 44  
 45        balance = leftSum / rightSum;
 46  
 47        if (lastBalance < 1 && balance >= 1) {
 48          return i;
 49        }
 50      }
 51      return values.length - 1;
 52    };
 53  
 54    const split = (tree) => {
 55      const values = isFinite(tree[0]) ? tree[1] : tree;
 56      const isLeaf = isLeafNode(tree);
 57  
 58      let midIndex;
 59  
 60      if (isLeaf) {
 61        midIndex = Math.floor(values.length / 2 + 0.01);
 62      } else {
 63        midIndex = findBalancePoint(tree);
 64      }
 65  
 66      let leftValues = values.slice(0, midIndex);
 67      let rightValues = values.slice(midIndex);
 68  
 69      let left = isLeaf ? leftValues : [sumNodes(leftValues), leftValues];
 70      let right = isLeaf ? rightValues : [sumNodes(rightValues), rightValues];
 71  
 72      return { left, right };
 73    };
 74  
 75    const setValuesAt = (idx, node, value) => {
 76      const isLeaf = isLeafNode(node);
 77      const values = isLeaf ? node : node[1];
 78  
 79      if (!isLeaf && !isArray(value)) {
 80        throw new Error();
 81      }
 82  
 83      const oldSize = isLeaf ? 1 : values[idx] ? getSum(values[idx]) : 0;
 84      const newSize = isLeaf ? 1 : getSum(value);
 85  
 86      // TODO mutable sets?
 87  
 88      const newValues = [...values];
 89      newValues[idx] = value;
 90      freezeObject(newValues);
 91      return isLeaf ? newValues : freezeObject([node[0] + newSize - oldSize, newValues]);
 92    };
 93  
 94    const addAt = (idx, tree, value) => {
 95      if (idx < 0) throw new Error('invalid argument');
 96  
 97      let path = findPath(idx, tree);
 98  
 99      let { node, index } = path.value;
100  
101      // left pushout vs right pushout?
102      let pushout = value;
103  
104      let values = [...getValues(node)];
105  
106      for (;;) {
107        if (pushout) {
108          values = [...values];
109          if (values.length + getValues(pushout).length > NODE_SIZE) {
110            values.splice(index, 0, pushout);
111            node = setValues(node, values);
112            const { left, right } = split(node);
113  
114            pushout = left;
115            node = right;
116          } else {
117            values.splice(index, 0, pushout);
118            node = setValues(node, values);
119            pushout = null;
120          }
121        }
122  
123        if (path.size === 1) {
124          if (pushout) {
125            return treeFrom(pushout, node);
126          } else {
127            return node;
128          }
129        }
130  
131        const poppedNode = node;
132        path = path.pop();
133        ({ node, index } = path.value);
134  
135        node = setValuesAt(index, node, poppedNode);
136        values = getValues(node);
137        path = path.replace({ node, index });
138      }
139    };
140  
141    const push = (tree, value) => {
142      return addAt(getSum(tree), tree, value);
143    };
144  
145    const collapses = (size) => {
146      return size > NODE_SIZE / 2 + 0.01;
147    };
148  
149    const nodeCollapses = (node) => {
150      return collapses(getValues(node).length);
151    };
152  
153    const nodeCanDonate = (node) => {
154      return collapses(getValues(node).length - 1);
155    };
156  
157    const pop = (tree) => {
158      let path = findPath(-1, tree);
159  
160      let { node, index } = path.value;
161  
162      const initialValues = [...getValues(node)].slice(0, -1);
163      let returnValue = isLeafNode(node) ? initialValues : [sumNodes(initialValues), initialValues];
164  
165      for (;;) {
166        let values = getValues(returnValue);
167        let adjustSibling = null;
168  
169        if (path.size > 1 && nodeCollapses(returnValue)) {
170          let { node: parentNode, index: parentIndex } = path.prev.value;
171          const prevSibling = getValues(parentNode)[parentIndex - 1];
172          const nextSibling = getValues(parentNode)[parentIndex + 1];
173          let targetSibling = nodeCanDonate(prevSibling)
174            ? prevSibling
175            : nodeCanDonate(nextSibling)
176            ? nextSibling
177            : null;
178          let targetSiblingIndex = targetSibling && (prevSibling ? parentIndex - 1 : parentIndex + 1);
179  
180          if (targetSibling) {
181            let targetValues = [...getValues(targetSibling)];
182  
183            const donationIdx = targetSibling === prevSibling ? targetValues.length - 1 : 0;
184            const donated = targetValues[donationIdx];
185            targetValues.splice(donationIdx, 1);
186  
187            adjustSibling = {
188              node: setValues(targetSibling, targetValues),
189              index: targetSiblingIndex,
190            };
191  
192            values = [...values];
193  
194            values.splice(targetSibling === prevSibling ? values.length : 0, 0, donated);
195  
196            returnValue = setValues(returnValue, values);
197          }
198        }
199  
200        if (path.size === 1) {
201          if (isFinite(returnValue[0]) && getSum(returnValue) <= NODE_SIZE) {
202            returnValue = returnValue[1].flat();
203          }
204          return returnValue;
205        }
206  
207        path = path.pop();
208        ({ node, index } = path.value);
209  
210        values = [...getValues(node)];
211  
212        values.splice(index, 1, returnValue);
213  
214        if (adjustSibling) {
215          const { index, node } = adjustSibling;
216          values.splice(index, 1, ...(node ? [node] : []));
217        }
218  
219        returnValue = node = setValues(node, values);
220      }
221    };
222  
223    const isValidNode = (node) => {
224      if (!isArray(node)) return false;
225      const values = isFinite(node[0]) ? node[1] : node;
226      return isArray(values); //  && values.length <= NODE_SIZE;
227    };
228  
229    const assertValidNode = (node) => {
230      if (!isValidNode(node)) throw new Error();
231      return true;
232    };
233  
234    const getValues = (node) => {
235      return node ? (isArray(node) ? (isFinite(node[0]) ? node[1] : node) : [node]) : [];
236    };
237  
238    const setValues = (node, values) => {
239      return isFinite(node[0]) ? treeFrom(...values) : values;
240    };
241  
242    const isLeafNode = (node) => {
243      return isArray(node) && !isFinite(node[0]);
244    };
245  
246    function* traverse(tree) {
247      let states = emptyStack.push({ node: tree, i: 0 });
248  
249      assertValidNode(tree);
250  
251      stack: while (states.size) {
252        const s = states.value;
253        const { node } = s;
254        const isLeaf = isLeafNode(node);
255  
256        const values = isLeaf ? node : node[1];
257  
258        if (isLeaf) {
259          for (let i = 0; i < values.length; i++) {
260            yield values[i];
261          }
262        } else {
263          for (let { i } = s; s.i < values.length; ) {
264            const node = values[i];
265            assertValidNode(node);
266            states = states.push({ node, i: 0 });
267            i = ++s.i;
268            continue stack;
269          }
270        }
271        states = states.pop();
272      }
273    }
274  
275    const getSum = (tree) => {
276      if (!isArray(tree)) {
277        return 1;
278      } else if (isFinite(tree[0])) {
279        return tree[0];
280      } else {
281        return tree.length;
282      }
283    };
284  
285    function* indexes(count, backwards = false) {
286      const increment = backwards ? -1 : 1;
287      for (let i = backwards ? count - 1 : 0; backwards ? i >= 0 : i < count; i += increment) {
288        yield i;
289      }
290    }
291  
292    const findPath = (idx, tree) => {
293      let path = emptyStack;
294  
295      let treeSum = getSum(tree);
296      let currentIdx = idx < 0 ? treeSum : 0;
297      let direction = idx < 0 ? -1 : 1;
298      let targetIdx = idx < 0 ? treeSum + idx : idx;
299  
300      let node = tree;
301      stack: while (node) {
302        assertValidNode(node);
303  
304        if (isLeafNode(node)) {
305          const startIdx = idx < 0 ? currentIdx - getSum(node) : currentIdx;
306          let index = isFinite(currentIdx) ? targetIdx - startIdx : currentIdx;
307          if (index < 0) {
308            index = -Infinity;
309          } else if (index >= node.length) {
310            index = Infinity;
311          }
312          return path.push({ index, node });
313        } else {
314          const values = node[1];
315          let candidateNode;
316          let i;
317  
318          for (i of indexes(values.length, idx < 0)) {
319            candidateNode = values[i];
320            const sum = getSum(candidateNode);
321            const nextCount = currentIdx + sum * direction;
322            if (idx < 0 ? nextCount <= targetIdx : nextCount > targetIdx || nextCount >= treeSum) {
323              path = path.push({ index: i, node });
324              node = candidateNode;
325              continue stack;
326            } else {
327              currentIdx += sum * direction;
328            }
329          }
330        }
331      }
332  
333      return null;
334    };
335  
336    const getAt = (idx, tree) => {
337      const v = findPath(idx, tree)?.value;
338      return v && v.node[v.index];
339    };
340  
341    const replaceAt = (idx, tree, value) => {
342      let path = findPath(idx, tree);
343  
344      let { node, index } = path.value;
345  
346      let returnValue = setValuesAt(index, node, value);
347  
348      for (;;) {
349        ({ node, index } = path.value);
350  
351        if (path.size > 1) {
352          path = path.pop();
353          ({ node, index } = path.value);
354  
355          returnValue = setValuesAt(index, node, returnValue);
356        } else {
357          return returnValue;
358        }
359      }
360    };
361  
362    const freeze = (tree) => {
363      let states = emptyStack.push({ node: tree, i: 0 });
364  
365      stack: while (states.size) {
366        const s = states.value;
367        const { node } = s;
368        const isLeaf = isLeafNode(node);
369  
370        assertValidNode(node);
371  
372        freezeObject(node);
373  
374        const values = isLeaf ? node : node[1];
375  
376        if (!isLeaf) {
377          freezeObject(values);
378  
379          for (let { i } = s; s.i < values.length; ) {
380            const node = values[i];
381            assertValidNode(node);
382            states = states.push({ node, i: 0 });
383            i = ++s.i;
384            continue stack;
385          }
386        }
387        states = states.pop();
388      }
389      return tree;
390    };
391  
392    return {
393      buildModule,
394      from,
395      fromValues,
396      findBalancePoint,
397      split,
398      collapses,
399      nodeCollapses,
400      nodeCanDonate,
401      pop,
402      push,
403      addAt,
404      isValidNode,
405      assertValidNode,
406      getValues,
407      setValues,
408      isLeafNode,
409      traverse,
410      getSum,
411      findPath,
412      getAt,
413      replaceAt,
414      freeze,
415    };
416  };