yjs-orderedtree/src/ytree.js

An ordered tree class for yjs. Lets you use a Y.Map like an ordered tree with insert, delete, and move operations. The children are ordered and the position of a child amongst its sibling can be manipulated

import * as Y from "yjs"; // eslint-disable-line
import { Heap } from "heap-js";
import { callAll } from "lib0/function"
import { edgeWithLargestCounter, insertBetween } from "./util.js";


/**
 * YTree uses https://madebyevan.com/algos/crdt-mutable-tree-hierarchy/ and https://madebyevan.com/algos/crdt-fractional-indexing/
 * It uses them to implement an ordered tree data structure, on top of YJS https://github.com/yjs/yjs, that maintains synchronization and consistency across multiple clients.
 * 
 * 
 */

/**
  *
  * @param {Y.Map} yMap
  * @returns {boolean}
  * 
  * @description Checks if the Ymap has been initialized for YTree operations. 
  * 
  * Specifically checks if root node exists.
  * 
  * This should be done any time you want to load a ytree and not create it.
  */
export function checkForYTree(yMap) {
  return yMap.has("root") && yMap.get("root").has("_parentHistory");
}

/**
 * @typedef {Object} ComputedMapNode
 * @property {string} id
 * @property {ComputedMapNode} parent
 * @property {Array} children
 * @property {Map} edges
 */

/**
 * Class representing a YTree.
 */
export class YTree {
  /**
   *
   * @param {Y.Map} yMap
   * 
   * 
   * Constructor is required to be called with a Y.Map instance that is bound to a Y.Doc. 
   * The YMap is required to be empty or initialized. 
   * If given an empty (uninitialized) YMap then it initializes the YMap for YTree operations by creating a root node. 
   * 
   * 
   * If you don't want to accidentally create a Ytree then use checkForYTree beforehand to ensure that a YTree has been initialized.
   */
  constructor(yMap) {

    /**
     * @type {Y.Map} Node Map
     * 
     * The Y.Map associated with the YTree.
     * It holds all the nodes of the YTree in a 'flat' way.
     * When a node is created, it should be given a globally unique key which can be generated by YTree.generateNodeKey.
     * Each node holds a child Y.Map which has two entries: parent history and value
     * Parent History holds entries for the node's previous parents, each entry holding the counter for that parent and the order index for when that node was part of the parent.
     * Value is for the user, can be whatever value a Y.Map allows.
     */
    this._ymap = null;

    /**
     * Check if yMap has already been initialized with root node. 
     * If yes, then simply set this._ymap to yMap and this.root to root noode.
     * If not, then set this._ymap to yMap then create a root node and add it to this._ymap.
     * 
     * If YMap was not initialized, then ensure that yMap does not hold any existing entries.
     * 
     */
    if (yMap === undefined) {
      throw new Error("[ytree] expected a yMap argument")
    }
    else if (yMap.doc === null) {
      throw new Error("[ytree] expected yMap to be bounded to a Y.Doc instance");
    }
    else if (checkForYTree(yMap)) {
      this._ymap = yMap;
    }
    else if (yMap.size === 0) {
      this._ymap = yMap;
      const root = this._ymap.set("root", new Y.Map());
      root.set("_parentHistory", new Y.Map());
    }
    else {
      throw new Error(
        "[ytree] expected yMap to either be initialized for ytree or empty"
      );
    }

    /**
     * @type {Y.Doc}
     * required
     */
    this._ydoc = this._ymap.doc;

    /**
     * @type {Map<string, ComputedMapNode>}
     * 
     * As virtual map that the YTree holds.
     * Keys are node keys, and their values are the node's parent, children and parent history
     * 
     * This should be recomputed every time parent history changes or new nodes are added or deleted.
     */
    this.computedMap = new Map();

    /**
     * @type {Array<function():void>}
     */
    this._callbacks = new Array();

    /**
     * When a node is added or deleted in the associated ymap,
     * or when parent history of any node is updated,
     * the computed map must be recalculated and all the callbacks in the list must be called.
     * 
     * TODO: allow for observing specific keys instead of every key.
     */
    this._ymap.observeDeep((events, transaction) => {
      events.forEach((event) => {

        if (event.path.length == 0 || (event.path.length == 2 && event.path[1] === "_parentHistory")) {
          if (event.path.length == 0) {
            event.changes.keys.forEach((change, key) => {
              if (change.action === "update") {
                throw new Error("[ytree] The node should not be updated");
              }
            });
            // Recalculate virtual/computed map whenever a node is added or deleted.
            this.recomputeParentsAndChildren();
            callAll(this._callbacks, []);
          }

          if (event.path.length == 2 && event.path[1] === "_parentHistory") {
            // Recalculate virtual/computed map whenever a node's parentHistory changes (reparented)
            this.recomputeParentsAndChildren();
            callAll(this._callbacks, []);
          }
        }
      });
      // }
    });

    this.recomputeParentsAndChildren();
  }

