lib/src/orm/model/TreeModel.ts

import {Model} from './Model'
import {collect, Collection, ErrorWithContext, Maybe} from '../../util'
import {HasSubtree} from './relation/HasSubtree'
import {Related} from './relation/decorators'
import {HasTreeParent} from './relation/HasTreeParent'
import {ModelBuilder} from './ModelBuilder'

/**
 * Model implementation with helpers for querying tree-structured data.
 *
 * This works by using a modified pre-order traversal to number the tree nodes
 * with a left- and right-side numbers. For example:
 *
 * ```txt
 *                     (1) A (14)
 *                         |
 *       (2) B (9)    (10) C (11)    (12) D (14)
 *           |
 * (3) E (6)   (7) G (8)
 *     |
 * (4) F (5)
 * ```
 *
 * These numbers are stored, by default, in `left_num` and `right_num` columns.
 * The `subtree()` method returns a `HasSubtree` relation which loads the subtree
 * of a model and recursively nests the nodes.
 *
 * You can use the `children()` helper method to get a collection of the immediate
 * children of this node, which also have the subtree set.
 *
 * To query the model without loading the entire subtree, use the `without()`
 * method on the `ModelBuilder`. For example:
 *
 * ```ts
 * MyModel.query<MyModel>().without('subtree')
 * ```
 */
export abstract class TreeModel<T extends TreeModel<T>> extends Model<T> {

    /** The table column where the left tree number is stored. */
    public static readonly leftTreeField = 'left_num'

    /** The table column where the right tree number is stored. */
    public static readonly rightTreeField = 'right_num'

    public static readonly parentIdField = 'parent_id'

    /** @override to include the tree fields */
    public static query<T2 extends Model<T2>>(): ModelBuilder<T2> {
        return super.query<T2>()
            .fields(this.rightTreeField, this.leftTreeField, this.parentIdField)
    }

    /**
     * @override to eager-load the subtree by default
     * @protected
     */
    protected with: (keyof T)[] = ['subtree']

    protected removedChildren: Collection<TreeModel<T>> = collect()

    /** @override to include the tree fields */
    public query(): ModelBuilder<T> {
        const ctor = this.constructor as typeof TreeModel
        return super.query()
            .fields(ctor.leftTreeField, ctor.rightTreeField, ctor.parentIdField)
    }

    /** Get the left tree number for this model. */
    public leftTreeNum(): Maybe<number> {
        const ctor = this.constructor as typeof TreeModel
        return this.getColumn(ctor.leftTreeField) as number
    }

    /** Get the right tree number for this model. */
    public rightTreeNum(): Maybe<number> {
        const ctor = this.constructor as typeof TreeModel
        return this.getColumn(ctor.rightTreeField) as number
    }

    /** Get the ID of this node's parent, if one exists. */
    public parentId(): Maybe<number> {
        const ctor = this.constructor as typeof TreeModel
        return this.getColumn(ctor.parentIdField) as number
    }

    /** Returns true if this node has no children. */
    public isLeaf(): boolean {
        const left = this.leftTreeNum()
        const right = this.rightTreeNum()
        return Boolean(left && right && (right - left === 1))
    }

    /** Returns true if the given `node` exists within the subtree of this node. */
    public contains(node: TreeModel<T>): boolean {
        const num = node.leftTreeNum()
        const left = this.leftTreeNum()
        const right = this.rightTreeNum()
        return Boolean(num && left && right && (left < num && right > num))
    }

    /** The subtree nodes of this model, recursively nested. */
    @Related()
    public subtree(): HasSubtree<T> {
        const ctor = this.constructor as typeof TreeModel
        return this.make<HasSubtree<T>>(HasSubtree, this, ctor.leftTreeField)
    }

    /** The parent node of this model, if one exists. */
    @Related()
    public parentNode(): HasTreeParent<T> {
        const ctor = this.constructor as typeof TreeModel
        return this.make<HasTreeParent<T>>(HasTreeParent, this, ctor.leftTreeField, ctor.parentIdField)
    }

    /** Get the immediate children of this model. */
    public children(): Collection<T> {
        return this.subtree().getValue()
    }

    /** Get the parent of this model, if one exists. */
    public parent(): Maybe<T> {
        return this.parentNode().getValue()
    }

    public root(): Maybe<TreeModel<T>> {
        // eslint-disable-next-line @typescript-eslint/no-this-alias
        let parent: Maybe<TreeModel<T>> = this
        while ( parent?.parent() ) {
            parent = parent?.parent()
        }
        return parent
    }

    public rootOrFail(): TreeModel<T> {
        const root = this.root()
        if ( !root ) {
            throw new ErrorWithContext('Unable to determine tree root', {
                node: this,
            })
        }
        return root
    }

