lib/src/orm/model/relation/HasSubtree.ts

import {TreeModel} from '../TreeModel'
import {Relation, RelationNotLoadedError} from './Relation'
import {collect, Collection, Maybe} from '../../../util'
import {RelationBuilder} from './RelationBuilder'
import {raw} from '../../dialect/SQLDialect'
import {AbstractBuilder} from '../../builder/AbstractBuilder'
import {ModelBuilder} from '../ModelBuilder'

/**
 * A relation that recursively loads the subtree of a model using
 * modified preorder traversal.
 */
export class HasSubtree<T extends TreeModel<T>> extends Relation<T, T, Collection<T>> {

    /**
     * When the relation is loaded, the immediate children of the node.
     * @protected
     */
    protected instances: Maybe<Collection<T>>

    constructor(
        protected readonly model: T,
        protected readonly leftTreeField: string,
    ) {
        super(model, model)
    }

    protected get parentValue(): any {
        return this.model.key()
    }

    public query(): RelationBuilder<T> {
        return this.builder()
            .select(raw('*'))
            .orderByAscending(this.leftTreeField)
    }

    public applyScope(where: AbstractBuilder<T>): void {
        const left = this.model.leftTreeNum()
        const right = this.model.rightTreeNum()
        if ( !left || !right ) {
            where.whereMatchNone()
            return
        }

        where.where(this.leftTreeField, '>', left)
            .where(this.leftTreeField, '<', right)
    }

    public buildEagerQuery(parentQuery: ModelBuilder<T>, result: Collection<T>): ModelBuilder<T> {
        const query = this.model.query().without('subtree')

        if ( result.isEmpty() ) {
            return query.whereMatchNone()
        }

        result.each(inst => {
            const left = inst.leftTreeNum()
            const right = inst.rightTreeNum()
            if ( !left || !right ) {
                return
            }

            query.where(where => {
                where.where(this.leftTreeField, '>', left)
                    .where(this.leftTreeField, '<', right)
            })
        })

        return query
    }

    public matchResults(possiblyRelated: Collection<T>): Collection<T> {
        const modelLeft = this.model.leftTreeNum()
        const modelRight = this.model.rightTreeNum()
        if ( !modelLeft || !modelRight ) {
            return collect()
        }

        return possiblyRelated.filter(inst => {
            const instLeft = inst.leftTreeNum()
            return Boolean(instLeft && instLeft > modelLeft && instLeft < modelRight)
        })
    }

    public setValue(related: Collection<T>): void {
        // `related` contains a flat collection of the subtree nodes, ordered by left key ascending
        // We will loop through the related nodes and recursively call `setValue` for our immediate
        // children to build the tree.

        type ReduceState = {
            currentChild: T,
            currentSubtree: Collection<T>,
        }

        const children = this.instances = collect()
        const firstChild = related.pop()
        if ( !firstChild ) {
            return
        }

        const finalState = related.reduce<ReduceState>((state: ReduceState, node: T) => {
            if ( state.currentChild.contains(node) ) {
                // `node` belongs in the subtree of `currentChild`, not this node
                state.currentSubtree.push(node)
                return state
            }

            // We've hit the end of the subtree for `currentChild`, so set the child's
            // subtree relation value and move on to the next child.
            state.currentChild.subtree().setValue(state.currentSubtree)
            children.push(state.currentChild)

            return {
                currentChild: node,
                currentSubtree: collect(),
            }
        }, {
            currentChild: firstChild,
            currentSubtree: collect(),
        })

        // Do this one last time, since the reducer isn't called for the last node in the collection
        if ( finalState ) {
            finalState.currentChild.subtree().setValue(finalState.currentSubtree)
            children.push(finalState.currentChild)
        }

        this.instances = children.sortBy(inst => inst.getOriginalValues()?.[this.leftTreeField])
    }

    public getValue(): Collection<T> {
        if ( !this.instances ) {
            throw new RelationNotLoadedError()
        }

        return this.instances
    }

    public isLoaded(): boolean {
        return Boolean(this.instances)
    }

    public get(): Promise<Collection<T>> {
        return this.fetch().collect()
    }
}