  /**
   * @returns string
   * @description Generates a globally unique node key
   * Concatenate ydoc client and a random number to generate global unique node id 
   */
  generateNodeKey() {
    const max = 10 ** 8;
    let randomInt = Math.floor(Math.random() * max);
    while (this._ymap.has(this._ydoc.clientID + "" + randomInt)) {
      randomInt = Math.floor(Math.random() * max);
    }

    return this._ydoc.clientID + "" + randomInt;
  }

  /**
   * Add to callback list
   * @param {function():void} f 
   */
  observe(f) {
    this._callbacks.push(f);
  }

  /**
   * Remove from callback list
   * @param {function} f 
   */
  unobserve(f) {
    const originalLength = this._callbacks.length;
    this._callbacks = this._callbacks.filter(g => f !== g)
    if (this._callbacks.length === originalLength) {
      console.error("[ytree] tried to remove callback that does not exist")
    }
  }

  /**
   * Get the Y.Map of the YTree.
   * @returns {Y.Map} The Y.Map of the YTree.
   */
  getYMap() {
    return this._ymap;
  }

  /**
   * @param {string} nodeKey - The key of the new node.
   * @param {object | boolean | string | number | Uint8Array | Y.AbstractType} value - The value to set for the new node.
   * @description sets a value on a node
  */
  setNodeValueFromKey(nodeKey, value) {
    if (!this.computedMap.has(nodeKey)) {
      throw new Error("[ytree] node with key: " + nodeKey + " does not exist");
    }

    if (!value) {
      throw new Error("[ytree] value is required")
    }
    this._ymap.get(nodeKey).set("value", value);
  }

  /**
   * @param {string} nodeKey - The key of the node.
   * @returns {object | boolean | string | number | Uint8Array | Y.AbstractType} value - The value of the node.
   * @description get node value from key
   */
  getNodeValueFromKey(nodeKey) {
    if (!this.computedMap.has(nodeKey)) {
      throw new Error("[ytree] node with key: " + nodeKey + " does not exist");
    }
    return this._ymap.get(nodeKey).get("value");
  }

  /**
   * @param {string} nodeKey 
   * @returns {Array<string>}
   * @description get a list of keys of the node's children from the virtual (computed) map 
   */
  getNodeChildrenFromKey(nodeKey) {
    if (!this.computedMap.has(nodeKey)) {
      throw new Error("[ytree] node with key: " + nodeKey + " does not exist");
    }
    return this.computedMap.get(nodeKey).children.map((child) => child.id);
  }

  /**
   * @param {string} nodeKey 
   * @return {string}
   * @description get the key of the parent of the node from the virtual (computed) map 
   */
  getNodeParentFromKey(nodeKey) {
    if (!this.computedMap.has(nodeKey)) {
      throw new Error("[ytree] node with key: " + nodeKey + " does not exist");
    }
    return this.computedMap.get(nodeKey).parent.id;
  }

