mol-state: added StateTreeSpine

src/mol-state/state.ts

@@ -18,6 +18,7 @@ import { TransientTree } from './tree/transient';
 import { LogEntry } from 'mol-util/log-entry';
 import { now, formatTimespan } from 'mol-util/now';
 import { ParamDefinition } from 'mol-util/param-definition';
+import { StateTreeSpine } from './tree/spine';
 
 export { State }
 
@@ -59,6 +60,7 @@ class State {
     build() { return new StateBuilder.Root(this._tree); }
 
     readonly cells: State.Cells = new Map();
+    private spine = new StateTreeSpine.Impl(this.cells);
 
     getSnapshot(): State.Snapshot {
         return { tree: StateTree.toJSON(this._tree) };
@@ -140,6 +142,8 @@ class State {
                     return cell && cell.obj;
                 }
             } finally {
+                this.spine.setSurrent();
+
                 if (updated) this.events.changed.next();
                 this.events.isUpdating.next(false);
 
@@ -164,6 +168,7 @@ class State {
             oldTree,
             tree: _tree,
             cells: this.cells as Map<StateTransform.Ref, StateObjectCell>,
+            spine: this.spine,
 
             results: [],
             stateChanges: [],
@@ -243,6 +248,7 @@ interface UpdateContext {
     oldTree: StateTree,
     tree: TransientTree,
     cells: Map<StateTransform.Ref, StateObjectCell>,
+    spine: StateTreeSpine.Impl,
 
     results: UpdateNodeResult[],
     stateChanges: StateTransform.Ref[],
@@ -593,12 +599,14 @@ async function updateNode(ctx: UpdateContext, currentRef: Ref): Promise<UpdateNo
     }
 
     let parentCell = transform.transformer.definition.from.length === 0
-        ? ctx.cells.get(current.transform.parent)
-        : StateSelection.findAncestorOfType(tree, ctx.cells, currentRef, transform.transformer.definition.from);
+    ? ctx.cells.get(current.transform.parent)
+    : StateSelection.findAncestorOfType(tree, ctx.cells, currentRef, transform.transformer.definition.from);
     if (!parentCell) {
         throw new Error(`No suitable parent found for '${currentRef}'`);
     }
 
+    ctx.spine.setSurrent(current);
+
     const parent = parentCell.obj!;
     current.sourceRef = parentCell.transform.ref;
 
@@ -653,7 +661,7 @@ function runTask<T>(t: T | Task<T>, ctx: RuntimeContext) {
 
 function createObject(ctx: UpdateContext, cell: StateObjectCell, transformer: StateTransformer, a: StateObject, params: any) {
     if (!cell.cache) cell.cache = Object.create(null);
-    return runTask(transformer.definition.apply({ a, params, cache: cell.cache }, ctx.parent.globalContext), ctx.taskCtx);
+    return runTask(transformer.definition.apply({ a, params, cache: cell.cache, spine: ctx.spine }, ctx.parent.globalContext), ctx.taskCtx);
 }
 
 async function updateObject(ctx: UpdateContext, cell: StateObjectCell,  transformer: StateTransformer, a: StateObject, b: StateObject, oldParams: any, newParams: any) {
@@ -661,5 +669,5 @@ async function updateObject(ctx: UpdateContext, cell: StateObjectCell,  transfor
         return StateTransformer.UpdateResult.Recreate;
     }
     if (!cell.cache) cell.cache = Object.create(null);
-    return runTask(transformer.definition.update({ a, oldParams, b, newParams, cache: cell.cache }, ctx.parent.globalContext), ctx.taskCtx);
+    return runTask(transformer.definition.update({ a, oldParams, b, newParams, cache: cell.cache, spine: ctx.spine }, ctx.parent.globalContext), ctx.taskCtx);
 }

src/mol-state/transformer.ts

@@ -10,6 +10,7 @@ import { StateTransform } from './transform';
 import { ParamDefinition as PD } from 'mol-util/param-definition';
 import { StateAction } from './action';
 import { capitalize } from 'mol-util/string';
+import { StateTreeSpine } from './tree/spine';
 
 export { Transformer as StateTransformer }
 
@@ -36,7 +37,8 @@ namespace Transformer {
         a: A,
         params: P,
         /** A cache object that is purged each time the corresponding StateObject is removed or recreated. */
-        cache: unknown
+        cache: unknown,
+        spine: StateTreeSpine
     }
 
     export interface UpdateParams<A extends StateObject = StateObject, B extends StateObject = StateObject, P extends {} = {}> {
@@ -45,7 +47,8 @@ namespace Transformer {
         oldParams: P,
         newParams: P,
         /** A cache object that is purged each time the corresponding StateObject is removed or recreated. */
-        cache: unknown
+        cache: unknown,
+        spine: StateTreeSpine
     }
 
     export interface AutoUpdateParams<A extends StateObject = StateObject, B extends StateObject = StateObject, P extends {} = {}> {

src/mol-state/tree/spine.ts

@@ -0,0 +1,55 @@
+/**
+ * Copyright (c) 2019 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { State } from '../state';
+import { StateTransform } from '../transform';
+import { StateObject, StateObjectCell } from '../object';
+
+export { StateTreeSpine }
+
+/** The tree spine allows access to ancestor of a node during reconciliation. */
+interface StateTreeSpine {
+    getAncestorOfType<T extends StateObject.Ctor>(type: T): StateObject.From<T> | undefined;
+    getRootOfType<T extends StateObject.Ctor>(type: T): StateObject.From<T> | undefined;
+}
+
+namespace StateTreeSpine {
+    export class Impl implements StateTreeSpine {
+        private current: StateObjectCell | undefined = void 0;
+        setSurrent(cell?: StateObjectCell) {
+            this.current = cell;
+        }
+
+        getAncestorOfType<T extends StateObject.Ctor>(t: T): StateObject.From<T> | undefined {
+            if (!this.current) return void 0;
+            let cell = this.current;
+            while (true) {
+                cell = this.cells.get(cell.transform.parent)!;
+                if (!cell.obj) return void 0;
+                if (cell.obj.type === t.type) return cell.obj as StateObject.From<T>;
+                if (cell.transform.ref === StateTransform.RootRef) return void 0;
+            }
+        }
+
+        getRootOfType<T extends StateObject.Ctor>(t: T): StateObject.From<T> | undefined {
+            if (!this.current) return void 0;
+            let cell = this.current;
+            let ret: StateObjectCell | undefined = void 0;
+            while (true) {
+                cell = this.cells.get(cell.transform.parent)!;
+                if (!cell.obj) return void 0;
+                if (cell.obj.type === t.type) {
+                    ret = cell;
+                }
+                if (cell.transform.ref === StateTransform.RootRef) return ret ? ret.obj as StateObject.From<T> : void 0;
+            }
+        }
+
+        constructor(private cells: State.Cells) {
+
+        }
+    }
+}