mol-state: wip cell state

src/mol-plugin/ui/state-tree.tsx

@@ -57,17 +57,19 @@ export class StateTreeNode extends PluginComponent<{ nodeRef: string, state: Sta
             }}>{obj.label}</a> {obj.description ? <small>{obj.description}</small> : void 0}</>;
         }
 
+        const cellState = this.props.state.tree.cellStates.get(this.props.nodeRef);
+
         const expander = <>
             [<a href='#' onClick={e => {
                 e.preventDefault();
                 PluginCommands.State.ToggleExpanded.dispatch(this.context, { state: this.props.state, ref: this.props.nodeRef });
-            }}>{cell.transform.cellState.isCollapsed ? '+' : '-'}</a>]
+            }}>{cellState.isCollapsed ? '+' : '-'}</a>]
         </>;
 
         const children = this.props.state.tree.children.get(this.props.nodeRef);
         return <div>
             {remove}{children.size === 0 ? void 0 : expander} {label}
-            {cell.transform.cellState.isCollapsed || children.size === 0
+            {cellState.isCollapsed || children.size === 0
                 ? void 0
                 : <div style={{ marginLeft: '7px', paddingLeft: '3px', borderLeft: '1px solid #999' }}>{children.map(c => <StateTreeNode state={this.props.state} nodeRef={c!} key={c} />)}</div>
             }

src/mol-state/object.ts

@@ -49,7 +49,6 @@ interface StateObjectCell {
 
     version: string
     status: StateObjectCell.Status,
-    visibility: StateObjectCell.Visibility,
 
     errorText?: string,
     obj?: StateObject
@@ -65,5 +64,14 @@ namespace StateObjectCell {
 
     export const DefaultState: State = { isHidden: false, isCollapsed: false };
 
-    export type Visibility = 'visible' | 'hidden' | 'partial'
+    export function areStatesEqual(a: State, b: State) {
+        return a.isHidden !== b.isHidden || a.isCollapsed !== b.isCollapsed;
+    }
+
+    export function isStateChange(a: State, b?: Partial<State>) {
+        if (!b) return false;
+        if (typeof b.isCollapsed !== 'undefined' && a.isCollapsed !== b.isCollapsed) return true;
+        if (typeof b.isHidden !== 'undefined' && a.isHidden !== b.isHidden) return true;
+        return false;
+    }
 }

src/mol-state/state.ts

@@ -30,7 +30,7 @@ class State {
     readonly globalContext: unknown = void 0;
     readonly events = {
         object: {
-            cellState: this.ev<State.ObjectEvent & { cell: StateObjectCell }>(),
+            cellState: this.ev<State.ObjectEvent>(),
             cellCreated: this.ev<State.ObjectEvent>(),
 
             updated: this.ev<State.ObjectEvent & { action: 'in-place' | 'recreate', obj: StateObject, oldObj?: StateObject }>(),
@@ -68,13 +68,13 @@ class State {
     }
 
     updateCellState(ref: Transform.Ref, stateOrProvider: ((old: StateObjectCell.State) => Partial<StateObjectCell.State>) | Partial<StateObjectCell.State>) {
-        const cell = this.cells.get(ref)!;
-        const state = typeof stateOrProvider === 'function'
-            ? stateOrProvider(cell.transform.cellState)
+        const update = typeof stateOrProvider === 'function'
+            ? stateOrProvider(this.tree.cellStates.get(ref))
             : stateOrProvider;
 
-        cell.transform = this._tree.setCellState(ref, state);
-        this.events.object.cellState.next({ state: this, ref, cell });
+        if (this._tree.updateCellState(ref, update)) {
+            this.events.object.cellState.next({ state: this, ref });
+        }
     }
 
     dispose() {
@@ -144,7 +144,6 @@ class State {
             sourceRef: void 0,
             obj: rootObject,
             status: 'ok',
-            visibility: 'visible',
             version: root.version,
             errorText: void 0
         });
@@ -182,7 +181,7 @@ interface UpdateContext {
     errorFree: boolean,
     taskCtx: RuntimeContext,
     oldTree: StateTree,
-    tree: StateTree,
+    tree: TransientTree,
     cells: Map<Transform.Ref, StateObjectCell>,
     transformCache: Map<Ref, unknown>,
 
@@ -231,6 +230,11 @@ async function update(ctx: UpdateContext) {
         roots = findUpdateRoots(ctx.cells, ctx.tree);
     }
 
+    // Ensure cell states stay consistent
+    if (!ctx.editInfo) {
+        syncStates(ctx);
+    }
+
     // Init empty cells where not present
     // this is done in "pre order", meaning that "parents" will be created 1st.
     initCells(ctx, roots);
@@ -264,21 +268,29 @@ function findUpdateRootsVisitor(n: Transform, _: any, s: { roots: Ref[], cells:
 }
 
 type FindDeletesCtx = { newTree: StateTree, cells: State.Cells, deletes: Ref[] }
-function _visitCheckDelete(n: Transform, _: any, ctx: FindDeletesCtx) {
+function checkDeleteVisitor(n: Transform, _: any, ctx: FindDeletesCtx) {
     if (!ctx.newTree.transforms.has(n.ref) && ctx.cells.has(n.ref)) ctx.deletes.push(n.ref);
 }
 function findDeletes(ctx: UpdateContext): Ref[] {
     const deleteCtx: FindDeletesCtx = { newTree: ctx.tree, cells: ctx.cells, deletes: [] };
-    StateTree.doPostOrder(ctx.oldTree, ctx.oldTree.root, deleteCtx, _visitCheckDelete);
+    StateTree.doPostOrder(ctx.oldTree, ctx.oldTree.root, deleteCtx, checkDeleteVisitor);
     return deleteCtx.deletes;
 }
 
+function syncStatesVisitor(n: Transform, tree: StateTree, oldState: StateTree.CellStates) {
+    if (!oldState.has(n.ref)) return;
+    (tree as TransientTree).updateCellState(n.ref, oldState.get(n.ref));
+}
+function syncStates(ctx: UpdateContext) {
+    StateTree.doPreOrder(ctx.tree, ctx.tree.root, ctx.oldTree.cellStates, syncStatesVisitor);
+}
+
 function setCellStatus(ctx: UpdateContext, ref: Ref, status: StateObjectCell.Status, errorText?: string) {
     const cell = ctx.cells.get(ref)!;
     const changed = cell.status !== status;
     cell.status = status;
     cell.errorText = errorText;
-    if (changed) ctx.parent.events.object.cellState.next({ state: ctx.parent, ref, cell });
+    if (changed) ctx.parent.events.object.cellState.next({ state: ctx.parent, ref });
 }
 
 function initCellStatusVisitor(t: Transform, _: any, ctx: UpdateContext) {
@@ -294,21 +306,17 @@ function initCellStatus(ctx: UpdateContext, roots: Ref[]) {
 
 function initCellsVisitor(transform: Transform, _: any, ctx: UpdateContext) {
     if (ctx.cells.has(transform.ref)) {
-        if (transform.cellState && transform.cellState.isHidden) {
-            ctx.cells.get(transform.ref)!.visibility = 'hidden';
-        }
         return;
     }
 
-    const obj: StateObjectCell = {
+    const cell: StateObjectCell = {
         transform,
         sourceRef: void 0,
         status: 'pending',
-        visibility: transform.cellState && transform.cellState.isHidden ? 'hidden' : 'visible',
         version: UUID.create22(),
         errorText: void 0
     };
-    ctx.cells.set(transform.ref, obj);
+    ctx.cells.set(transform.ref, cell);
     ctx.parent.events.object.cellCreated.next({ state: ctx.parent, ref: transform.ref });
 }
 

src/mol-state/transform.ts

@@ -4,18 +4,17 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import { StateObject, StateObjectCell } from './object';
+import { StateObject } from './object';
 import { Transformer } from './transformer';
 import { UUID } from 'mol-util';
 
 export interface Transform<A extends StateObject = StateObject, B extends StateObject = StateObject, P = unknown> {
     readonly parent: Transform.Ref,
     readonly transformer: Transformer<A, B, P>,
-    readonly params: P,
     readonly props: Transform.Props,
     readonly ref: Transform.Ref,
-    readonly version: string,
-    readonly cellState: StateObjectCell.State
+    readonly params: P,
+    readonly version: string
 }
 
 export namespace Transform {
@@ -31,8 +30,7 @@ export namespace Transform {
 
     export interface Options {
         ref?: string,
-        props?: Props,
-        cellState?: Partial<StateObjectCell.State>
+        props?: Props
     }
 
     export function create<A extends StateObject, B extends StateObject, P>(parent: Ref, transformer: Transformer<A, B, P>, params?: P, options?: Options): Transform<A, B, P> {
@@ -40,11 +38,10 @@ export namespace Transform {
         return {
             parent,
             transformer,
-            params: params || {} as any,
             props: (options && options.props) || { },
             ref,
-            version: UUID.create22(),
-            cellState: { ...StateObjectCell.DefaultState, ...(options && options.cellState) }
+            params: params || {} as any,
+            version: UUID.create22()
         }
     }
 
@@ -52,10 +49,6 @@ export namespace Transform {
         return { ...t, params, version: UUID.create22() };
     }
 
-    export function withCellState<T>(t: Transform, state: Partial<StateObjectCell.State>): Transform {
-        return { ...t, cellState: { ...t.cellState, ...state } };
-    }
-
     export function createRoot(): Transform {
         return create(RootRef, Transformer.ROOT, {}, { ref: RootRef });
     }
@@ -66,8 +59,7 @@ export namespace Transform {
         params: any,
         props: Props,
         ref: string,
-        version: string,
-        cellState: StateObjectCell.State,
+        version: string
     }
 
     function _id(x: any) { return x; }
@@ -81,8 +73,7 @@ export namespace Transform {
             params: pToJson(t.params),
             props: t.props,
             ref: t.ref,
-            version: t.version,
-            cellState: t.cellState,
+            version: t.version
         };
     }
 
@@ -97,8 +88,7 @@ export namespace Transform {
             params: pFromJson(t.params),
             props: t.props,
             ref: t.ref as Ref,
-            version: t.version,
-            cellState: t.cellState,
+            version: t.version
         };
     }
 }

src/mol-state/tree/builder.ts

@@ -6,7 +6,7 @@
 
 import { StateTree } from './immutable';
 import { TransientTree } from './transient';
-import { StateObject } from '../object';
+import { StateObject, StateObjectCell } from '../object';
 import { Transform } from '../transform';
 import { Transformer } from '../transformer';
 
@@ -51,9 +51,9 @@ namespace StateTreeBuilder {
     export class To<A extends StateObject> implements StateTreeBuilder {
         get editInfo() { return this.state.editInfo; }
 
-        apply<T extends Transformer<A, any, any>>(tr: T, params?: Transformer.Params<T>, props?: Partial<Transform.Options>): To<Transformer.To<T>> {
-            const t = tr.apply(this.ref, params, props);
-            this.state.tree.add(t);
+        apply<T extends Transformer<A, any, any>>(tr: T, params?: Transformer.Params<T>, options?: Partial<Transform.Options>, initialCellState?: Partial<StateObjectCell.State>): To<Transformer.To<T>> {
+            const t = tr.apply(this.ref, params, options);
+            this.state.tree.add(t, initialCellState);
             this.editInfo.count++;
             this.editInfo.lastUpdate = t.ref;
             return new To(this.state, t.ref, this.root);

src/mol-state/tree/immutable.ts

@@ -8,6 +8,7 @@ import { Map as ImmutableMap, OrderedSet } from 'immutable';
 import { Transform } from '../transform';
 import { TransientTree } from './transient';
 import { StateTreeBuilder } from './builder';
+import { StateObjectCell } from 'mol-state/object';
 
 export { StateTree }
 
@@ -17,8 +18,10 @@ export { StateTree }
  */
 interface StateTree {
     readonly root: Transform,
-    readonly transforms: StateTree.Nodes,
+    readonly transforms: StateTree.Transforms,
     readonly children: StateTree.Children,
+    readonly cellStates: StateTree.CellStates,
+
     asTransient(): TransientTree,
     build(): StateTreeBuilder.Root
 }
@@ -40,8 +43,9 @@ namespace StateTree {
         get(ref: Ref): T
     }
 
-    export interface Nodes extends _Map<Transform> {}
+    export interface Transforms extends _Map<Transform> {}
     export interface Children extends _Map<ChildSet> { }
+    export interface CellStates extends _Map<StateObjectCell.State> { }
 
     class Impl implements StateTree {
         get root() { return this.transforms.get(Transform.RootRef)! }
@@ -54,7 +58,7 @@ namespace StateTree {
             return new StateTreeBuilder.Root(this);
         }
 
-        constructor(public transforms: StateTree.Nodes, public children: Children) {
+        constructor(public transforms: StateTree.Transforms, public children: Children, public cellStates: CellStates) {
         }
     }
 
@@ -63,11 +67,11 @@ namespace StateTree {
      */
     export function createEmpty(): StateTree {
         const root = Transform.createRoot();
-        return create(ImmutableMap([[root.ref, root]]), ImmutableMap([[root.ref, OrderedSet()]]));
+        return create(ImmutableMap([[root.ref, root]]), ImmutableMap([[root.ref, OrderedSet()]]), ImmutableMap([[root.ref, StateObjectCell.DefaultState]]));
     }
 
-    export function create(nodes: Nodes, children: Children): StateTree {
-        return new Impl(nodes, children);
+    export function create(nodes: Transforms, children: Children, cellStates: CellStates): StateTree {
+        return new Impl(nodes, children, cellStates);
     }
 
     type VisitorCtx = { tree: StateTree, state: any, f: (node: Transform, tree: StateTree, state: any) => boolean | undefined | void };
@@ -118,45 +122,45 @@ namespace StateTree {
         return doPostOrder<Transform[]>(tree, root, [], _subtree);
     }
 
-
-    // function _visitChildToJson(this: Ref[], ref: Ref) { this.push(ref); }
-    // interface ToJsonCtx { nodes: Transform.Serialized[] }
-    function _visitNodeToJson(node: Transform, tree: StateTree, ctx: Transform.Serialized[]) {
+    function _visitNodeToJson(node: Transform, tree: StateTree, ctx: [Transform.Serialized, StateObjectCell.State][]) {
         // const children: Ref[] = [];
         // tree.children.get(node.ref).forEach(_visitChildToJson as any, children);
-        ctx.push(Transform.toJSON(node));
+        ctx.push([Transform.toJSON(node), tree.cellStates.get(node.ref)]);
     }
 
     export interface Serialized {
-        /** Tree nodes serialized in pre-order */
-        nodes: Transform.Serialized[]
+        /** Transforms serialized in pre-order */
+        transforms: [Transform.Serialized, StateObjectCell.State][]
     }
 
     export function toJSON<T>(tree: StateTree): Serialized {
-        const nodes: Transform.Serialized[] = [];
-        doPreOrder(tree, tree.root, nodes, _visitNodeToJson);
-        return { nodes };
+        const transforms: [Transform.Serialized, StateObjectCell.State][] = [];
+        doPreOrder(tree, tree.root, transforms, _visitNodeToJson);
+        return { transforms };
     }
 
     export function fromJSON<T>(data: Serialized): StateTree {
         const nodes = ImmutableMap<Ref, Transform>().asMutable();
         const children = ImmutableMap<Ref, OrderedSet<Ref>>().asMutable();
+        const cellStates = ImmutableMap<Ref, StateObjectCell.State>().asMutable();
 
-        for (const s of data.nodes) {
-            const node = Transform.fromJSON(s);
-            nodes.set(node.ref, node);
+        for (const t of data.transforms) {
+            const transform = Transform.fromJSON(t[0]);
+            nodes.set(transform.ref, transform);
+            cellStates.set(transform.ref, t[1]);
 
-            if (!children.has(node.ref)) {
-                children.set(node.ref, OrderedSet<Ref>().asMutable());
+            if (!children.has(transform.ref)) {
+                children.set(transform.ref, OrderedSet<Ref>().asMutable());
             }
 
-            if (node.ref !== node.parent) children.get(node.parent).add(node.ref);
+            if (transform.ref !== transform.parent) children.get(transform.parent).add(transform.ref);
         }
 
-        for (const s of data.nodes) {
-            children.set(s.ref, children.get(s.ref).asImmutable());
+        for (const t of data.transforms) {
+            const ref = t[0].ref;
+            children.set(ref, children.get(ref).asImmutable());
         }
 
-        return create(nodes.asImmutable(), children.asImmutable());
+        return create(nodes.asImmutable(), children.asImmutable(), cellStates.asImmutable());
     }
 }

src/mol-state/tree/transient.ts

@@ -10,18 +10,19 @@ import { StateTree } from './immutable';
 import { StateTreeBuilder } from './builder';
 import { StateObjectCell } from 'mol-state/object';
 import { shallowEqual } from 'mol-util/object';
-import { UUID } from 'mol-util';
 
 export { TransientTree }
 
 class TransientTree implements StateTree {
     transforms = this.tree.transforms as ImmutableMap<Transform.Ref, Transform>;
     children = this.tree.children as ImmutableMap<Transform.Ref, OrderedSet<Transform.Ref>>;
+    cellStates = this.tree.cellStates as ImmutableMap<Transform.Ref, StateObjectCell.State>;
 
     private changedNodes = false;
     private changedChildren = false;
+    private changedStates = false;
+
     private _childMutations: Map<Transform.Ref, OrderedSet<Transform.Ref>> | undefined = void 0;
-    private _transformMutations: Map<Transform.Ref, Transform> | undefined = void 0;
 
     private get childMutations() {
         if (this._childMutations) return this._childMutations;
@@ -29,6 +30,24 @@ class TransientTree implements StateTree {
         return this._childMutations;
     }
 
+    private changeStates() {
+        if (this.changedStates) return;
+        this.changedStates = true;
+        this.cellStates = this.cellStates.asMutable();
+    }
+
+    private changeNodes() {
+        if (this.changedNodes) return;
+        this.changedNodes = true;
+        this.transforms = this.transforms.asMutable();
+    }
+
+    private changeChildren() {
+        if (this.changedChildren) return;
+        this.changedChildren = true;
+        this.children = this.children.asMutable();
+    }
+
     get root() { return this.transforms.get(Transform.RootRef)! }
 
     build(): StateTreeBuilder.Root {
@@ -40,10 +59,7 @@ class TransientTree implements StateTree {
     }
 
     private addChild(parent: Transform.Ref, child: Transform.Ref) {
-        if (!this.changedChildren) {
-            this.changedChildren = true;
-            this.children = this.children.asMutable();
-        }
+        this.changeChildren();
 
         if (this.childMutations.has(parent)) {
             this.childMutations.get(parent)!.add(child);
@@ -56,10 +72,7 @@ class TransientTree implements StateTree {
     }
 
     private removeChild(parent: Transform.Ref, child: Transform.Ref) {
-        if (!this.changedChildren) {
-            this.changedChildren = true;
-            this.children = this.children.asMutable();
-        }
+        this.changeChildren();
 
         if (this.childMutations.has(parent)) {
             this.childMutations.get(parent)!.remove(child);
@@ -74,16 +87,15 @@ class TransientTree implements StateTree {
     private clearRoot() {
         const parent = Transform.RootRef;
         if (this.children.get(parent).size === 0) return;
+
+        this.changeChildren();
+
         const set = OrderedSet<Transform.Ref>();
-        if (!this.changedChildren) {
-            this.changedChildren = true;
-            this.children = this.children.asMutable();
-        }
         this.children.set(parent, set);
         this.childMutations.set(parent, set);
     }
 
-    add(transform: Transform) {
+    add(transform: Transform, initialState?: Partial<StateObjectCell.State>) {
         const ref = transform.ref;
 
         if (this.transforms.has(transform.ref)) {
@@ -106,12 +118,18 @@ class TransientTree implements StateTree {
             this.children.set(transform.ref, OrderedSet());
         }
 
-        if (!this.changedNodes) {
-            this.changedNodes = true;
-            this.transforms = this.transforms.asMutable();
+        this.changeNodes();
+        this.transforms.set(ref, transform);
+
+        if (!this.cellStates.has(ref)) {
+            this.changeStates();
+            if (StateObjectCell.isStateChange(StateObjectCell.DefaultState, initialState)) {
+                this.cellStates.set(ref, { ...StateObjectCell.DefaultState, ...initialState });
+            } else {
+                this.cellStates.set(ref, StateObjectCell.DefaultState);
+            }
         }
 
-        this.transforms.set(ref, transform);
         return this;
     }
 
@@ -129,48 +147,25 @@ class TransientTree implements StateTree {
             }
         }
 
-        if (this._transformMutations && this._transformMutations.has(transform.ref)) {
-            const mutated = this._transformMutations.get(transform.ref)!;
-            (mutated.params as any) = params;
-            (mutated.version as UUID) = UUID.create22();
-        } else {
-            this.set(Transform.withParams(transform, params));
+        if (!this.changedNodes) {
+            this.changedNodes = true;
+            this.transforms = this.transforms.asMutable();
         }
 
+        this.transforms.set(transform.ref, Transform.withParams(transform, params));
         return true;
     }
 
-    setCellState(ref: Transform.Ref, state: Partial<StateObjectCell.State>) {
+    updateCellState(ref: Transform.Ref, state: Partial<StateObjectCell.State>) {
         ensurePresent(this.transforms, ref);
 
-        if (this._transformMutations && this._transformMutations.has(ref)) {
-            const transform = this._transformMutations.get(ref)!;
-            const old = transform.cellState;
-            (transform.cellState as StateObjectCell.State) = { ...old, ...state };
-            return transform;
-        } else {
-            const transform = this.transforms.get(ref);
-            const newT = Transform.withCellState(transform, state);
-            this.set(newT);
-            return newT;
-        }
-    }
+        const old = this.cellStates.get(ref);
+        if (!StateObjectCell.isStateChange(old, state)) return false;
 
-    private set(transform: Transform) {
-        ensurePresent(this.transforms, transform.ref);
+        this.changeStates();
+        this.cellStates.set(ref, { ...old, ...state });
 
-        if (!this.changedNodes) {
-            this.changedNodes = true;
-            this.transforms = this.transforms.asMutable();
-        }
-
-        if (!this._transformMutations) {
-            this._transformMutations = new Map();
-        }
-        this._transformMutations.set(transform.ref, transform);
-
-        this.transforms.set(transform.ref, transform);
-        return this;
+        return true;
     }
 
     remove(ref: Transform.Ref): Transform[] {
@@ -187,14 +182,14 @@ class TransientTree implements StateTree {
             this.removeChild(node.parent, node.ref);
         }
 
-        if (!this.changedNodes) {
-            this.changedNodes = true;
-            this.transforms = this.transforms.asMutable();
-        }
+        this.changeNodes();
+        this.changeChildren();
+        this.changeStates();
 
         for (const n of st) {
             this.transforms.delete(n.ref);
             this.children.delete(n.ref);
+            this.cellStates.delete(n.ref);
             if (this._childMutations) this._childMutations.delete(n.ref);
         }
 
@@ -202,11 +197,12 @@ class TransientTree implements StateTree {
     }
 
     asImmutable() {
-        if (!this.changedNodes && !this.changedChildren && !this._childMutations) return this.tree;
+        if (!this.changedNodes && !this.changedChildren && !this.changedStates && !this._childMutations) return this.tree;
         if (this._childMutations) this._childMutations.forEach(fixChildMutations, this.children);
         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.changedStates ? this.cellStates.asImmutable() : this.cellStates);
     }
 
     constructor(private tree: StateTree) {
@@ -224,7 +220,7 @@ function parentNotPresent(ref: Transform.Ref) {
     throw new Error(`Parent '${ref}' must be present in the tree.`);
 }
 
-function ensurePresent(nodes: StateTree.Nodes, ref: Transform.Ref) {
+function ensurePresent(nodes: StateTree.Transforms, ref: Transform.Ref) {
     if (!nodes.has(ref)) {
         throw new Error(`Node '${ref}' is not present in the tree.`);
     }