state & plugin: refactoring

src/mol-plugin/behavior/behavior.ts

@@ -22,8 +22,8 @@ interface PluginBehavior<P = unknown> {
 }
 
 namespace PluginBehavior {
-    export class Root extends PluginStateObject.Create({ name: 'Root', shortName: 'R', typeClass: 'Root', description: 'Where everything begins.' }) { }
-    export class Behavior extends PluginStateObject.CreateBehavior<PluginBehavior>({ name: 'Behavior', shortName: 'B', description: 'Modifies plugin functionality.' }) { }
+    export class Root extends PluginStateObject.Create({ name: 'Root', typeClass: 'Root' }) { }
+    export class Behavior extends PluginStateObject.CreateBehavior<PluginBehavior>({ name: 'Behavior' }) { }
 
     export interface Ctor<P = undefined> { new(ctx: PluginContext, params?: P): PluginBehavior<P> }
 
@@ -49,7 +49,7 @@ namespace PluginBehavior {
             params: params.params,
             apply({ params: p }, ctx: PluginContext) {
                 const label = params.label ? params.label(p) : { label: params.display.name, description: params.display.description };
-                return new Behavior(label, new params.ctor(ctx, p));
+                return new Behavior(new params.ctor(ctx, p), label);
             },
             update({ b, newParams }) {
                 return Task.create('Update Behavior', async () => {

src/mol-plugin/command/state.ts

@@ -5,10 +5,10 @@
  */
 
 import { PluginCommand } from './command';
-import { Transform, StateTree, State } from 'mol-state';
+import { Transform, State } from 'mol-state';
 
 export const SetCurrentObject = PluginCommand<{ state: State, ref: Transform.Ref }>('ms-data', 'set-current-object');
-export const Update = PluginCommand<{ state: State, tree: StateTree }>('ms-data', 'update');
+export const Update = PluginCommand<{ state: State, tree: State.Tree | State.Builder }>('ms-data', 'update');
 
 // export const UpdateObject = PluginCommand<{ ref: Transform.Ref, params: any }>('ms-data', 'update-object');
 

src/mol-plugin/context.ts

@@ -4,7 +4,7 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import { StateTree, Transformer, Transform, State } from 'mol-state';
+import { Transformer, Transform, State } from 'mol-state';
 import { Canvas3D } from 'mol-canvas3d/canvas3d';
 import { StateTransforms } from './state/transforms';
 import { PluginStateObject as PSO } from './state/base';
@@ -96,12 +96,12 @@ export class PluginContext {
     }
 
     applyTransform(state: State, a: Transform.Ref, transformer: Transformer, params: any) {
-        const tree = state.tree.build().to(a).apply(transformer, params).getTree();
+        const tree = state.tree.build().to(a).apply(transformer, params);
         return PluginCommands.State.Update.dispatch(this, { state, tree });
     }
 
     updateTransform(state: State, a: Transform.Ref, params: any) {
-        const tree = StateTree.updateParams(state.tree, a, params);
+        const tree = state.build().to(a).update(params);
         return PluginCommands.State.Update.dispatch(this, { state, tree });
     }
 
@@ -122,8 +122,9 @@ export class PluginContext {
         const newTree = b.toRoot()
             .apply(StateTransforms.Data.Download, { url })
             .apply(StateTransforms.Data.ParseCif)
-            .apply(StateTransforms.Model.ParseModelsFromMmCif, {}, { ref: 'models' })
-            .apply(StateTransforms.Model.CreateStructureFromModel, { modelIndex: 0 }, { ref: 'structure' })
+            .apply(StateTransforms.Model.ParseTrajectoryFromMmCif, {}, { ref: 'trajectory' })
+            .apply(StateTransforms.Model.CreateModelFromTrajectory, { modelIndex: 0 }, { ref: 'model' })
+            .apply(StateTransforms.Model.CreateStructureFromModel, { }, { ref: 'structure' })
             .apply(StateTransforms.Model.CreateStructureAssembly)
             .apply(StateTransforms.Model.CreateStructureSelection, { query, label: 'ALA residues' })
             .apply(StateTransforms.Visuals.CreateStructureRepresentation)
@@ -134,8 +135,8 @@ export class PluginContext {
 
     private initEvents() {
         merge(this.events.state.data.object.created, this.events.state.behavior.object.created).subscribe(o => {
-            console.log('creating', o.obj.type);
             if (!PSO.isBehavior(o.obj)) return;
+            console.log('registering behavior', o.obj.label);
             o.obj.data.register();
         });
 
@@ -159,11 +160,15 @@ export class PluginContext {
         this.canvas3d.requestDraw(true);
     }
 
-    _test_nextModel() {
-        const models = this.state.data.select('models')[0].obj as SO.Molecule.Models;
-        const idx = (this.state.data.tree.nodes.get('structure')!.params as Transformer.Params<typeof StateTransforms.Model.CreateStructureFromModel>).modelIndex;
-        const newTree = StateTree.updateParams(this.state.data.tree, 'structure', { modelIndex: (idx + 1) % models.data.length });
-        return this.runTask(this.state.data.update(newTree));
+    async _test_nextModel() {
+        const traj = this.state.data.select('trajectory')[0].obj as SO.Molecule.Trajectory;
+        //const modelIndex = (this.state.data.select('model')[0].transform.params as CreateModelFromTrajectory.Params).modelIndex;
+        const newTree = this.state.data.build().to('model').update(
+            StateTransforms.Model.CreateModelFromTrajectory,
+            old => ({ modelIndex: (old.modelIndex + 1) % traj.data.length }))
+            .getTree();
+        // const newTree = StateTree.updateParams(this.state.data.tree, 'model', { modelIndex: (modelIndex + 1) % traj.data.length });
+        await this.runTask(this.state.data.update(newTree));
         // this.viewer.requestDraw(true);
     }
 

src/mol-plugin/state.ts

@@ -37,8 +37,8 @@ class PluginState {
     }
 
     constructor(private plugin: import('./context').PluginContext) {
-        this.data = State.create(new SO.Root({ label: 'Root' }, { }), { globalContext: plugin });
-        this.behavior = State.create(new PluginBehavior.Root({ label: 'Root' }, { }), { globalContext: plugin });
+        this.data = State.create(new SO.Root({ }), { globalContext: plugin });
+        this.behavior = State.create(new PluginBehavior.Root({ }), { globalContext: plugin });
     }
 }
 

src/mol-plugin/state/base.ts

@@ -11,27 +11,26 @@ import { PluginBehavior } from 'mol-plugin/behavior/behavior';
 export type TypeClass = 'root' | 'data' | 'prop'
 
 export namespace PluginStateObject {
-    export type Any = StateObject<Props, any, TypeInfo>
+    export type Any = StateObject<any, TypeInfo>
 
     export type TypeClass = 'Root' | 'Group' | 'Data' | 'Object' | 'Representation3D' | 'Behavior'
-    export interface TypeInfo { name: string, shortName: string, description: string, typeClass: TypeClass }
-    export interface Props { label: string, description?: string }
+    export interface TypeInfo { name: string, typeClass: TypeClass }
 
-    export const Create = StateObject.factory<TypeInfo, Props>();
+    export const Create = StateObject.factory<TypeInfo>();
 
-    export function isRepresentation3D(o?: Any): o is StateObject<Props, Representation.Any, TypeInfo> {
+    export function isRepresentation3D(o?: Any): o is StateObject<Representation.Any, TypeInfo> {
         return !!o && o.type.typeClass === 'Representation3D';
     }
 
-    export function isBehavior(o?: Any): o is StateObject<Props, PluginBehavior, TypeInfo> {
+    export function isBehavior(o?: Any): o is StateObject<PluginBehavior, TypeInfo> {
         return !!o && o.type.typeClass === 'Behavior';
     }
 
-    export function CreateRepresentation3D<T extends Representation.Any>(type: { name: string, shortName: string, description: string }) {
+    export function CreateRepresentation3D<T extends Representation.Any>(type: { name: string }) {
         return Create<T>({ ...type, typeClass: 'Representation3D' })
     }
 
-    export function CreateBehavior<T extends PluginBehavior>(type: { name: string, shortName: string, description: string }) {
+    export function CreateBehavior<T extends PluginBehavior>(type: { name: string }) {
         return Create<T>({ ...type, typeClass: 'Behavior' })
     }
 }

src/mol-plugin/state/objects.ts

@@ -14,15 +14,15 @@ import { VolumeData } from 'mol-model/volume';
 const _create = PluginStateObject.Create, _createRepr3D = PluginStateObject.CreateRepresentation3D
 
 namespace PluginStateObjects {
-    export class Root extends _create({ name: 'Root', shortName: 'R', typeClass: 'Root', description: 'Where everything begins.' }) { }
+    export class Root extends _create({ name: 'Root', typeClass: 'Root' }) { }
 
-    export class Group extends _create({ name: 'Group', shortName: 'G', typeClass: 'Group', description: 'A group on entities.' }) { }
+    export class Group extends _create({ name: 'Group', typeClass: 'Group' }) { }
 
     export namespace Data {
-        export class String extends _create<string>({ name: 'String Data', typeClass: 'Data', shortName: 'S_D', description: 'A string.' }) { }
-        export class Binary extends _create<Uint8Array>({ name: 'Binary Data', typeClass: 'Data', shortName: 'B_D', description: 'A binary blob.' }) { }
-        export class Json extends _create<any>({ name: 'JSON Data', typeClass: 'Data', shortName: 'JS_D', description: 'Represents JSON data.' }) { }
-        export class Cif extends _create<CifFile>({ name: 'Cif File', typeClass: 'Data', shortName: 'CF', description: 'Represents parsed CIF data.' }) { }
+        export class String extends _create<string>({ name: 'String Data', typeClass: 'Data', }) { }
+        export class Binary extends _create<Uint8Array>({ name: 'Binary Data', typeClass: 'Data' }) { }
+        export class Json extends _create<any>({ name: 'JSON Data', typeClass: 'Data' }) { }
+        export class Cif extends _create<CifFile>({ name: 'CIF File', typeClass: 'Data' }) { }
 
         // TODO
         // export class MultipleRaw extends _create<{
@@ -31,14 +31,15 @@ namespace PluginStateObjects {
     }
 
     export namespace Molecule {
-        export class Models extends _create<ReadonlyArray<_Model>>({ name: 'Molecule Model', typeClass: 'Object', shortName: 'M_M', description: 'A model of a molecule.' }) { }
-        export class Structure extends _create<_Structure>({ name: 'Molecule Structure', typeClass: 'Object', shortName: 'M_S', description: 'A structure of a molecule.' }) { }
-        export class Representation3D extends _createRepr3D<StructureRepresentation<any>>({ name: 'Molecule Structure 3D Representation', shortName: 'S_R', description: 'A 3D representation of a molecular structure.' }) { }
+        export class Trajectory extends _create<ReadonlyArray<_Model>>({ name: 'Trajectory', typeClass: 'Object' }) { }
+        export class Model extends _create<_Model>({ name: 'Model', typeClass: 'Object' }) { }
+        export class Structure extends _create<_Structure>({ name: 'Structure', typeClass: 'Object' }) { }
+        export class Representation3D extends _createRepr3D<StructureRepresentation<any>>({ name: 'Structure 3D' }) { }
     }
 
     export namespace Volume {
-        export class Data extends _create<VolumeData>({ name: 'Volume Data', typeClass: 'Object', shortName: 'V_D', description: 'Volume Data.' }) { }
-        export class Representation3D extends _createRepr3D<VolumeRepresentation<any>>({ name: 'Volume 3D Representation', shortName: 'V_R', description: 'A 3D representation of volumetric data.' }) { }
+        export class Data extends _create<VolumeData>({ name: 'Volume Data', typeClass: 'Object' }) { }
+        export class Representation3D extends _createRepr3D<VolumeRepresentation<any>>({ name: 'Volume 3D' }) { }
     }
 
 }

src/mol-plugin/state/transforms/data.ts

@@ -35,8 +35,8 @@ const Download = PluginStateTransform.Create<SO.Root, SO.Data.String | SO.Data.B
             // TODO: track progress
             const data = await globalCtx.fetch(p.url, p.isBinary ? 'binary' : 'string');
             return p.isBinary
-                ? new SO.Data.Binary({ label: p.label ? p.label : p.url }, data as Uint8Array)
-                : new SO.Data.String({ label: p.label ? p.label : p.url }, data as string);
+                ? new SO.Data.Binary(data as Uint8Array, { label: p.label ? p.label : p.url })
+                : new SO.Data.String(data as string, { label: p.label ? p.label : p.url });
         });
     }
 });
@@ -55,7 +55,7 @@ const ParseCif = PluginStateTransform.Create<SO.Data.String | SO.Data.Binary, SO
         return Task.create('Parse CIF', async ctx => {
             const parsed = await (SO.Data.String.is(a) ? CIF.parse(a.data) : CIF.parseBinary(a.data)).runInContext(ctx);
             if (parsed.isError) throw new Error(parsed.message);
-            return new SO.Data.Cif({ label: 'CIF File' }, parsed.result);
+            return new SO.Data.Cif(parsed.result);
         });
     }
 });

src/mol-plugin/state/transforms/model.ts

@@ -13,16 +13,16 @@ import Expression from 'mol-script/language/expression';
 import { compile } from 'mol-script/runtime/query/compiler';
 import { Mat4 } from 'mol-math/linear-algebra';
 
-export { ParseModelsFromMmCif }
-namespace ParseModelsFromMmCif { export interface Params { blockHeader?: string } }
-const ParseModelsFromMmCif = PluginStateTransform.Create<SO.Data.Cif, SO.Molecule.Models, ParseModelsFromMmCif.Params>({
-    name: 'parse-models-from-mmcif',
+export { ParseTrajectoryFromMmCif }
+namespace ParseTrajectoryFromMmCif { export interface Params { blockHeader?: string } }
+const ParseTrajectoryFromMmCif = PluginStateTransform.Create<SO.Data.Cif, SO.Molecule.Trajectory, ParseTrajectoryFromMmCif.Params>({
+    name: 'parse-trajectory-from-mmcif',
     display: {
         name: 'Models from mmCIF',
         description: 'Identify and create all separate models in the specified CIF data block'
     },
     from: [SO.Data.Cif],
-    to: [SO.Molecule.Models],
+    to: [SO.Molecule.Trajectory],
     params: {
         default: a => ({ blockHeader: a.data.blocks[0].header }),
         controls(a) {
@@ -40,35 +40,57 @@ const ParseModelsFromMmCif = PluginStateTransform.Create<SO.Data.Cif, SO.Molecul
             if (!block) throw new Error(`Data block '${[header]}' not found.`);
             const models = await Model.create(Format.mmCIF(block)).runInContext(ctx);
             if (models.length === 0) throw new Error('No models found.');
-            const label = models.length === 1 ? `${models[0].label}` : `${models[0].label} (${models.length} models)`;
-            return new SO.Molecule.Models({ label }, models);
+            const label = { label: models[0].label, description: `${models.length} model${models.length === 1 ? '' : 's'}` };
+            return new SO.Molecule.Trajectory(models, label);
         });
     }
 });
 
-export { CreateStructureFromModel }
-namespace CreateStructureFromModel { export interface Params { modelIndex: number, transform3d?: Mat4 } }
-const CreateStructureFromModel = PluginStateTransform.Create<SO.Molecule.Models, SO.Molecule.Structure, CreateStructureFromModel.Params>({
-    name: 'create-structure-from-model',
+export { CreateModelFromTrajectory }
+namespace CreateModelFromTrajectory { export interface Params { modelIndex: number } }
+const CreateModelFromTrajectory = PluginStateTransform.Create<SO.Molecule.Trajectory, SO.Molecule.Model, CreateModelFromTrajectory.Params>({
+    name: 'create-model-from-trajectory',
     display: {
-        name: 'Structure from Model',
+        name: 'Model from Trajectory',
         description: 'Create a molecular structure from the specified model.'
     },
-    from: [SO.Molecule.Models],
-    to: [SO.Molecule.Structure],
+    from: [SO.Molecule.Trajectory],
+    to: [SO.Molecule.Model],
     params: {
         default: () => ({ modelIndex: 0 }),
         controls: a => ({ modelIndex: PD.Range('Model Index', 'Model Index', 0, 0, Math.max(0, a.data.length - 1), 1) })
     },
     isApplicable: a => a.data.length > 0,
     apply({ a, params }) {
+        console.log('parans', params);
         if (params.modelIndex < 0 || params.modelIndex >= a.data.length) throw new Error(`Invalid modelIndex ${params.modelIndex}`);
-        let s = Structure.ofModel(a.data[params.modelIndex]);
+        const model = a.data[params.modelIndex];
+        const label = { label: `Model ${model.modelNum}` };
+        return new SO.Molecule.Model(model, label);
+    }
+});
+
+export { CreateStructureFromModel }
+namespace CreateStructureFromModel { export interface Params { transform3d?: Mat4 } }
+const CreateStructureFromModel = PluginStateTransform.Create<SO.Molecule.Model, SO.Molecule.Structure, CreateStructureFromModel.Params>({
+    name: 'create-structure-from-model',
+    display: {
+        name: 'Structure from Model',
+        description: 'Create a molecular structure from the specified model.'
+    },
+    from: [SO.Molecule.Model],
+    to: [SO.Molecule.Structure],
+    apply({ a, params }) {
+        let s = Structure.ofModel(a.data);
         if (params.transform3d) s = Structure.transform(s, params.transform3d);
-        return new SO.Molecule.Structure({ label: `Model ${s.models[0].modelNum}`, description: s.elementCount === 1 ? '1 element' : `${s.elementCount} elements` }, s);
+        const label = { label: a.data.label, description: s.elementCount === 1 ? '1 element' : `${s.elementCount} elements` };
+        return new SO.Molecule.Structure(s, label);
     }
 });
 
+function structureDesc(s: Structure) {
+    return s.elementCount === 1 ? '1 element' : `${s.elementCount} elements`;
+}
 
 export { CreateStructureAssembly }
 namespace CreateStructureAssembly { export interface Params { /** if not specified, use the 1st */ id?: string } }
@@ -98,7 +120,8 @@ const CreateStructureAssembly = PluginStateTransform.Create<SO.Molecule.Structur
             if (!asm) throw new Error(`Assembly '${id}' not found`);
 
             const s = await StructureSymmetry.buildAssembly(a.data, id!).runInContext(ctx);
-            return new SO.Molecule.Structure({ label: `Assembly ${id}`, description: s.elementCount === 1 ? '1 element' : `${s.elementCount} elements` }, s);
+            const label = { label: `Assembly ${id}`, description: structureDesc(s) };
+            return new SO.Molecule.Structure(s, label);
         })
     }
 });
@@ -118,6 +141,7 @@ const CreateStructureSelection = PluginStateTransform.Create<SO.Molecule.Structu
         const compiled = compile<StructureSelection>(params.query);
         const result = compiled(new QueryContext(a.data));
         const s = StructureSelection.unionStructure(result);
-        return new SO.Molecule.Structure({ label: `${params.label || 'Selection'}`, description: s.elementCount === 1 ? '1 element' : `${s.elementCount} elements` }, s);
+        const label = { label: `${params.label || 'Selection'}`, description: structureDesc(s) };
+        return new SO.Molecule.Structure(s, label);
     }
 });