  /**
   * Creates a Node with the given parent ad the value
   * @param {string} parentKey - The key of the parent node.
   * @param {string} nodeKey - The key of the new node.
   * @param {object | boolean | string | number | Uint8Array | Y.AbstractType} value - The value to set for the new node.
   * @description Create a new node in the YTree.
   */
  createNode(parentKey, nodeKey, value) {
    this.recomputeParentsAndChildren();

    if (!this.computedMap.has(parentKey)) {
      throw new Error("[ytree] Parent with key: " + parentKey + " does not exist");
    }

    if (this.computedMap.has(nodeKey)) {
      throw new Error("[ytree] Node with key: " + nodeKey + " already exists");
    }

    const order_index = insertBetween(this._getHighestOrderIndex(this.getNodeChildrenFromKey(parentKey), parentKey), '');

    this._ydoc.transact(() => {
      const node = new Y.Map();
      node.set("value", value);

      const parentHistory = new Y.Map();
      parentHistory.set(parentKey, { counter: 0, order: order_index });

      node.set("_parentHistory", parentHistory);

      this._ymap.set(nodeKey, node);
    })
  }

  /**
   * @param {string} nodeKey - The key of the new node.
   * @description deletes the node and its descendants
   */
  deleteNodeAndDescendants(nodeKey) {
    /**
     * Make sure that the computed Map is up to date
     */
    this.recomputeParentsAndChildren();

    /** @type {Array<String>} */
    const allDescendants = new Array();
    this.getAllDescendants(nodeKey, allDescendants);



    const concatenatedDescendants = allDescendants.join(", ");

    this._ydoc.transact(() => {
      // Delete all descendants
      for (const key of allDescendants) {
        this._ymap.delete(key);
      }
    });
  }

  /**
   * @param {string} childKey - The key of the node.
   * @param {string} parentKey - The key of the new parent node.
   * @description Reparents or moves a child to a new parent
   * https://madebyevan.com/algos/crdt-mutable-tree-hierarchy/
   * Before the move operation is done, it must be ensured that the old parent and the new parent are both still rooted.
   * This is done by simply just going up the branch, making sure that a node (old/new parent and their parents and so on)'s parent is the key in the node's parent history that has the highest counter.
   * If that's not the casefor the node, then it is made so by adding an operation, operation which should set the parent as the key in the node's parent history with the highest counter.
   */
  moveChildToParent(childKey, parentKey) {
    if (!this.computedMap.has(parentKey)) {
      throw new Error("[ytree] Parent with key: " + parentKey + " does not exist");
    }

    if (!this.computedMap.has(childKey)) {
      throw new Error("[ytree] Child with key: " + childKey + " does not exist");
    }


    const new_order_index = insertBetween(this._getHighestOrderIndex(this.getNodeChildrenFromKey(parentKey), parentKey), '');

    if (
      this.isNodeUnderOtherNode(
        this.computedMap.get(parentKey),
        this.computedMap.get(childKey)
      )
    ) {
      throw new Error("[ytree] Cannot reparent a node to itself or its descendants.");
    }

    const ensureNodeIsRooted = (node) => {
      while (node) {
        const parent = node.parent;
        if (!parent) break; // Stop at the root
        const edge = edgeWithLargestCounter(this.computedMap.get(node.id));
        if (edge !== parent.id) {
          edits.push({ child: node, parent });
        }
        node = parent;
      }
    };

    // Ensure both the old and new parent paths remain rooted after reparenting
    const edits = [];
    const child = this.computedMap.get(childKey);
    const newParent = this.computedMap.get(parentKey);
    ensureNodeIsRooted(child.parent);
    ensureNodeIsRooted(newParent);

    // Add the actual reparenting operation
    edits.push({ child, parent: newParent });

    // Apply all database edits
    this._ydoc.transact(() => {
      for (const { child, parent } of edits) {
        let maxCounter = -1;
        for (const { counter, order } of this._ymap
          .get(child.id)
          .get("_parentHistory")
          .values()) {
          maxCounter = Math.max(maxCounter, counter);
        }

        if (parent.id === parentKey && child.id === childKey) {
          this._ymap
            .get(child.id)
            .get("_parentHistory")
            .set(parent.id, { counter: maxCounter + 1, order: new_order_index });
          continue;
        }

        const order_index = this._ymap
          .get(child.id)
          .get("_parentHistory").get(parent.id).order;

        this._ymap
          .get(child.id)
          .get("_parentHistory")
          .set(parent.id, { counter: maxCounter + 1, order: order_index });
      }
    })
  }

