wip mol-state refactoring

src/mol-plugin/state.ts

@@ -99,7 +99,7 @@ class PluginState {
     constructor(private plugin: import('./context').PluginContext) {
         this.snapshots = new PluginStateSnapshotManager(plugin);
         this.dataState = State.create(new SO.Root({ }), { globalContext: plugin });
-        this.behaviorState = State.create(new PluginBehavior.Root({ }), { globalContext: plugin, rootProps: { isLocked: true } });
+        this.behaviorState = State.create(new PluginBehavior.Root({ }), { globalContext: plugin, rootState: { isLocked: true } });
 
         this.dataState.behaviors.currentObject.subscribe(o => {
             if (this.behavior.kind.value === 'data') this.behavior.currentObject.next(o);

src/mol-state/manager.ts

@@ -1,7 +0,0 @@
-/**
- * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
- *
- * @author David Sehnal <david.sehnal@gmail.com>
- */
-
-// TODO manage snapshots etc

src/mol-state/object.ts

@@ -19,7 +19,7 @@ interface StateObject<D = any, T extends StateObject.Type = StateObject.Type<any
     readonly label: string,
     readonly description?: string,
     // assigned by reconciler to be StateTransform.props.tag
-    readonly tag?: string
+    readonly tags?: string[]
 }
 
 namespace StateObject {
@@ -62,6 +62,7 @@ interface StateObjectCell<T extends StateObject = StateObject, F extends StateTr
     sourceRef: StateTransform.Ref | undefined,
 
     status: StateObjectCell.Status,
+    state: StateTransform.State,
 
     params: {
         definition: ParamDefinition.Params,
@@ -79,24 +80,6 @@ namespace StateObjectCell {
 
     export type Obj<C extends StateObjectCell> = C extends StateObjectCell<infer T> ? T : never
     export type Transform<C extends StateObjectCell> = C extends StateObjectCell<any, infer T> ? T : never
-
-    export interface State {
-        isHidden: boolean,
-        isCollapsed: boolean
-    }
-
-    export const DefaultState: State = { isHidden: false, isCollapsed: false };
-
-    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;
-    }
 }
 
 // TODO: improve the API?

src/mol-state/state.ts

@@ -33,7 +33,7 @@ class State {
     readonly globalContext: unknown = void 0;
     readonly events = {
         cell: {
-            stateUpdated: this.ev<State.ObjectEvent & { cellState: StateObjectCell.State }>(),
+            stateUpdated: this.ev<State.ObjectEvent & { cell: StateObjectCell }>(),
             created: this.ev<State.ObjectEvent & { cell: StateObjectCell }>(),
             removed: this.ev<State.ObjectEvent & { parent: StateTransform.Ref }>(),
         },
@@ -55,7 +55,6 @@ class State {
 
     get tree(): StateTree { return this._tree; }
     get transforms() { return (this._tree as StateTree).transforms; }
-    get cellStates() { return (this._tree as StateTree).cellStates; }
     get current() { return this.behaviors.currentObject.value.ref; }
 
     build() { return new StateBuilder.Root(this.tree, this); }
@@ -64,6 +63,7 @@ class State {
     private spine = new StateTreeSpine.Impl(this.cells);
 
     getSnapshot(): State.Snapshot {
+        this.cells.forEach(c => this._tree.updateState(c.transform.ref, c.state));
         return { tree: StateTree.toJSON(this._tree) };
     }
 
@@ -76,13 +76,15 @@ class State {
         this.behaviors.currentObject.next({ state: this, ref });
     }
 
-    updateCellState(ref: StateTransform.Ref, stateOrProvider: ((old: StateObjectCell.State) => Partial<StateObjectCell.State>) | Partial<StateObjectCell.State>) {
-        const update = typeof stateOrProvider === 'function'
-            ? stateOrProvider(this.tree.cellStates.get(ref))
-            : stateOrProvider;
+    updateCellState(ref: StateTransform.Ref, stateOrProvider: ((old: StateTransform.State) => Partial<StateTransform.State>) | Partial<StateTransform.State>) {
+        const cell = this.cells.get(ref);
+        if (!cell) return;
 
-        if (this._tree.updateCellState(ref, update)) {
-            this.events.cell.stateUpdated.next({ state: this, ref, cellState: this.tree.cellStates.get(ref) });
+        const update = typeof stateOrProvider === 'function' ? stateOrProvider(cell.state) : stateOrProvider;
+
+        if (StateTransform.assignState(cell.state, update)) {
+            // this._tree.updateCellState(ref, update)) {
+            this.events.cell.stateUpdated.next({ state: this, ref, cell });
         }
     }
 
@@ -165,10 +167,6 @@ class State {
 
             if (updated) this.events.changed.next();
             this.events.isUpdating.next(false);
-
-            for (const ref of ctx.stateChanges) {
-                this.events.cell.stateUpdated.next({ state: this, ref, cellState: this.tree.cellStates.get(ref) });
-            }
         }
     }
 
@@ -189,7 +187,6 @@ class State {
             spine: this.spine,
 
             results: [],
-            stateChanges: [],
 
             options: { ...StateUpdateDefaultOptions, ...options },
 
@@ -203,8 +200,8 @@ class State {
         return ctx;
     }
 
-    constructor(rootObject: StateObject, params?: { globalContext?: unknown, rootProps?: StateTransform.Props }) {
-        this._tree = StateTree.createEmpty(StateTransform.createRoot(params && params.rootProps)).asTransient();
+    constructor(rootObject: StateObject, params?: { globalContext?: unknown, rootState?: StateTransform.State }) {
+        this._tree = StateTree.createEmpty(StateTransform.createRoot(params && params.rootState)).asTransient();
         const tree = this._tree;
         const root = tree.root;
 
@@ -213,6 +210,7 @@ class State {
             sourceRef: void 0,
             obj: rootObject,
             status: 'ok',
+            state: { ...root.state },
             errorText: void 0,
             params: {
                 definition: {},
@@ -245,7 +243,7 @@ namespace State {
         doNotUpdateCurrent: boolean
     }
 
-    export function create(rootObject: StateObject, params?: { globalContext?: unknown, rootProps?: StateTransform.Props }) {
+    export function create(rootObject: StateObject, params?: { globalContext?: unknown, rootState?: StateTransform.State }) {
         return new State(rootObject, params);
     }
 }
@@ -271,7 +269,6 @@ interface UpdateContext {
     spine: StateTreeSpine.Impl,
 
     results: UpdateNodeResult[],
-    stateChanges: StateTransform.Ref[],
 
     // suppress timing messages
     options: State.UpdateOptions,
@@ -319,12 +316,6 @@ async function update(ctx: UpdateContext) {
         roots = findUpdateRoots(ctx.cells, ctx.tree);
     }
 
-    let newCellStates: StateTree.CellStates;
-    if (!ctx.editInfo) {
-        newCellStates = ctx.tree.cellStatesSnapshot();
-        syncOldStates(ctx);
-    }
-
     // Init empty cells where not present
     // this is done in "pre order", meaning that "parents" will be created 1st.
     const addedCells = initCells(ctx, roots);
@@ -353,7 +344,7 @@ async function update(ctx: UpdateContext) {
 
     // Sync cell states
     if (!ctx.editInfo) {
-        syncNewStates(ctx, newCellStates!);
+        syncNewStates(ctx);
     }
 
     let newCurrent: StateTransform.Ref | undefined = ctx.newCurrent;
@@ -363,7 +354,7 @@ async function update(ctx: UpdateContext) {
             ctx.parent.events.object.created.next({ state: ctx.parent, ref: update.ref, obj: update.obj! });
             if (!ctx.newCurrent) {
                 const transform = ctx.tree.transforms.get(update.ref);
-                if (!(transform.props && transform.props.isGhost) && update.obj !== StateObject.Null) newCurrent = update.ref;
+                if (!transform.state.isGhost && update.obj !== StateObject.Null) newCurrent = update.ref;
             }
         } else if (update.action === 'updated') {
             ctx.parent.events.object.updated.next({ state: ctx.parent, ref: update.ref, action: 'in-place', obj: update.obj });
@@ -415,25 +406,14 @@ function findDeletes(ctx: UpdateContext): Ref[] {
     return deleteCtx.deletes;
 }
 
-function syncOldStatesVisitor(n: StateTransform, tree: StateTree, oldState: StateTree.CellStates) {
-    if (oldState.has(n.ref)) {
-        (tree as TransientTree).updateCellState(n.ref, oldState.get(n.ref));
-    }
-}
-function syncOldStates(ctx: UpdateContext) {
-    StateTree.doPreOrder(ctx.tree, ctx.tree.root, ctx.oldTree.cellStates, syncOldStatesVisitor);
+function syncNewStatesVisitor(n: StateTransform, tree: StateTree, ctx: UpdateContext) {
+    const cell = ctx.cells.get(n.ref);
+    if (!cell || !StateTransform.assignState(cell.state, n.state)) return;
+    ctx.parent.events.cell.stateUpdated.next({ state: ctx.parent, ref: n.ref, cell });
 }
 
-function syncNewStatesVisitor(n: StateTransform, tree: StateTree, ctx: { newState: StateTree.CellStates, changes: StateTransform.Ref[] }) {
-    if (ctx.newState.has(n.ref)) {
-        const changed = (tree as TransientTree).updateCellState(n.ref, ctx.newState.get(n.ref));
-        if (changed) {
-            ctx.changes.push(n.ref);
-        }
-    }
-}
-function syncNewStates(ctx: UpdateContext, newState: StateTree.CellStates) {
-    StateTree.doPreOrder(ctx.tree, ctx.tree.root, { newState, changes: ctx.stateChanges }, syncNewStatesVisitor);
+function syncNewStates(ctx: UpdateContext) {
+    StateTree.doPreOrder(ctx.tree, ctx.tree.root, ctx, syncNewStatesVisitor);
 }
 
 function setCellStatus(ctx: UpdateContext, ref: Ref, status: StateObjectCell.Status, errorText?: string) {
@@ -441,7 +421,7 @@ function setCellStatus(ctx: UpdateContext, ref: Ref, status: StateObjectCell.Sta
     const changed = cell.status !== status;
     cell.status = status;
     cell.errorText = errorText;
-    if (changed) ctx.parent.events.cell.stateUpdated.next({ state: ctx.parent, ref, cellState: ctx.tree.cellStates.get(ref) });
+    if (changed) ctx.parent.events.cell.stateUpdated.next({ state: ctx.parent, ref, cell });
 }
 
 function initCellStatusVisitor(t: StateTransform, _: any, ctx: UpdateContext) {
@@ -465,6 +445,7 @@ function initCellsVisitor(transform: StateTransform, _: any, { ctx, added }: Ini
         transform,
         sourceRef: void 0,
         status: 'pending',
+        state: { ...transform.state },
         errorText: void 0,
         params: void 0,
         cache: void 0
@@ -505,7 +486,7 @@ function _findNewCurrent(tree: StateTree, ref: Ref, deletes: Set<Ref>, cells: Ma
         }
 
         const t = tree.transforms.get(s.value);
-        if (t.props && t.props.isGhost) continue;
+        if (t.state.isGhost) continue;
         if (s.value === ref) {
             seenRef = true;
             if (!deletes.has(ref)) prevCandidate = ref;
@@ -671,7 +652,7 @@ async function updateNode(ctx: UpdateContext, currentRef: Ref): Promise<UpdateNo
 
 function updateTag(obj: StateObject | undefined, transform: StateTransform) {
     if (!obj || obj === StateObject.Null) return;
-    (obj.tag as string | undefined) = transform.props.tag;
+    (obj.tags as string[] | undefined) = transform.tags;
 }
 
 function runTask<T>(t: T | Task<T>, ctx: RuntimeContext) {

src/mol-state/state/builder.ts

@@ -36,7 +36,7 @@ namespace StateBuilder {
         | { kind: 'add', transform: StateTransform }
         | { kind: 'update', ref: string, params: any }
         | { kind: 'delete', ref: string }
-        | { kind: 'insert', ref: string, transform: StateTransform, initialCellState?: Partial<StateObjectCell.State> }
+        | { kind: 'insert', ref: string, transform: StateTransform }
 
     function buildTree(state: BuildState) {
         if (!state.state || state.state.tree === state.editInfo.sourceTree) {
@@ -52,7 +52,7 @@ namespace StateBuilder {
                 case 'delete': tree.remove(a.ref); break;
                 case 'insert': {
                     const children = tree.children.get(a.ref).toArray();
-                    tree.add(a.transform, a.initialCellState);
+                    tree.add(a.transform);
                     for (const c of children) {
                         tree.changeParent(c, a.transform.ref);
                     }
@@ -89,7 +89,7 @@ namespace StateBuilder {
             this.state.actions.push({ kind: 'delete', ref });
             return this;
         }
-        getTree(): StateTree { return buildTree(this.state); } //this.state.tree.asImmutable(); }
+        getTree(): StateTree { return buildTree(this.state); }
         constructor(tree: StateTree, state?: State) { this.state = { state, tree: tree.asTransient(), actions: [], editInfo: { sourceTree: tree, count: 0, lastUpdate: void 0 } } }
     }
 
@@ -102,9 +102,9 @@ namespace StateBuilder {
          * Apply the transformed to the parent node
          * If no params are specified (params <- undefined), default params are lazily resolved.
          */
-        apply<T extends StateTransformer<A, any, any>>(tr: T, params?: StateTransformer.Params<T>, options?: Partial<StateTransform.Options>, initialCellState?: Partial<StateObjectCell.State>): To<StateTransformer.To<T>> {
+        apply<T extends StateTransformer<A, any, any>>(tr: T, params?: StateTransformer.Params<T>, options?: Partial<StateTransform.Options>): To<StateTransformer.To<T>> {
             const t = tr.apply(this.ref, params, options);
-            this.state.tree.add(t, initialCellState);
+            this.state.tree.add(t);
             this.editInfo.count++;
             this.editInfo.lastUpdate = t.ref;
 
@@ -116,20 +116,20 @@ namespace StateBuilder {
         /**
          * A helper to greate a group-like state object and keep the current type.
          */
-        group<T extends StateTransformer<A, any, any>>(tr: T, params?: StateTransformer.Params<T>, options?: Partial<StateTransform.Options>, initialCellState?: Partial<StateObjectCell.State>): To<A> {
-            return this.apply(tr, params, options, initialCellState) as To<A>;
+        group<T extends StateTransformer<A, any, any>>(tr: T, params?: StateTransformer.Params<T>, options?: Partial<StateTransform.Options>): To<A> {
+            return this.apply(tr, params, options) as To<A>;
         }
 
         /**
          * Inserts a new transform that does not change the object type and move the original children to it.
          */
-        insert<T extends StateTransformer<A, A, any>>(tr: T, params?: StateTransformer.Params<T>, options?: Partial<StateTransform.Options>, initialCellState?: Partial<StateObjectCell.State>): To<StateTransformer.To<T>> {
+        insert<T extends StateTransformer<A, A, any>>(tr: T, params?: StateTransformer.Params<T>, options?: Partial<StateTransform.Options>): To<StateTransformer.To<T>> {
             // cache the children
             const children = this.state.tree.children.get(this.ref).toArray();
 
             // add the new node
             const t = tr.apply(this.ref, params, options);
-            this.state.tree.add(t, initialCellState);
+            this.state.tree.add(t);
 
             // move the original children to the new node
             for (const c of children) {
@@ -139,7 +139,7 @@ namespace StateBuilder {
             this.editInfo.count++;
             this.editInfo.lastUpdate = t.ref;
 
-            this.state.actions.push({ kind: 'insert', ref: this.ref, transform: t, initialCellState });
+            this.state.actions.push({ kind: 'insert', ref: this.ref, transform: t });
 
             return new To(this.state, t.ref, this.root);
         }
@@ -193,7 +193,7 @@ namespace StateBuilder {
         toRoot<A extends StateObject>() { return this.root.toRoot<A>(); }
         delete(ref: StateTransform.Ref) { return this.root.delete(ref); }
 
-        getTree(): StateTree { return buildTree(this.state); } //this.state.tree.asImmutable(); }
+        getTree(): StateTree { return buildTree(this.state); }
 
         constructor(private state: BuildState, ref: StateTransform.Ref, private root: Root) {
             this.ref = ref;

src/mol-state/state/selection.ts

@@ -268,8 +268,12 @@ namespace StateSelection {
     }
 
     function _findUniqueTagsInSubtree(n: StateTransform, _: any, s: { refs: { [name: string]: StateTransform.Ref }, tags: Set<string> }) {
-        if (n.props.tag && s.tags.has(n.props.tag)) {
-            s.refs[n.props.tag] = n.ref;
+        if (n.tags) {
+            for (const t of n.tags) {
+                if (!s.tags.has(t)) continue;
+                s.refs[t] = n.ref;
+                break;
+            }
         }
         return true;
     }
@@ -279,7 +283,7 @@ namespace StateSelection {
     }
 
     function _findTagInSubtree(n: StateTransform, _: any, s: { ref: string | undefined, tag: string }) {
-        if (n.props.tag === s.tag) {
+        if (n.tags && n.tags.indexOf(s.tag) >= 0) {
             s.ref = n.ref;
             return false;
         }

src/mol-state/transform.ts

@@ -12,7 +12,8 @@ export { Transform as StateTransform }
 interface Transform<T extends StateTransformer = StateTransformer> {
     readonly parent: Transform.Ref,
     readonly transformer: T,
-    readonly props: Transform.Props,
+    readonly state: Transform.State,
+    readonly tags?: string[],
     readonly ref: Transform.Ref,
     readonly params?: StateTransformer.Params<T>,
     readonly version: string
@@ -24,24 +25,60 @@ namespace Transform {
 
     export const RootRef = '-=root=-' as Ref;
 
-    export interface Props {
-        tag?: string
+    export interface State {
+        // is the cell shown in the UI
         isGhost?: boolean,
-        // determine if the corresponding cell can be deleted by the user.
-        isLocked?: boolean
+        // can the corresponding be deleted by the user.
+        isLocked?: boolean,
+        // is the representation associated with the cell hidden
+        isHidden?: boolean,
+        // is the tree node collapsed?
+        isCollapsed?: boolean
+    }
+
+    export function areStatesEqual(a: State, b: State) {
+        return !!a.isHidden !== !!b.isHidden || !!a.isCollapsed !== !!b.isCollapsed
+            || !!a.isGhost !== !!b.isGhost || !!a.isLocked !== !!b.isLocked;
+    }
+
+    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;
+        if (typeof b.isGhost !== 'undefined' && a.isGhost !== b.isGhost) return true;
+        if (typeof b.isLocked !== 'undefined' && a.isLocked !== b.isLocked) return true;
+        return false;
+    }
+
+    export function assignState(a: State, b?: Partial<State>): boolean {
+        if (!b) return false;
+        let changed = false;
+        for (const k of Object.keys(b)) {
+            const s = (b as any)[k], t = (a as any)[k];
+            if (!!s === !!t) continue;
+            changed = true;
+            (a as any)[k] = s;
+        }
+        return changed;
     }
 
     export interface Options {
         ref?: string,
-        props?: Props
+        tags?: string | string[],
+        state?: State
     }
 
     export function create<T extends StateTransformer>(parent: Ref, transformer: T, params?: StateTransformer.Params<T>, options?: Options): Transform<T> {
         const ref = options && options.ref ? options.ref : UUID.create22() as string as Ref;
+        let tags: string[] | undefined = void 0;
+        if (options && options.tags) {
+            tags = typeof options.tags === 'string' ? [options.tags] : options.tags;
+        }
         return {
             parent,
             transformer,
-            props: (options && options.props) || { },
+            state: (options && options.state) || { },
+            tags,
             ref,
             params,
             version: UUID.create22()
@@ -52,23 +89,25 @@ namespace Transform {
         return { ...t, params, version: UUID.create22() };
     }
 
-    export function withParent(t: Transform, parent: Ref): Transform {
-        return { ...t, parent, version: UUID.create22() };
+    export function withState(t: Transform, state?: Partial<State>): Transform {
+        if (!state) return t;
+        return { ...t, state: { ...t.state, ...state } };
     }
 
-    export function withNewVersion(t: Transform): Transform {
-        return { ...t, version: UUID.create22() };
+    export function withParent(t: Transform, parent: Ref): Transform {
+        return { ...t, parent, version: UUID.create22() };
     }
 
-    export function createRoot(props?: Props): Transform {
-        return create(RootRef, StateTransformer.ROOT, {}, { ref: RootRef, props });
+    export function createRoot(state?: State): Transform {
+        return create(RootRef, StateTransformer.ROOT, {}, { ref: RootRef, state });
     }
 
     export interface Serialized {
         parent: string,
         transformer: string,
         params: any,
-        props: Props,
+        state?: State,
+        tags?: string[],
         ref: string,
         version: string
     }
@@ -78,11 +117,19 @@ namespace Transform {
         const pToJson = t.transformer.definition.customSerialization
             ? t.transformer.definition.customSerialization.toJSON
             : _id;
+        let state: any = void 0;
+        for (const k of Object.keys(t.state)) {
+            const s = (t.state as any)[k];
+            if (!s) continue;
+            if (!state) state = { };
+            state[k] = true;
+        }
         return {
             parent: t.parent,
             transformer: t.transformer.id,
             params: t.params ? pToJson(t.params) : void 0,
-            props: t.props,
+            state,
+            tags: t.tags,
             ref: t.ref,
             version: t.version
         };
@@ -97,7 +144,8 @@ namespace Transform {
             parent: t.parent as Ref,
             transformer,
             params: t.params ? pFromJson(t.params) : void 0,
-            props: t.props,
+            state: t.state || { },
+            tags: t.tags,
             ref: t.ref as Ref,
             version: t.version
         };

src/mol-state/tree/immutable.ts

@@ -7,7 +7,6 @@
 import { Map as ImmutableMap, OrderedSet } from 'immutable';
 import { StateTransform } from '../transform';
 import { TransientTree } from './transient';
-import { StateObjectCell } from 'mol-state/object';
 
 export { StateTree }
 
@@ -19,7 +18,6 @@ interface StateTree {
     readonly root: StateTransform,
     readonly transforms: StateTree.Transforms,
     readonly children: StateTree.Children,
-    readonly cellStates: StateTree.CellStates,
 
     asTransient(): TransientTree
 }
@@ -43,7 +41,6 @@ namespace StateTree {
 
     export interface Transforms extends _Map<StateTransform> {}
     export interface Children extends _Map<ChildSet> { }
-    export interface CellStates extends _Map<StateObjectCell.State> { }
 
     class Impl implements StateTree {
         get root() { return this.transforms.get(StateTransform.RootRef)! }
@@ -52,7 +49,7 @@ namespace StateTree {
             return new TransientTree(this);
         }
 
-        constructor(public transforms: StateTree.Transforms, public children: Children, public cellStates: CellStates) {
+        constructor(public transforms: StateTree.Transforms, public children: Children) {
         }
     }
 
@@ -61,11 +58,11 @@ namespace StateTree {
      */
     export function createEmpty(customRoot?: StateTransform): StateTree {
         const root = customRoot || StateTransform.createRoot();
-        return create(ImmutableMap([[root.ref, root]]), ImmutableMap([[root.ref, OrderedSet()]]), ImmutableMap([[root.ref, StateObjectCell.DefaultState]]));
+        return create(ImmutableMap([[root.ref, root]]), ImmutableMap([[root.ref, OrderedSet()]]));
     }
 
-    export function create(nodes: Transforms, children: Children, cellStates: CellStates): StateTree {
-        return new Impl(nodes, children, cellStates);
+    export function create(nodes: Transforms, children: Children): StateTree {
+        return new Impl(nodes, children);
     }
 
     type VisitorCtx = { tree: StateTree, state: any, f: (node: StateTransform, tree: StateTree, state: any) => boolean | undefined | void };
@@ -116,19 +113,19 @@ namespace StateTree {
         return doPostOrder<StateTransform[]>(tree, root, [], _subtree);
     }
 
-    function _visitNodeToJson(node: StateTransform, tree: StateTree, ctx: [StateTransform.Serialized, StateObjectCell.State][]) {
+    function _visitNodeToJson(node: StateTransform, tree: StateTree, ctx: StateTransform.Serialized[]) {
         // const children: Ref[] = [];
         // tree.children.get(node.ref).forEach(_visitChildToJson as any, children);
-        ctx.push([StateTransform.toJSON(node), tree.cellStates.get(node.ref)]);
+        ctx.push(StateTransform.toJSON(node));
     }
 
     export interface Serialized {
         /** Transforms serialized in pre-order */
-        transforms: [StateTransform.Serialized, StateObjectCell.State][]
+        transforms: StateTransform.Serialized[]
     }
 
     export function toJSON(tree: StateTree): Serialized {
-        const transforms: [StateTransform.Serialized, StateObjectCell.State][] = [];
+        const transforms: StateTransform.Serialized[] = [];
         doPreOrder(tree, tree.root, transforms, _visitNodeToJson);
         return { transforms };
     }
@@ -136,12 +133,10 @@ namespace StateTree {
     export function fromJSON(data: Serialized): StateTree {
         const nodes = ImmutableMap<Ref, StateTransform>().asMutable();
         const children = ImmutableMap<Ref, OrderedSet<Ref>>().asMutable();
-        const cellStates = ImmutableMap<Ref, StateObjectCell.State>().asMutable();
 
         for (const t of data.transforms) {
-            const transform = StateTransform.fromJSON(t[0]);
+            const transform = StateTransform.fromJSON(t);
             nodes.set(transform.ref, transform);
-            cellStates.set(transform.ref, t[1]);
 
             if (!children.has(transform.ref)) {
                 children.set(transform.ref, OrderedSet<Ref>().asMutable());
@@ -151,19 +146,18 @@ namespace StateTree {
         }
 
         for (const t of data.transforms) {
-            const ref = t[0].ref;
+            const ref = t.ref;
             children.set(ref, children.get(ref).asImmutable());
         }
 
-        return create(nodes.asImmutable(), children.asImmutable(), cellStates.asImmutable());
+        return create(nodes.asImmutable(), children.asImmutable());
     }
 
     export function dump(tree: StateTree) {
         console.log({
             tr: (tree.transforms as ImmutableMap<any, any>).keySeq().toArray(),
             tr1: (tree.transforms as ImmutableMap<any, any>).valueSeq().toArray().map(t => t.ref),
-            ch: (tree.children as ImmutableMap<any, any>).keySeq().toArray(),
-            cs: (tree.cellStates as ImmutableMap<any, any>).keySeq().toArray()
+            ch: (tree.children as ImmutableMap<any, any>).keySeq().toArray()
         });
     }
 }

src/mol-state/tree/transient.ts

@@ -7,7 +7,6 @@
 import { Map as ImmutableMap, OrderedSet } from 'immutable';
 import { StateTransform } from '../transform';
 import { StateTree } from './immutable';
-import { StateObjectCell } from 'mol-state/object';
 import { shallowEqual } from 'mol-util/object';
 
 export { TransientTree }
@@ -15,11 +14,9 @@ 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>>;
-    cellStates = this.tree.cellStates as ImmutableMap<StateTransform.Ref, StateObjectCell.State>;
 
     private changedNodes = false;
     private changedChildren = false;
-    private changedStates = false;
 
     private _childMutations: Map<StateTransform.Ref, OrderedSet<StateTransform.Ref>> | undefined = void 0;
 
@@ -29,12 +26,6 @@ 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;
@@ -49,10 +40,6 @@ class TransientTree implements StateTree {
 
     get root() { return this.transforms.get(StateTransform.RootRef)! }
 
-    cellStatesSnapshot() {
-        return this.cellStates.asImmutable();
-    }
-
     asTransient() {
         return this.asImmutable().asTransient();
     }
@@ -104,15 +91,7 @@ class TransientTree implements StateTree {
         this.transforms.set(ref, StateTransform.withParent(old, newParent));
     }
 
-    updateVersion(ref: StateTransform.Ref) {
-        ensurePresent(this.transforms, ref);
-
-        const t = this.transforms.get(ref);
-        this.changeNodes();
-        this.transforms.set(ref, StateTransform.withNewVersion(t));
-    }
-
-    add(transform: StateTransform, initialState?: Partial<StateObjectCell.State>) {
+    add(transform: StateTransform) {
         const ref = transform.ref;
 
         if (this.transforms.has(transform.ref)) {
@@ -138,15 +117,6 @@ class TransientTree implements StateTree {
         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);
-            }
-        }
-
         return this;
     }
 
@@ -169,14 +139,15 @@ class TransientTree implements StateTree {
         return true;
     }
 
-    updateCellState(ref: StateTransform.Ref, state: Partial<StateObjectCell.State>) {
+    updateState(ref: StateTransform.Ref, state?: Partial<StateTransform.State>) {
         ensurePresent(this.transforms, ref);
 
-        const old = this.cellStates.get(ref);
-        if (!StateObjectCell.isStateChange(old, state)) return false;
+        const old = this.transforms.get(ref);
+        if (!StateTransform.isStateChange(old.state, state)) return false;
 
-        this.changeStates();
-        this.cellStates.set(ref, { ...old, ...state });
+        this.changeNodes();
+        // TODO: cache these changes?
+        this.transforms.set(ref, StateTransform.withState(old, state));
 
         return true;
     }
@@ -197,12 +168,10 @@ class TransientTree implements StateTree {
 
         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);
         }
 
@@ -210,12 +179,11 @@ class TransientTree implements StateTree {
     }
 
     asImmutable() {
-        if (!this.changedNodes && !this.changedChildren && !this.changedStates && !this._childMutations) return this.tree;
+        if (!this.changedNodes && !this.changedChildren && !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.changedStates ? this.cellStates.asImmutable() : this.cellStates);
+            this.changedChildren ? this.children.asImmutable() : this.children);
     }
 
     constructor(private tree: StateTree) {