src/mol-plugin/state/transforms/visuals.ts

@@ -23,7 +23,7 @@ const CreateStructureRepresentation = PluginStateTransform.Create<SO.Molecule.St
         return Task.create('Structure Representation', async ctx => {
             const repr = BallAndStickRepresentation(); // CartoonRepresentation();
             await repr.createOrUpdate({ webgl: plugin.canvas3d.webgl }, DefaultBallAndStickProps, a.data).runInContext(ctx);
-            return new SO.Molecule.Representation3D({ label: 'Visual Repr.' }, repr);
+            return new SO.Molecule.Representation3D(repr);
         });
     },
     update({ a, b }, plugin: PluginContext) {

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

@@ -44,18 +44,16 @@ export class StateTreeNode extends React.Component<{ plugin: PluginContext, node
             }}>{name}</a>: <i>{cell.errorText}</i></>;
         } else {
             const obj = cell.obj as PluginStateObject.Any;
-            const props = obj.props;
-            const type = obj.type;
-            label = <>[<span title={type.description}>{ type.shortName }</span>] <a href='#' onClick={e => {
+            label = <><a href='#' onClick={e => {
                 e.preventDefault();
                 PluginCommands.State.SetCurrentObject.dispatch(this.props.plugin, { state: this.props.state, ref: this.props.nodeRef });
-            }}>{props.label}</a> {props.description ? <small>{props.description}</small> : void 0}</>;
+            }}>{obj.label}</a> {obj.description ? <small>{obj.description}</small> : void 0}</>;
         }
 
         const children = this.props.state.tree.children.get(this.props.nodeRef);
 
         return <div>