  /**
   * This holds the main logic of the function. Read here to understand it: https://madebyevan.com/algos/crdt-mutable-tree-hierarchy/
   * This is largely the same as the source code from the website. 
   */
  recomputeParentsAndChildren() {
    this.computedMap.clear();
    const computedMap = this.computedMap;

    for (const nodeKey of this._ymap.keys()) {
      computedMap.set(nodeKey, {

        id: nodeKey,
        parent: null,
        children: [],
        edges: new Map(),
      });
    }

    for (const node of computedMap.values()) {
      const edges = node.edges;
      for (const [key, { counter, order }] of this._ymap
        .get(node.id)
        .get("_parentHistory")
        .entries()) {
        if (computedMap.has(key)) {
          edges.set(key, counter);
        }
      }
    }

    computedMap.forEach((node, nodeKey) => {
      if (nodeKey === "root") return;

      node.parent = computedMap.get(edgeWithLargestCounter(node));
      node.parent.children.push(node);
    });


    const allRootDescendants = new Set();
    const descendantStack = ["root"];
    while (descendantStack.length > 0) {
      const current = descendantStack.pop();
      allRootDescendants.add(current);
      const children = this.computedMap.get(current)?.children || [];
      for (const child of children) {
        descendantStack.push(child.id);
      }
    }

    const nonRootedNodesTest = new Set();

    for (let [nodeKey, node] of computedMap.entries()) {
      if (!allRootDescendants.has(nodeKey))
        nonRootedNodesTest.add(node);

      // clean up children
      node.children = [];
    }

    const nonRootedNodes = new Set();

    for (let node of computedMap.values()) {
      if (!this.isNodeUnderOtherNode(node, computedMap.get("root"))) {
        while (node && !nonRootedNodes.has(node)) {
          nonRootedNodes.add(node);
          node = node.parent;
        }
      }
    }

    if (nonRootedNodes.size > 0) {
      const deferredEdges = new Map();
      const readyEdges = new Heap((a, b) => {
        const counterDelta = b.counter - a.counter;
        if (counterDelta !== 0) return counterDelta;
        if (a.parent.id < b.parent.id) return -1;
        if (a.parent.id > b.parent.id) return 1;
        if (a.child.id < b.child.id) return -1;
        if (a.child.id > b.child.id) return 1;
        return 0;
      });

      for (const child of nonRootedNodes.values()) {
        for (const [parentKey, counter] of child.edges.entries()) {
          const parent = computedMap.get(parentKey);
          if (!nonRootedNodes.has(parent)) {
            readyEdges.push({ child, parent, counter });
          } else {
            let edges = deferredEdges.get(parent);
            if (!edges) {
              edges = [];
              deferredEdges.set(parent, edges);
            }
            edges.push({ child, parent, counter });
          }
        }
      }

      for (let top; (top = readyEdges.pop());) {
        // Skip nodes that have already been reattached
        const child = top.child;
        if (!nonRootedNodes.has(child)) continue;

        // Reattach this node
        child.parent = top.parent;
        nonRootedNodes.delete(child);

        // Activate all deferred edges for this node
        const edges = deferredEdges.get(child);
        if (edges) for (const edge of edges) readyEdges.push(edge);
      }
    }

    // Add items as children of their parents so that the rest of the app
    // can easily traverse down the tree for drawing and hit-testing
    for (const node of computedMap.values()) {
      if (node.parent) {
        node.parent.children.push(node);
      }
    }

    // Sort each node's children by their identifiers so that all peers
    // display the same tree. In this demo, the ordering of siblings
    // under the same parent is considered unimportant. If this is
    // important for your app, you will need to use another CRDT in
    // combination with this CRDT to handle the ordering of siblings.
    for (const node of computedMap.values()) {
      node.children.sort((a, b) => {
        if (a.id < b.id) return -1;
        if (a.id > b.id) return 1;
        return 0;
      });
    }
  }

  /**
 *
 * @param {string} nodeKey
 * @param {Array<String>} allDescendants
 * 
 * @description Gets all the descendants of the node.
 */
  getAllDescendants(nodeKey, allDescendants) {
    const stack = [nodeKey];
    while (stack.length > 0) {
      const current = stack.pop();
      allDescendants.push(current);
      const children = this.computedMap.get(current)?.children || [];
      for (const child of children) {
        stack.push(child.id);
      }
    }
  }