    /** Get the nearest node in the parental line that is an ancestor of both this node and the `other` node. */
    public commonAncestorWith(other: TreeModel<T>): Maybe<TreeModel<T>> {
        if ( this.contains(other) ) {
            // Handle the special case when this node is an ancestor of the other.
            return this
        }

        if ( other.contains(this) ) {
            // Handle the special case when this node is a descendant of the other.
            return other
        }

        // Otherwise, walk up this node's ancestral line and try to find a shared ancestor.
        // It's getting too anthropological up in here.
        let parent = this.parent()
        while ( parent ) {
            if ( parent.contains(other) ) {
                return parent
            }

            parent = parent.parent()
        }
    }

    /**
     * Update the preorder traversal numbering for this node and its subtree.
     * After renumbering, you probably also want to save the subtree nodes to persist
     * the tree into the database.
     * @example
     * ```ts
     * await model.renumber().saveSubtree()
     * ```
     */
    public renumber(): this {
        // Assume our leftTreeNum is -correct-.
        // Given that, renumber the children recursively, then set our rightTreeNum
        const myLeftNum = this.leftTreeNum()
        if ( !myLeftNum ) {
            return this
        }

        const ctor = this.constructor as typeof TreeModel
        let nextChildLeftNum = myLeftNum + 1
        let myRightNum = myLeftNum + 1

        this.children()
            .each(child => {
                child.setColumn(ctor.leftTreeField, nextChildLeftNum)
                child.setColumn(ctor.parentIdField, this.key())
                child.parentNode().setValue(this as unknown as T)
                child.renumber()
                myRightNum = (child.rightTreeNum() ?? 0) + 1
                nextChildLeftNum = (child.rightTreeNum() ?? 0) + 1
            })

        this.setColumn(ctor.rightTreeField, myRightNum)
        return this
    }

    /** Get a flat collection of this node & its subtree nodes, in pre-order. */
    public flatten(): Collection<TreeModel<T>> {
        return this.flattenLevel()
    }

    /** Recursive helper for `flatten()`. */
    private flattenLevel(topLevel = true): Collection<TreeModel<T>> {
        const flat = collect<TreeModel<T>>(topLevel ? [this] : [])
        this.children().each(child => {
            flat.push(child)
            flat.concat(child.flattenLevel(false))
        })
        return flat
    }

    /** Call the `save()` method on all nodes in this node & its subtree. */
    public async saveSubtree(): Promise<void> {
        await this.save()
        await this.removedChildren.awaitMapCall('saveSubtree')
        this.removedChildren = collect()
        await this.children()
            .map(child => child.saveSubtree())
            .awaitAll()
    }

    /**
     * Append the `other` node as a child of this node.
     * Returns the common ancestor from which point the tree was renumbered.
     * You should save the subtree after this.
     * @example
     * ```ts
     * await nodeA.appendChild(nodeB).saveSubtree()
     * ```
     * @param other
     * @param before - if provided, `other` will be inserted before the `before` child of this node. Otherwise, it will be added as the last child.
     */
    public appendChild(other: T, before: Maybe<T> = undefined): TreeModel<T> {
        // Determine the index where we should insert the node in our children
        const idx = (before ? this.children().search(before) : undefined) ?? this.children().length

        // Insert the child at that index
        this.subtree().appendSubtree(idx, other)

        // If the child has a parent:
        const parent = other.parent()
        if ( parent ) {
            // Remove the child from the parent's children
            parent.subtree().removeSubtree(other)

            // Get the common ancestor of the child's parent and this node
            const ancestor = this.commonAncestorWith(other)
            if ( ancestor ) {
                // Renumber from that ancestor
                return ancestor.renumber()
            }
        }

        // If the child has no parent or a common ancestor could not be found,
        // renumber the entire tree since the total number of nodes has changed
        return this.rootOrFail().renumber()
    }

    /**
     * Remove this node from the tree structure.
     * If the node existed in the tree structure and was removed, this method
     * returns the tree node that was renumbered. You should save the subtree after this.
     * @example
     * ```ts
     * await nodeA.removeFromTree().saveSubtree()
     * ```
     */
    public removeFromTree(): Maybe<TreeModel<T>> {
        const parent = this.parent()
        if ( parent ) {
            const ctor = this.constructor as typeof TreeModel
            parent.subtree().removeSubtree(this)
            parent.removedChildren.push(this)
            this.setColumn(ctor.leftTreeField, null)
            this.setColumn(ctor.rightTreeField, null)
            this.setColumn(ctor.parentIdField, null)
            return this.rootOrFail().renumber()
        }
    }
}