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'
import {Inject, Injectable} from '../../../di'
import {Logging} from '../../../service/Logging'

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

    @Inject()
    protected readonly logging!: Logging

    /**
     * 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)
    }

    public flatten(): Collection<T> {
        const children = this.getValue()
        const subtrees = children.reduce((subtree, child) => subtree.concat(child.subtree().flatten()), collect())
        return children.concat(subtrees)
    }

    /** Manually load the subtree. */
    public async load(): Promise<void> {
        this.setValue(await this.get())
    }

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

    public query(): RelationBuilder<T> {
        return this.builder()
            .tap(b => this.model.applyScopes(b))
            .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')

        this.logging.debug(`Building eager query for parent: ${parentQuery}`)
        this.logging.debug(result)

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

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

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

        this.logging.debug(`Built eager query: ${query}`)

        return query
    }

    public matchResults(possiblyRelated: Collection<T>): Collection<T> {
        this.logging.debug('Matching possibly related: ' + possiblyRelated.length)
        this.logging.verbose(possiblyRelated)
        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 appendSubtree(idx: number, subtree: T): void {
        if ( !this.instances ) {
            throw new RelationNotLoadedError()
        }

        this.instances = this.instances.put(idx, subtree)
    }

    public removeSubtree(subtree: TreeModel<T>): void {
        if ( !this.instances ) {
            throw new RelationNotLoadedError()
        }

        this.instances = this.instances.filter(x => x.isNot(subtree))
    }

    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)
        }

        // Set the parent relation on the immediate children we identified
        children.each(child => child.parentNode().setValue(this.model))

        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()
    }
}