-            {remove}{label}
+            {remove} {label}
             {children.size === 0
                 ? void 0
                 : <div style={{ marginLeft: '7px', paddingLeft: '3px', borderLeft: '1px solid #999' }}>{children.map(c => <StateTreeNode plugin={this.props.plugin} state={this.props.state} nodeRef={c!} key={c} />)}</div>

src/mol-state/object.ts

@@ -9,28 +9,34 @@ import { Transform } from './transform';
 
 export { StateObject, StateObjectCell }
 
-interface StateObject<P = any, D = any, T = any> {
+interface StateObject<D = any, T extends StateObject.Type = { name: string, typeClass: any }> {
     readonly id: UUID,
-    readonly type: StateObject.Type<T>,
-    readonly props: P,
-    readonly data: D
+    readonly type: T,
+    readonly data: D,
+    readonly label: string,
+    readonly description?: string,
 }
 
 namespace StateObject {
-    export function factory<Type, CommonProps>() {
-        return <D = { }, P = {}>(type: Type) => create<P & CommonProps, D, Type>(type);
+    export function factory<T extends Type>() {
+        return <D = { }>(type: T) => create<D, T>(type);
     }
 
-    export type Type<I = unknown> = I
+    export type Type<Cls extends string = string> = { name: string, typeClass: Cls }
     export type Ctor = { new(...args: any[]): StateObject, type: any }
 
-    export function create<Props, Data, Type>(type: Type) {
-        return class implements StateObject<Props, Data, Type> {
+    export function create<Data, T extends Type>(type: T) {
+        return class implements StateObject<Data, T> {
             static type = type;
-            static is(obj?: StateObject): obj is StateObject<Props, Data, Type> { return !!obj && type === obj.type; }
+            static is(obj?: StateObject): obj is StateObject<Data, T> { return !!obj && type === obj.type; }
             id = UUID.create();
             type = type;
-            constructor(public props: Props, public data: Data) { }
+            label: string;
+            description?: string;
+            constructor(public data: Data, props?: { label: string, description?: string }) {
+                this.label = props && props.label || type.name;
+                this.description = props && props.description;
+            }
         }
     }
 }

src/mol-state/state.ts

@@ -12,6 +12,7 @@ import { UUID } from 'mol-util';
 import { RuntimeContext, Task } from 'mol-task';
 import { StateSelection } from './state/selection';
 import { RxEventHelper } from 'mol-util/rx-event-helper';
+import { StateTreeBuilder } from './tree/builder';
 
 export { State }
 
@@ -43,6 +44,8 @@ class State {
     get tree() { return this._tree; }
     get current() { return this._current; }
 
+    build() { return this._tree.build(); }
+
     readonly cells: State.Cells = new Map();
 
     getSnapshot(): State.Snapshot {
@@ -77,22 +80,19 @@ class State {
         return StateSelection.select(selector(StateSelection.Generators), this)
     }
 
-    query(q: StateSelection.Query) {
-        return q(this);
-    }
-
-    update(tree: StateTree): Task<void> {
+    update(tree: StateTree | StateTreeBuilder): Task<void> {
         // TODO: support cell state
+        const _tree = StateTreeBuilder.is(tree) ? tree.getTree() : tree;
         return Task.create('Update Tree', async taskCtx => {
             try {
                 const oldTree = this._tree;
-                this._tree = tree;
+                this._tree = _tree;
 
                 const ctx: UpdateContext = {
                     parent: this,
                     taskCtx,
                     oldTree,
-                    tree,
+                    tree: _tree,
                     cells: this.cells as Map<Transform.Ref, StateObjectCell>,
                     transformCache: this.transformCache
                 };
@@ -123,6 +123,9 @@ class State {
 namespace State {
     export type Cells = ReadonlyMap<Transform.Ref, StateObjectCell>
 
+    export type Tree = StateTree
+    export type Builder = StateTreeBuilder
+
     export interface ObjectEvent {
         state: State,
         ref: Ref
@@ -202,6 +205,7 @@ function setCellStatus(ctx: UpdateContext, ref: Ref, status: StateObjectCell.Sta
 }
 
 function _initCellStatusVisitor(t: Transform, _: any, ctx: UpdateContext) {
+    ctx.cells.get(t.ref)!.transform = t;
     setCellStatus(ctx, t.ref, 'pending');
 }
 

src/mol-state/tree/builder.ts

@@ -21,6 +21,10 @@ namespace StateTreeBuilder {
         tree: TransientTree
     }
 
+    export function is(obj: any): obj is StateTreeBuilder {
+        return !!obj && typeof (obj as StateTreeBuilder).getTree === 'function';
+    }
+
     export class Root implements StateTreeBuilder {
         private state: State;
         to<A extends StateObject>(ref: Transform.Ref) { return new To<A>(this.state, ref, this); }
@@ -40,6 +44,20 @@ namespace StateTreeBuilder {
             return new To(this.state, t.ref, this.root);
         }
 
+        update<T extends Transformer<A, any, any>>(transformer: T, params: (old: Transformer.Params<T>) => Transformer.Params<T>): Root
+        update(params: any): Root
+        update<T extends Transformer<A, any, any>>(paramsOrTransformer: T, provider?: (old: Transformer.Params<T>) => Transformer.Params<T>) {
+            const old = this.state.tree.nodes.get(this.ref)!;
+            let params: any;
+            if (provider) {
+                params = provider(old.params as any);
+            } else {
+                params = paramsOrTransformer;
+            }
+            this.state.tree.set(Transform.updateParams(old, params));
+            return this.root;
+        }
+
         and() { return this.root; }
 
         getTree(): StateTree { return this.state.tree.asImmutable(); }

src/mol-state/tree/immutable.ts

@@ -8,7 +8,6 @@ import { Map as ImmutableMap, OrderedSet } from 'immutable';
 import { Transform } from '../transform';
 import { TransientTree } from './transient';
 import { StateTreeBuilder } from './builder';
-import { Transformer } from '../transformer';
 
 export { StateTree }
 
@@ -66,12 +65,6 @@ namespace StateTree {
         return new Impl(nodes, children);
     }
 
-    export function updateParams<T extends Transformer = Transformer>(tree: StateTree, ref: Transform.Ref, params: Transformer.Params<T>): StateTree {
-        const t = tree.nodes.get(ref)!;
-        const newTransform = Transform.updateParams(t, params);
-        return tree.asTransient().set(newTransform).asImmutable();
-    }
-
     type VisitorCtx = { tree: StateTree, state: any, f: (node: Node, tree: StateTree, state: any) => boolean | undefined | void };
 
     function _postOrderFunc(this: VisitorCtx, c: Ref | undefined) { _doPostOrder(this, this.tree.nodes.get(c!)); }

src/perf-tests/state.ts

@@ -1,133 +1,133 @@
-import { State, StateObject, StateTree, Transformer } from 'mol-state';
-import { Task } from 'mol-task';
-import * as util from 'util';
-
-export type TypeClass = 'root' | 'shape' | 'prop'
-export interface ObjProps { label: string }
-export interface TypeInfo { name: string, class: TypeClass }
-
-const _obj = StateObject.factory<TypeInfo, ObjProps>()
-const _transform = Transformer.factory('test');
-
-export class Root extends _obj({ name: 'Root', class: 'root' }) { }
-export class Square extends _obj<{ a: number }>({ name: 'Square', class: 'shape' }) { }
-export class Circle extends _obj<{ r: number }>({ name: 'Circle', class: 'shape' }) { }
-export class Area extends _obj<{ volume: number }>({ name: 'Area', class: 'prop' }) { }
-
-export const CreateSquare = _transform<Root, Square, { a: number }>({
-    name: 'create-square',
-    from: [Root],
-    to: [Square],
-    apply({ params: p }) {
-        return new Square({ label: `Square a=${p.a}` }, p);
-    },
-    update({ b, newParams: p }) {
-        b.props.label = `Square a=${p.a}`
-        b.data.a = p.a;
-        return Transformer.UpdateResult.Updated;
-    }
-});
-
-export const CreateCircle = _transform<Root, Circle, { r: number }>({
-    name: 'create-circle',
-    from: [Root],
-    to: [Square],
-    apply({ params: p }) {
-        return new Circle({ label: `Circle r=${p.r}` }, p);
-    },
-    update({ b, newParams: p }) {
-        b.props.label = `Circle r=${p.r}`
-        b.data.r = p.r;
-        return Transformer.UpdateResult.Updated;
-    }
-});
-
-export const CaclArea = _transform<Square | Circle, Area, {}>({
-    name: 'calc-area',
-    from: [Square, Circle],
-    to: [Area],
-    apply({ a }) {
-        if (a instanceof Square) return new Area({ label: 'Area' }, { volume: a.data.a * a.data.a });
-        else if (a instanceof Circle) return new Area({ label: 'Area' }, { volume: a.data.r * a.data.r * Math.PI });
-        throw new Error('Unknown object type.');
-    },
-    update({ a, b }) {
-        if (a instanceof Square) b.data.volume = a.data.a * a.data.a;
-        else if (a instanceof Circle) b.data.volume = a.data.r * a.data.r * Math.PI;
-        else throw new Error('Unknown object type.');
-        return Transformer.UpdateResult.Updated;
-    }
-});
-
-export async function runTask<A>(t: A | Task<A>): Promise<A> {
-    if ((t as any).run) return await (t as Task<A>).run();
-    return t as A;
-}
-
-function hookEvents(state: State) {
-    state.events.object.created.subscribe(e => console.log('created:', e.ref));
-    state.events.object.removed.subscribe(e => console.log('removed:', e.ref));
-    state.events.object.replaced.subscribe(e => console.log('replaced:', e.ref));
-    state.events.object.cellState.subscribe(e => console.log('stateChanged:', e.ref, e.cell.status));
-    state.events.object.updated.subscribe(e => console.log('updated:', e.ref));
-}
-
-export async function testState() {
-    const state = State.create(new Root({ label: 'Root' }, { }));
-    hookEvents(state);
-
-    const tree = state.tree;
-    const builder = tree.build();
-    builder.toRoot<Root>()
-        .apply(CreateSquare, { a: 10 }, { ref: 'square' })
-        .apply(CaclArea);
-    const tree1 = builder.getTree();
-
-    printTTree(tree1);
-
-    const tree2 = StateTree.updateParams<typeof CreateSquare>(tree1, 'square', { a: 15 });
-    printTTree(tree1);
-    printTTree(tree2);
-
-    await state.update(tree1).run();
-    console.log('----------------');
-    console.log(util.inspect(state.cells, true, 3, true));
-
-    console.log('----------------');
-    const jsonString = JSON.stringify(StateTree.toJSON(tree2), null, 2);
-    const jsonData = JSON.parse(jsonString);
-    printTTree(tree2);
-    console.log(jsonString);
-    const treeFromJson = StateTree.fromJSON(jsonData);
-    printTTree(treeFromJson);
-
-    console.log('----------------');
-    await state.update(treeFromJson).run();
-    console.log(util.inspect(state.cells, true, 3, true));
-
-    console.log('----------------');
-
-    const sel = state.select('square');
-    console.log(sel);
-}
-
-testState();
-
-
-//test();
-
-export function printTTree(tree: StateTree) {
-    let lines: string[] = [];
-    function print(offset: string, ref: any) {
-        const t = tree.nodes.get(ref)!;
-        const tr = t;
-
-        const name = tr.transformer.id;
-        lines.push(`${offset}|_ (${ref}) ${name} ${tr.params ? JSON.stringify(tr.params) : ''}, v${t.version}`);
-        offset += '   ';
-
-        tree.children.get(ref).forEach(c => print(offset, c!));
-    }
-    print('', tree.root.ref);
-    console.log(lines.join('\n'));
-}
+// import { State, StateObject, StateTree, Transformer } from 'mol-state';
+// import { Task } from 'mol-task';
+// import * as util from 'util';
+
+// export type TypeClass = 'root' | 'shape' | 'prop'
+// export interface ObjProps { label: string }
+// export interface TypeInfo { name: string, class: TypeClass }
+
+// const _obj = StateObject.factory<TypeInfo, ObjProps>()
+// const _transform = Transformer.factory('test');
+
+// export class Root extends _obj({ name: 'Root', class: 'root' }) { }
+// export class Square extends _obj<{ a: number }>({ name: 'Square', class: 'shape' }) { }
+// export class Circle extends _obj<{ r: number }>({ name: 'Circle', class: 'shape' }) { }
+// export class Area extends _obj<{ volume: number }>({ name: 'Area', class: 'prop' }) { }
+
+// export const CreateSquare = _transform<Root, Square, { a: number }>({
+//     name: 'create-square',
+//     from: [Root],
+//     to: [Square],
+//     apply({ params: p }) {
+//         return new Square({ label: `Square a=${p.a}` }, p);
+//     },
+//     update({ b, newParams: p }) {
+//         b.props.label = `Square a=${p.a}`
+//         b.data.a = p.a;
+//         return Transformer.UpdateResult.Updated;
+//     }
+// });
+
+// export const CreateCircle = _transform<Root, Circle, { r: number }>({
+//     name: 'create-circle',
+//     from: [Root],
+//     to: [Square],
+//     apply({ params: p }) {
+//         return new Circle({ label: `Circle r=${p.r}` }, p);
+//     },
+//     update({ b, newParams: p }) {
+//         b.props.label = `Circle r=${p.r}`
+//         b.data.r = p.r;
+//         return Transformer.UpdateResult.Updated;
+//     }
+// });
+
+// export const CaclArea = _transform<Square | Circle, Area, {}>({
+//     name: 'calc-area',
+//     from: [Square, Circle],
+//     to: [Area],
+//     apply({ a }) {
+//         if (a instanceof Square) return new Area({ label: 'Area' }, { volume: a.data.a * a.data.a });
+//         else if (a instanceof Circle) return new Area({ label: 'Area' }, { volume: a.data.r * a.data.r * Math.PI });
+//         throw new Error('Unknown object type.');
+//     },
+//     update({ a, b }) {
+//         if (a instanceof Square) b.data.volume = a.data.a * a.data.a;
+//         else if (a instanceof Circle) b.data.volume = a.data.r * a.data.r * Math.PI;
+//         else throw new Error('Unknown object type.');
+//         return Transformer.UpdateResult.Updated;
+//     }
+// });
+
+// export async function runTask<A>(t: A | Task<A>): Promise<A> {
+//     if ((t as any).run) return await (t as Task<A>).run();
+//     return t as A;
+// }
+
+// function hookEvents(state: State) {
+//     state.events.object.created.subscribe(e => console.log('created:', e.ref));
+//     state.events.object.removed.subscribe(e => console.log('removed:', e.ref));
+//     state.events.object.replaced.subscribe(e => console.log('replaced:', e.ref));
+//     state.events.object.cellState.subscribe(e => console.log('stateChanged:', e.ref, e.cell.status));
+//     state.events.object.updated.subscribe(e => console.log('updated:', e.ref));
+// }
+
+// export async function testState() {
+//     const state = State.create(new Root({ label: 'Root' }, { }));
+//     hookEvents(state);
+
+//     const tree = state.tree;
+//     const builder = tree.build();
+//     builder.toRoot<Root>()
+//         .apply(CreateSquare, { a: 10 }, { ref: 'square' })
+//         .apply(CaclArea);
+//     const tree1 = builder.getTree();
+
+//     printTTree(tree1);
+
+//     const tree2 = StateTree.updateParams<typeof CreateSquare>(tree1, 'square', { a: 15 });
+//     printTTree(tree1);
+//     printTTree(tree2);
+
+//     await state.update(tree1).run();
+//     console.log('----------------');
+//     console.log(util.inspect(state.cells, true, 3, true));
+
+//     console.log('----------------');
+//     const jsonString = JSON.stringify(StateTree.toJSON(tree2), null, 2);
+//     const jsonData = JSON.parse(jsonString);
+//     printTTree(tree2);
+//     console.log(jsonString);
+//     const treeFromJson = StateTree.fromJSON(jsonData);
+//     printTTree(treeFromJson);
+
+//     console.log('----------------');
+//     await state.update(treeFromJson).run();
+//     console.log(util.inspect(state.cells, true, 3, true));
+
+//     console.log('----------------');
+
+//     const sel = state.select('square');
+//     console.log(sel);
+// }
+
+// testState();
+
+
+// //test();
+
+// export function printTTree(tree: StateTree) {
+//     let lines: string[] = [];
+//     function print(offset: string, ref: any) {
+//         const t = tree.nodes.get(ref)!;
+//         const tr = t;
+
+//         const name = tr.transformer.id;
+//         lines.push(`${offset}|_ (${ref}) ${name} ${tr.params ? JSON.stringify(tr.params) : ''}, v${t.version}`);
+//         offset += '   ';
+
+//         tree.children.get(ref).forEach(c => print(offset, c!));
+//     }
+//     print('', tree.root.ref);
+//     console.log(lines.join('\n'));
+// }