  /**
   * 
   * @param {ComputedMapNode} node 
   * @param {ComputedMapNode} other 
   * @returns {boolean}
   * 
   * @description Returns true if and only if "node" is in the subtree under "other".
   * This function is safe to call in the presence of parent cycles.
   * https://cp-algorithms.com/others/tortoise_and_hare.html
   */
  isNodeUnderOtherNode(node, other) {
    if (node === other) return true;
    let tortoise = node;
    let hare = node.parent;
    while (hare && hare !== other) {
      if (tortoise === hare) return false; // Cycle detected
      hare = hare.parent;
      if (!hare || hare === other) break;
      tortoise = tortoise.parent;
      hare = hare.parent;
    }
    return hare === other;
  }

  /* Functions for ordering */

  /**
   * @param {string} nodeKey
   * @description sets the nodes's order index below its siblings
   */
  setNodeOrderToStart(nodeKey) {
    if (nodeKey === "root") {
      throw new Error("[ytree] cannot set order of root node")
    }

    const parent = this.computedMap.get(nodeKey).parent;
    const children = this.getNodeChildrenFromKey(parent.id)

    const order_index = insertBetween('', this._getLowestOrderIndex(children, parent.id));
    const parentHistory = this._ymap.get(nodeKey).get("_parentHistory");
    const parentCounter = parentHistory.get(parent.id).counter;
    parentHistory.set(parent.id, { counter: parentCounter, order: order_index });
  }

  /**
   * 
   * @param {string} nodeKey
   * @description sets the node's order index above its siblings
   */
  setNodeOrderToEnd(nodeKey) {
    if (nodeKey === "root") {
      throw new Error("[ytree] cannot set order of root node")
    }

    const parent = this.computedMap.get(nodeKey).parent;
    const children = this.getNodeChildrenFromKey(parent.id)

    const order_index = insertBetween(this._getHighestOrderIndex(children, parent.id), '');
    const parentHistory = this._ymap.get(nodeKey).get("_parentHistory");
    const parentCounter = parentHistory.get(parent.id).counter;
    parentHistory.set(parent.id, { counter: parentCounter, order: order_index })
  }

  /**
   * 
   * @param {string} nodeKey 
   * @param {string} target 
   * 
   * @description sets the the node's order index to be right after the target
   * 
   * The node and target must share the same parent.
   */
  setNodeAfter(nodeKey, target) {
    if (nodeKey === "root") {
      throw new Error("[ytree] cannot set order of root node")
    }

    const parent = this.computedMap.get(nodeKey).parent;

    if (this.computedMap.get(target).parent !== parent) {
      throw new Error("[ytree] expected nodes with keys nodeKey and target to have same parent")
    }

    const children = this.getNodeChildrenFromKey(parent.id);
    const order_index = insertBetween(
      this._ymap.get(target).get("_parentHistory").get(parent.id).order,
      this.getNextOrderIndex(target, children, parent.id)
    );

    const parentHistory = this._ymap.get(nodeKey).get("_parentHistory");
    const parentCounter = parentHistory.get(parent.id).counter;
    parentHistory.set(parent.id, { counter: parentCounter, order: order_index })
  }


  /**
   * 
   * @param {string} nodeKey 
   * @param {string} target 
   * 
   * @description sets the the node's order index to be right before the target
   * 
   * The node and target must share the same parent.
   */
  setNodeBefore(nodeKey, target) {
    if (nodeKey === "root") {
      throw new Error("[ytree] cannot set order of root node")
    }

    const parent = this.computedMap.get(nodeKey).parent;

    if (this.computedMap.get(target).parent !== parent) {
      throw new Error("[ytree] expected nodes with keys nodeKey and target to have same parent")
    }

    const children = this.getNodeChildrenFromKey(parent.id);
    const order_index = insertBetween(
      this.getPreviousOrderIndex(target, children, parent.id),
      this._ymap.get(target).get("_parentHistory").get(parent.id).order
    );

    const parentHistory = this._ymap.get(nodeKey).get("_parentHistory");
    const parentCounter = parentHistory.get(parent.id).counter;
    parentHistory.set(parent.id, { counter: parentCounter, order: order_index })
  }


