mol-state: dependent transforms

src/mol-state/tree/immutable.ts

@@ -18,6 +18,8 @@ interface StateTree {
     readonly root: StateTransform,
     readonly transforms: StateTree.Transforms,
     readonly children: StateTree.Children,
+    /** Refs to all nodes that depend on the given key */
+    readonly dependencies: StateTree.Dependencies,
 
     asTransient(): TransientTree
 }
@@ -41,6 +43,7 @@ namespace StateTree {
 
     export interface Transforms extends _Map<StateTransform> {}
     export interface Children extends _Map<ChildSet> { }
+    export interface Dependencies extends _Map<ChildSet> { }
 
     class Impl implements StateTree {
         get root() { return this.transforms.get(StateTransform.RootRef)! }
@@ -49,7 +52,7 @@ namespace StateTree {
             return new TransientTree(this);
         }
 
-        constructor(public transforms: StateTree.Transforms, public children: Children) {
+        constructor(public transforms: Transforms, public children: Children, public dependencies: Dependencies) {
         }
     }
 
@@ -58,11 +61,11 @@ namespace StateTree {
      */
     export function createEmpty(customRoot?: StateTransform): StateTree {
         const root = customRoot || StateTransform.createRoot();
-        return create(ImmutableMap([[root.ref, root]]), ImmutableMap([[root.ref, OrderedSet()]]));
+        return create(ImmutableMap([[root.ref, root]]), ImmutableMap([[root.ref, OrderedSet()]]), ImmutableMap());
     }
 
-    export function create(nodes: Transforms, children: Children): StateTree {
-        return new Impl(nodes, children);
+    export function create(nodes: Transforms, children: Children, dependencies: Dependencies): StateTree {
+        return new Impl(nodes, children, dependencies);
     }
 
     type VisitorCtx = { tree: StateTree, state: any, f: (node: StateTransform, tree: StateTree, state: any) => boolean | undefined | void };
@@ -133,6 +136,7 @@ namespace StateTree {
     export function fromJSON(data: Serialized): StateTree {
         const nodes = ImmutableMap<Ref, StateTransform>().asMutable();
         const children = ImmutableMap<Ref, OrderedSet<Ref>>().asMutable();
+        const dependencies = ImmutableMap<Ref, OrderedSet<Ref>>().asMutable();
 
         for (const t of data.transforms) {
             const transform = StateTransform.fromJSON(t);
@@ -145,12 +149,28 @@ namespace StateTree {
             if (transform.ref !== transform.parent) children.get(transform.parent).add(transform.ref);
         }
 
+        const dependent = new Set<Ref>();
         for (const t of data.transforms) {
             const ref = t.ref;
             children.set(ref, children.get(ref).asImmutable());
+
+            if (!t.dependsOn) continue;
+
+            for (const d of t.dependsOn) {
+                dependent.add(d);
+                if (!dependencies.has(d)) {
+                    dependencies.set(d, OrderedSet<Ref>([ref]).asMutable());
+                } else {
+                    dependencies.get(d).add(ref);
+                }
+            }
         }
 
-        return create(nodes.asImmutable(), children.asImmutable());
+        dependent.forEach(d => {
+            dependencies.set(d, dependencies.get(d).asImmutable());
+        });
+
+        return create(nodes.asImmutable(), children.asImmutable(), dependencies.asImmutable());
     }
 
     export function dump(tree: StateTree) {

src/mol-state/tree/transient.ts

@@ -14,11 +14,14 @@ export { TransientTree }
 class TransientTree implements StateTree {
     transforms = this.tree.transforms as ImmutableMap<StateTransform.Ref, StateTransform>;
     children = this.tree.children as ImmutableMap<StateTransform.Ref, OrderedSet<StateTransform.Ref>>;
+    dependencies = this.tree.dependencies as ImmutableMap<StateTransform.Ref, OrderedSet<StateTransform.Ref>>;
 
     private changedNodes = false;
     private changedChildren = false;
+    private changedDependencies = false;
 
     private _childMutations: Map<StateTransform.Ref, OrderedSet<StateTransform.Ref>> | undefined = void 0;
+    private _dependencyMutations: Map<StateTransform.Ref, OrderedSet<StateTransform.Ref>> | undefined = void 0;
     private _stateUpdates: Set<StateTransform.Ref> | undefined = void 0;
 
     private get childMutations() {
@@ -27,6 +30,12 @@ class TransientTree implements StateTree {
         return this._childMutations;
     }
 
+    private get dependencyMutations() {
+        if (this._dependencyMutations) return this._dependencyMutations;
+        this._dependencyMutations = new Map();
+        return this._dependencyMutations;
+    }
+
     private changeNodes() {
         if (this.changedNodes) return;
         this.changedNodes = true;
@@ -39,6 +48,12 @@ class TransientTree implements StateTree {
         this.children = this.children.asMutable();
     }
 
+    private changeDependencies() {
+        if (this.changedDependencies) return;
+        this.changedDependencies = true;
+        this.dependencies = this.dependencies.asMutable();
+    }
+
     get root() { return this.transforms.get(StateTransform.RootRef)! }
 
     asTransient() {
@@ -82,6 +97,26 @@ class TransientTree implements StateTree {
         this.childMutations.set(parent, set);
     }
 
+    private mutateDependency(parent: StateTransform.Ref, child: StateTransform.Ref, action: 'add' | 'remove') {
+        let set = this.dependencyMutations.get(parent);
+
+        if (!set) {
+            const src = this.dependencies.get(parent);
+            if (!src && action === 'remove') return;
+
+            this.changeDependencies();
+            set = src ? src.asMutable() : OrderedSet<string>().asMutable();
+            this.dependencyMutations.set(parent, set);
+            this.dependencies.set(parent, set);
+        }
+
+        if (action === 'add') {
+            set.add(child);
+        } else {
+            set.remove(child);
+        }
+    }
+
     changeParent(ref: StateTransform.Ref, newParent: StateTransform.Ref) {
         ensurePresent(this.transforms, ref);
 
@@ -118,6 +153,12 @@ class TransientTree implements StateTree {
         this.changeNodes();
         this.transforms.set(ref, transform);
 
+        if (transform.dependsOn) {
+            for (const d of transform.dependsOn) {
+                this.mutateDependency(d, ref, 'add');
+            }
+        }
+
         return this;
     }
 
@@ -178,6 +219,21 @@ class TransientTree implements StateTree {
             this.transforms.delete(n.ref);
             this.children.delete(n.ref);
             if (this._childMutations) this._childMutations.delete(n.ref);
+
+            if (this.dependencies.has(n.ref)) {
+                this.changeDependencies();
+                this.dependencies.delete(n.ref);
+                if (this._dependencyMutations) this._dependencyMutations.delete(n.ref);
+            }
+        }
+
+        for (const n of st) {
+            if (!n.dependsOn) continue;
+
+            for (const d of n.dependsOn) {
+                if (!this.transforms.has(d)) continue;
+                this.mutateDependency(d, n.ref, 'remove');
+            }
         }
 
         return st;
@@ -186,9 +242,11 @@ class TransientTree implements StateTree {
     asImmutable() {
         if (!this.changedNodes && !this.changedChildren && !this._childMutations) return this.tree;
         if (this._childMutations) this._childMutations.forEach(fixChildMutations, this.children);
+        if (this._dependencyMutations) this._dependencyMutations.forEach(fixDependencyMutations, this.dependencies);
         return StateTree.create(
             this.changedNodes ? this.transforms.asImmutable() : this.transforms,
-            this.changedChildren ? this.children.asImmutable() : this.children);
+            this.changedChildren ? this.children.asImmutable() : this.children,
+            this.changedDependencies ? this.dependencies.asImmutable() : this.dependencies);
     }
 
     constructor(private tree: StateTree) {
@@ -196,7 +254,14 @@ class TransientTree implements StateTree {
     }
 }
 
-function fixChildMutations(this: ImmutableMap<StateTransform.Ref, OrderedSet<StateTransform.Ref>>, m: OrderedSet<StateTransform.Ref>, k: StateTransform.Ref) { this.set(k, m.asImmutable()); }
+function fixChildMutations(this: ImmutableMap<StateTransform.Ref, OrderedSet<StateTransform.Ref>>, m: OrderedSet<StateTransform.Ref>, k: StateTransform.Ref) {
+    this.set(k, m.asImmutable());
+}
+
+function fixDependencyMutations(this: ImmutableMap<StateTransform.Ref, OrderedSet<StateTransform.Ref>>, m: OrderedSet<StateTransform.Ref>, k: StateTransform.Ref) {
+    if (m.size === 0) this.delete(k);
+    else this.set(k, m.asImmutable());
+}
 
 function alreadyPresent(ref: StateTransform.Ref) {
     throw new Error(`Transform '${ref}' is already present in the tree.`);