mol-state: fixes to transform state handling

src/mol-plugin/behavior/dynamic/labels.ts

@@ -118,7 +118,7 @@ export const SceneLabels = PluginBehavior.create<SceneLabelsProps>({
             for (const s of structures) {
                 const rootStructure = getRootStructure(s, state)
                 if (!rootStructure || !SO.Molecule.Structure.is(rootStructure.obj)) continue
-                if (!s.transform.state.isHidden) {
+                if (!s.state.isHidden) {
                     rootStructures.add(rootStructure.obj)
                 }
             }

src/mol-plugin/behavior/static/representation.ts

@@ -92,5 +92,5 @@ export function UpdateRepresentationVisibility(ctx: PluginContext) {
 }
 
 function updateVisibility(cell: StateObjectCell, r: Representation<any>) {
-    r.setState({ visible: !cell.transform.state.isHidden });
+    r.setState({ visible: !cell.state.isHidden });
 }
\ No newline at end of file

src/mol-plugin/behavior/static/state.ts

@@ -87,7 +87,7 @@ export function ToggleExpanded(ctx: PluginContext) {
 }
 
 export function ToggleVisibility(ctx: PluginContext) {
-    PluginCommands.State.ToggleVisibility.subscribe(ctx, ({ state, ref }) => setVisibility(state, ref, !state.cells.get(ref)!.transform.state.isHidden));
+    PluginCommands.State.ToggleVisibility.subscribe(ctx, ({ state, ref }) => setVisibility(state, ref, !state.cells.get(ref)!.state.isHidden));
 }
 
 function setVisibility(state: State, root: StateTransform.Ref, value: boolean) {

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

@@ -54,8 +54,8 @@ class StateTreeNode extends PluginUIComponent<{ cell: StateObjectCell, depth: nu
         this.subscribe(this.plugin.events.state.cell.stateUpdated, e => {
             if (this.props.cell === e.cell && this.is(e) && e.state.cells.has(this.ref)) {
                 this.forceUpdate();
-                // if (!!this.props.cell.transform.state.isCollapsed !== this.state.isCollapsed) {
-                //     this.setState({ isCollapsed: !!e.cell.transform.state.isCollapsed });
+                // if (!!this.props.cell.state.isCollapsed !== this.state.isCollapsed) {
+                //     this.setState({ isCollapsed: !!e.cell.state.isCollapsed });
                 // }
             }
         });
@@ -74,12 +74,12 @@ class StateTreeNode extends PluginUIComponent<{ cell: StateObjectCell, depth: nu
     }
 
     state = {
-        isCollapsed: !!this.props.cell.transform.state.isCollapsed
+        isCollapsed: !!this.props.cell.state.isCollapsed
     }
 
     static getDerivedStateFromProps(props: _Props<StateTreeNode>, state: _State<StateTreeNode>): _State<StateTreeNode> | null {
-        if (!!props.cell.transform.state.isCollapsed === state.isCollapsed) return null;
-        return { isCollapsed: !!props.cell.transform.state.isCollapsed };
+        if (!!props.cell.state.isCollapsed === state.isCollapsed) return null;
+        return { isCollapsed: !!props.cell.state.isCollapsed };
     }
 
     render() {
@@ -88,8 +88,8 @@ class StateTreeNode extends PluginUIComponent<{ cell: StateObjectCell, depth: nu
             return null;
         }
 
-        const cellState = cell.transform.state;
-        const showLabel = cell.status !== 'ok' || !cell.transform.state.isGhost;
+        const cellState = cell.state;
+        const showLabel = cell.status !== 'ok' || !cell.state.isGhost;
         const children = cell.parent.tree.children.get(this.ref);
         const newDepth = showLabel ? this.props.depth + 1 : this.props.depth;
 
@@ -140,7 +140,7 @@ class StateTreeNodeLabel extends PluginUIComponent<
             if (e.state.transforms.has(this.ref)) {
                 this.setState({
                     isCurrent: this.props.cell.parent.current === this.ref,
-                    isCollapsed: !!this.props.cell.transform.state.isCollapsed
+                    isCollapsed: !!this.props.cell.state.isCollapsed
                 });
             }
         });
@@ -148,12 +148,12 @@ class StateTreeNodeLabel extends PluginUIComponent<
 
     state = {
         isCurrent: this.props.cell.parent.current === this.ref,
-        isCollapsed: !!this.props.cell.transform.state.isCollapsed
+        isCollapsed: !!this.props.cell.state.isCollapsed
     }
 
     static getDerivedStateFromProps(props: _Props<StateTreeNodeLabel>, state: _State<StateTreeNodeLabel>): _State<StateTreeNodeLabel> | null {
         const isCurrent = props.cell.parent.current === props.cell.transform.ref;
-        const isCollapsed = !!props.cell.transform.state.isCollapsed;
+        const isCollapsed = !!props.cell.state.isCollapsed;
 
         if (state.isCollapsed === isCollapsed && state.isCurrent === isCurrent) return null;
         return { isCurrent, isCollapsed };
@@ -227,7 +227,7 @@ class StateTreeNodeLabel extends PluginUIComponent<
         }
 
         const children = cell.parent.tree.children.get(this.ref);
-        const cellState = cell.transform.state;
+        const cellState = cell.state;
 
         const visibility = <button onClick={this.toggleVisible} className={`msp-btn msp-btn-link msp-tree-visibility${cellState.isHidden ? ' msp-tree-visibility-hidden' : ''}`}>
             <span className='msp-icon msp-icon-visual-visibility' />
@@ -244,7 +244,7 @@ class StateTreeNodeLabel extends PluginUIComponent<
             {children.size > 0 &&  <button onClick={this.toggleExpanded} className='msp-btn msp-btn-link msp-tree-toggle-exp-button'>
                 <span className={`msp-icon msp-icon-${cellState.isCollapsed ? 'expand' : 'collapse'}`} />
             </button>}
-            {!cell.transform.state.isLocked && <button onClick={this.remove} className='msp-btn msp-btn-link msp-tree-remove-button'>
+            {!cell.state.isLocked && <button onClick={this.remove} className='msp-btn msp-btn-link msp-tree-remove-button'>
                 <span className='msp-icon msp-icon-remove' />
             </button>}{visibility}
         </div>;

src/mol-state/object.ts

@@ -64,7 +64,7 @@ interface StateObjectCell<T extends StateObject = StateObject, F extends StateTr
     sourceRef: StateTransform.Ref | undefined,
 
     status: StateObjectCell.Status,
-    // state: StateTransform.State,
+    state: StateTransform.State,
 
     params: {
         definition: ParamDefinition.Params,

src/mol-state/state.ts

@@ -79,10 +79,10 @@ class State {
         const cell = this.cells.get(ref);
         if (!cell) return;
 
-        const update = typeof stateOrProvider === 'function' ? stateOrProvider(cell.transform.state) : stateOrProvider;
+        const update = typeof stateOrProvider === 'function' ? stateOrProvider(cell.state) : stateOrProvider;
 
-        if (this._tree.updateState(cell.transform.ref, update)) {
-            cell.transform = this._tree.transforms.get(cell.transform.ref);
+        if (StateTransform.assignState(cell.state, update)) {
+            cell.transform = this._tree.assignState(cell.transform.ref, update);
             this.events.cell.stateUpdated.next({ state: this, ref, cell });
         }
     }
@@ -210,6 +210,7 @@ class State {
             sourceRef: void 0,
             obj: rootObject,
             status: 'ok',
+            state: { ...root.state },
             errorText: void 0,
             params: {
                 definition: {},
@@ -407,7 +408,7 @@ function findDeletes(ctx: UpdateContext): Ref[] {
 
 function syncNewStatesVisitor(n: StateTransform, tree: StateTree, ctx: UpdateContext) {
     const cell = ctx.cells.get(n.ref);
-    if (!cell || !StateTransform.syncState(cell.transform.state, n.state)) return;
+    if (!cell || !StateTransform.syncState(cell.state, n.state)) return;
     ctx.parent.events.cell.stateUpdated.next({ state: ctx.parent, ref: n.ref, cell });
 }
 
@@ -445,6 +446,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

src/mol-state/tree/transient.ts

@@ -140,22 +140,21 @@ class TransientTree implements StateTree {
         return true;
     }
 
-    updateState(ref: StateTransform.Ref, state?: Partial<StateTransform.State>) {
+    assignState(ref: StateTransform.Ref, state?: Partial<StateTransform.State>) {
         ensurePresent(this.transforms, ref);
 
         const old = this.transforms.get(ref);
-        if (!StateTransform.isStateChange(old.state, state)) return false;
-
-        if (this._stateUpdates && this._stateUpdates.has(old.ref)) {
+        if (this._stateUpdates && this._stateUpdates.has(ref)) {
             StateTransform.assignState(old.state, state);
+            return old;
         } else {
             if (!this._stateUpdates) this._stateUpdates = new Set();
             this._stateUpdates.add(old.ref);
             this.changeNodes();
-            this.transforms.set(ref, StateTransform.withState(old, state));
+            const updated = StateTransform.withState(old, state);
+            this.transforms.set(ref, updated);
+            return updated;
         }
-
-        return true;
     }
 
     remove(ref: StateTransform.Ref): StateTransform[] {