  /**
   * Constructs an ordered array of objects based on their positions.
   *
   * @param {Array<string>} children - The array of objects to be ordered.
   * @param {string} parentKey
   * @returns {Array<string>} - An array of objects sorted by their positions.
   */
  sortChildrenByOrder(children, parentKey) {

    // It was suggested here that incremental sorting would be better. 
    // (maintaining all the children of a node in an ordered list, and whenever an order index of a child changes then simply just change it's position in the ordered list)
    // YTree crawls at a snail's at 1000 nodes, and the benefit of incremental sorting at this scale  is not worth the increase in code complexity.
    return children.sort((a, b) => {
      // Sort by position using the object identifier as a tie-breaker
      const a_order = this._ymap.get(a).get("_parentHistory").get(parentKey).order;
      const b_order = this._ymap.get(b).get("_parentHistory").get(parentKey).order;

      if (a_order < b_order) return -1;
      if (a_order > b_order) return 1;
      if (a < b) return -1;
      if (a > b) return 1;
      return 0;
    });
  }

  /**
   * gets the next order index in the array from the current item
   *
   * @param {Array<string>} children
   * @param {string} key
   * @param {string} parentKey
   * @returns {String}
   */
  getNextOrderIndex(key, children, parentKey) {

    if (children.length === 0) return "";

    const currentItemOrderIndex = this._ymap.get(key).get("_parentHistory").get(parentKey).order


    const nextOrderIndex = children.reduce((accumulator, currentValue) => {
      const current_order = this._ymap.get(currentValue).get("_parentHistory").get(parentKey).order
      if (accumulator === "") {
        if (current_order > currentItemOrderIndex) {
          return current_order;
        } else {
          return accumulator;
        }
      }

      if (
        current_order < accumulator &&
        current_order > currentItemOrderIndex
      ) {
        return current_order;
      }

      return accumulator;
    }, "");

    return nextOrderIndex;
  }

  /**
   * gets the next order index in the array from the current item
   *
   * @param {Array<string>} children
   * @param {string} key
   * @param {string} parentKey
   * @returns {String}
   */
  getPreviousOrderIndex(key, children, parentKey) {
    if (children.length === 0) return "";

    const currentItemOrderIndex = this._ymap.get(key).get("_parentHistory").get(parentKey).order


    const previousOrderIndex = children.reduce((accumulator, currentValue) => {
      const current_order = this._ymap.get(currentValue).get("_parentHistory").get(parentKey).order
      if (accumulator === "") {
        if (current_order < currentItemOrderIndex) {
          return current_order;
        } else {
          return accumulator;
        }
      }

      if (
        current_order > accumulator &&
        current_order < currentItemOrderIndex
      ) {
        return current_order;
      }

      return accumulator;
    }, "");

    return previousOrderIndex;
  }

  /**
   * 
   * @param {Array<string>} children 
   * @param {string} parentKey
   * @returns {string}
   */
  _getHighestOrderIndex(children, parentKey) {
    if (children.length === 0) return "";

    const first_order_index = this._ymap.get(children[0]).get("_parentHistory").get(parentKey).order

    const highestOrderIndex = children.reduce((accumulator, currentValue) => {
      const current_order = this._ymap.get(currentValue).get("_parentHistory").get(parentKey).order

      if (current_order > accumulator) {
        return current_order;
      }

      return accumulator;
    }, first_order_index);

    return highestOrderIndex;
  }

  /**
   * 
   * @param {Array<string>} children 
   * @param {string} parentKey
   * @returns {string}
   */
  _getLowestOrderIndex(children, parentKey) {
    if (children.length === 0) return "";

    const lowestOrderIndex = children.reduce((accumulator, currentValue) => {
      const current_order = this._ymap.get(currentValue).get("_parentHistory").get(parentKey).order

      if (current_order < accumulator) {
        return current_order;
      }

      return accumulator;
    }, this._ymap.get(children[0]).get("_parentHistory").get(parentKey).order);

    return lowestOrderIndex;
  }
};