Add TreeModel and HasSubtree implementation

src/orm/model/relation/HasSubtree.ts

@@ -0,0 +1,147 @@
+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()
+    }
+}