mol-state: wip

src/mol-plugin/behavior/built-in/data.ts

@@ -6,7 +6,6 @@
 
 import { PluginBehavior } from '../behavior';
 import { PluginCommands } from 'mol-plugin/command';
-import { StateTree } from 'mol-state';
 
 export const SetCurrentObject = PluginBehavior.create({
     name: 'set-current-data-object-behavior',
@@ -23,7 +22,7 @@ export const Update = PluginBehavior.create({
 export const RemoveObject = PluginBehavior.create({
     name: 'remove-object-data-behavior',
     ctor: PluginBehavior.simpleCommandHandler(PluginCommands.Data.RemoveObject, ({ ref }, ctx) => {
-        const tree = StateTree.build(ctx.state.data.tree).delete(ref).getTree();
+        const tree = ctx.state.data.tree.build().delete(ref).getTree();
         ctx.state.updateData(tree);
     }),
     display: { name: 'Remove Object Handler', group: 'Data' }

src/mol-plugin/context.ts

@@ -83,7 +83,7 @@ export class PluginContext {
     }
 
     async _test_initBehaviours() {
-        const tree = StateTree.build(this.state.behavior.tree)
+        const tree = this.state.behavior.tree.build()
             .toRoot().apply(PluginBehaviors.Data.SetCurrentObject, { ref: PluginBehaviors.Data.SetCurrentObject.id })
             .and().toRoot().apply(PluginBehaviors.Data.Update, { ref: PluginBehaviors.Data.Update.id })
             .and().toRoot().apply(PluginBehaviors.Data.RemoveObject, { ref: PluginBehaviors.Data.RemoveObject.id })
@@ -96,7 +96,7 @@ export class PluginContext {
     }
 
     _test_applyTransform(a: Transform.Ref, transformer: Transformer, params: any) {
-        const tree = StateTree.build(this.state.data.tree).to(a).apply(transformer, params).getTree();
+        const tree = this.state.data.tree.build().to(a).apply(transformer, params).getTree();
         PluginCommands.Data.Update.dispatch(this, { tree });
     }
 
@@ -106,7 +106,7 @@ export class PluginContext {
     }
 
     _test_createState(url: string) {
-        const b = StateTree.build(this.state.data.tree);
+        const b = this.state.data.tree.build();
 
         const query = MolScriptBuilder.struct.generator.atomGroups({
             // 'atom-test': MolScriptBuilder.core.rel.eq([

src/mol-plugin/state.ts

@@ -26,8 +26,8 @@ class PluginState {
     }
 
     async setSnapshot(snapshot: PluginState.Snapshot) {
-        await this.behavior.setSnapshot(snapshot.behaviour);
-        await this.data.setSnapshot(snapshot.data);
+        await this.plugin.runTask(this.behavior.setSnapshot(snapshot.behaviour));
+        await this.plugin.runTask(this.data.setSnapshot(snapshot.data));
         this.plugin.canvas3d.camera.setState(snapshot.canvas3d.camera);
         this.plugin.canvas3d.requestDraw(true);
     }

src/mol-state/context.ts

@@ -34,15 +34,13 @@ class StateContext {
     };
 
     readonly globalContext: unknown;
-    readonly defaultCellState: unknown;
 
     dispose() {
         this.ev.dispose();
     }
 
-    constructor(params: { globalContext: unknown, defaultCellState: unknown }) {
+    constructor(params: { globalContext: unknown }) {
         this.globalContext = params.globalContext;
-        this.defaultCellState = params.defaultCellState;
         this.behaviors.currentObject.next({ ref: Transform.RootRef });
     }
 }
\ No newline at end of file

src/mol-state/object.ts

@@ -50,9 +50,16 @@ namespace StateObjectCell {
     export type Status = 'ok' | 'error' | 'pending' | 'processing'
 
     export interface State {
-        isVisible: boolean,
+        isObjectHidden: boolean,
         isHidden: boolean,
-        isBound: boolean,
-        isExpanded: boolean
+        isBinding: boolean,
+        isCollapsed: boolean
     }
+
+    export const DefaultState: State = {
+        isObjectHidden: false,
+        isHidden: false,
+        isBinding: false,
+        isCollapsed: false
+    };
 }
\ No newline at end of file

src/mol-state/selection.ts

@@ -6,7 +6,7 @@
 
 import { StateObject, StateObjectCell } from './object';
 import { State } from './state';
-import { ImmutableTree } from './tree';
+import { StateTree } from './tree';
 import { Transform } from './transform';
 
 namespace StateSelection {
@@ -150,7 +150,7 @@ namespace StateSelection {
     export function subtree(b: Selector) {
         return flatMap(b, (n, s) => {
             const nodes = [] as string[];
-            ImmutableTree.doPreOrder(s.tree, s.tree.nodes.get(n.ref), nodes, (x, _, ctx) => { ctx.push(x.ref) });
+            StateTree.doPreOrder(s.tree, s.tree.nodes.get(n.ref), nodes, (x, _, ctx) => { ctx.push(x.ref) });
             return nodes.map(x => s.cells.get(x)!);
         });
     }

src/mol-state/state.ts

@@ -5,17 +5,18 @@
  */
 
 import { StateObject, StateObjectCell } from './object';
-import { StateTree, ImmutableTree } from './tree';
+import { StateTree } from './tree';
 import { Transform } from './transform';
 import { Transformer } from './transformer';
 import { StateContext } from './context';
 import { UUID } from 'mol-util';
 import { RuntimeContext, Task } from 'mol-task';
+import { StateSelection } from './selection';
 
 export { State }
 
 class State {
-    private _tree: StateTree = StateTree.create();
+    private _tree: StateTree = StateTree.createEmpty();
     private _current: Transform.Ref = this._tree.root.ref;
     private transformCache = new Map<Transform.Ref, unknown>();
 
@@ -26,24 +27,12 @@ class State {
     readonly context: StateContext;
 
     getSnapshot(): State.Snapshot {
-        const props = Object.create(null);
-        const keys = this.cells.keys();
-        while (true) {
-            const key = keys.next();
-            if (key.done) break;
-            const o = this.cells.get(key.value)!;
-            props[key.value] = { ...o.state };
-        }
-        return {
-            tree: StateTree.toJSON(this._tree),
-            props
-        };
+        return { tree: StateTree.toJSON(this._tree) };
     }
 
     setSnapshot(snapshot: State.Snapshot) {
         const tree = StateTree.fromJSON(snapshot.tree);
-        // TODO: support props and async
-        return this.update(tree).run();
+        return this.update(tree);
     }
 
     setCurrent(ref: Transform.Ref) {
@@ -51,10 +40,22 @@ class State {
         this.context.behaviors.currentObject.next({ ref });
     }
 
+    updateCellState(ref: Transform.Ref, state?: Partial<StateObjectCell.State>) {
+        // TODO
+    }
+
     dispose() {
         this.context.dispose();
     }
 
+    select(selector: StateSelection.Selector) {
+        return StateSelection.select(selector, this);
+    }
+
+    get selector() {
+        return StateSelection;
+    }
+
     update(tree: StateTree): Task<void> {
         // TODO: support props
         return Task.create('Update Tree', async taskCtx => {
@@ -78,22 +79,20 @@ class State {
         });
     }
 
-    constructor(rootObject: StateObject, params?: { globalContext?: unknown, defaultCellState?: unknown }) {
+    constructor(rootObject: StateObject, params?: { globalContext?: unknown }) {
         const tree = this._tree;
         const root = tree.root;
-        const defaultCellState = (params && params.defaultCellState) || { }
 
         this.cells.set(root.ref, {
             ref: root.ref,
             obj: rootObject,
             status: 'ok',
             version: root.version,
-            state: { ...defaultCellState }
+            state: { ...StateObjectCell.DefaultState }
         });
 
         this.context = new StateContext({
-            globalContext: params && params.globalContext,
-            defaultCellState
+            globalContext: params && params.globalContext
         });
     }
 }
@@ -102,8 +101,7 @@ namespace State {
     export type Cells = Map<Transform.Ref, StateObjectCell>
 
     export interface Snapshot {
-        readonly tree: StateTree.Serialized,
-        readonly props: { [key: string]: unknown }
+        readonly tree: StateTree.Serialized
     }
 
     export function create(rootObject: StateObject, params?: { globalContext?: unknown, defaultObjectProps?: unknown }) {
@@ -133,7 +131,8 @@ namespace State {
         // TODO: handle current object change
     }
 
-        initObjectState(ctx, roots);
+    initCells(ctx, roots);
+    initObjectStatus(ctx, roots);
 
     for (const root of roots) {
         await updateSubtree(ctx, root);
@@ -142,21 +141,16 @@ namespace State {
 
 function findUpdateRoots(cells: State.Cells, tree: StateTree) {
     const findState = { roots: [] as Ref[], cells };
-        ImmutableTree.doPreOrder(tree, tree.root, findState, _findUpdateRoots);
+    StateTree.doPreOrder(tree, tree.root, findState, _findUpdateRoots);
     return findState.roots;
 }
 
 function _findUpdateRoots(n: Transform, _: any, s: { roots: Ref[], cells: Map<Ref, StateObjectCell> }) {
-        if (!s.cells.has(n.ref)) {
+    const cell = s.cells.get(n.ref);
+    if (!cell || cell.version !== n.version) {
         s.roots.push(n.ref);
         return false;
     }
-        const o = s.cells.get(n.ref)!;
-        if (o.version !== n.version) {
-            s.roots.push(n.ref);
-            return false;
-        }
-
     return true;
 }
 
@@ -166,37 +160,52 @@ namespace State {
 }
 function findDeletes(ctx: UpdateContext): Ref[] {
     const deleteCtx: FindDeletesCtx = { newTree: ctx.tree, cells: ctx.cells, deletes: [] };
-        ImmutableTree.doPostOrder(ctx.oldTree, ctx.oldTree.root, deleteCtx, _visitCheckDelete);
+    StateTree.doPostOrder(ctx.oldTree, ctx.oldTree.root, deleteCtx, _visitCheckDelete);
     return deleteCtx.deletes;
 }
 
-    function setObjectState(ctx: UpdateContext, ref: Ref, status: StateObjectCell.Status, errorText?: string) {
-        let changed = false;
-        if (ctx.cells.has(ref)) {
+function setObjectStatus(ctx: UpdateContext, ref: Ref, status: StateObjectCell.Status, errorText?: string) {
     const obj = ctx.cells.get(ref)!;
-            changed = obj.status !== status;
+    const changed = obj.status !== status;
     obj.status = status;
     obj.errorText = errorText;
-        } else {
-            const obj: StateObjectCell = { ref, status, version: UUID.create(), errorText, state: { ...ctx.stateCtx.defaultCellState } };
-            ctx.cells.set(ref, obj);
-            changed = true;
-        }
     if (changed) ctx.stateCtx.events.object.stateChanged.next({ ref });
 }
 
-    function _initVisitor(t: Transform, _: any, ctx: UpdateContext) {
-        setObjectState(ctx, t.ref, 'pending');
+function _initObjectStatusVisitor(t: Transform, _: any, ctx: UpdateContext) {
+    setObjectStatus(ctx, t.ref, 'pending');
 }
+
 /** Return "resolve set" */
-    function initObjectState(ctx: UpdateContext, roots: Ref[]) {
+function initObjectStatus(ctx: UpdateContext, roots: Ref[]) {
     for (const root of roots) {
-            ImmutableTree.doPreOrder(ctx.tree, ctx.tree.nodes.get(root), ctx, _initVisitor);
+        StateTree.doPreOrder(ctx.tree, ctx.tree.nodes.get(root), ctx, _initObjectStatusVisitor);
+    }
+}
+
+function _initCellsVisitor(transform: Transform, _: any, ctx: UpdateContext) {
+    if (ctx.cells.has(transform.ref)) return;
+
+    const obj: StateObjectCell = {
+        ref: transform.ref,
+        status: 'pending',
+        version: UUID.create(),
+        errorText: void 0,
+        state: { ...StateObjectCell.DefaultState, ...transform.cellState }
+    };
+    ctx.cells.set(transform.ref, obj);
+
+    // TODO: created event???
+}
+
+function initCells(ctx: UpdateContext, roots: Ref[]) {
+    for (const root of roots) {
+        StateTree.doPreOrder(ctx.tree, ctx.tree.nodes.get(root), ctx, _initCellsVisitor);
     }
 }
 
 function doError(ctx: UpdateContext, ref: Ref, errorText: string) {
-        setObjectState(ctx, ref, 'error', errorText);
+    setObjectStatus(ctx, ref, 'error', errorText);
     const wrap = ctx.cells.get(ref)!;
     if (wrap.obj) {
         ctx.stateCtx.events.object.removed.next({ ref });
@@ -225,11 +234,11 @@ namespace State {
 }
 
 async function updateSubtree(ctx: UpdateContext, root: Ref) {
-        setObjectState(ctx, root, 'processing');
+    setObjectStatus(ctx, root, 'processing');
 
     try {
         const update = await updateNode(ctx, root);
-            setObjectState(ctx, root, 'ok');
+        setObjectStatus(ctx, root, 'ok');
         if (update.action === 'created') {
             ctx.stateCtx.events.object.created.next({ ref: root, obj: update.obj! });
         } else if (update.action === 'updated') {
@@ -258,16 +267,12 @@ namespace State {
     if (!oldTree.nodes.has(currentRef) || !cells.has(currentRef)) {
         // console.log('creating...', transform.transformer.id, oldTree.nodes.has(currentRef), objects.has(currentRef));
         const obj = await createObject(ctx, currentRef, transform.transformer, parent, transform.params);
-            cells.set(currentRef, {
-                ref: currentRef,
-                obj,
-                status: 'ok',
-                version: transform.version,
-                state: { ...ctx.stateCtx.defaultCellState, ...transform.cellState }
-            });
+        const cell = cells.get(currentRef)!;
+        cell.obj = obj;
+        cell.version = transform.version;
+
         return { action: 'created', obj };
     } else {
-            // console.log('updating...', transform.transformer.id);
         const current = cells.get(currentRef)!;
         const oldParams = oldTree.nodes.get(currentRef)!.params;
 
@@ -277,29 +282,23 @@ namespace State {
 
         switch (updateKind) {
             case Transformer.UpdateResult.Recreate: {
-                    const obj = await createObject(ctx, currentRef, transform.transformer, parent, transform.params);
-                    cells.set(currentRef, {
-                        ref: currentRef,
-                        obj,
-                        status: 'ok',
-                        version: transform.version,
-                        state: { ...ctx.stateCtx.defaultCellState, ...current.state, ...transform.cellState }
-                    });
-                    return { action: 'replaced', oldObj: current.obj!, newObj: obj };
+                const oldObj = current.obj;
+                const newObj = await createObject(ctx, currentRef, transform.transformer, parent, transform.params);
+                current.obj = newObj;
+                current.version = transform.version;
+                return { action: 'replaced', oldObj, newObj: newObj };
             }
             case Transformer.UpdateResult.Updated:
                 current.version = transform.version;
-                    current.state = { ...ctx.stateCtx.defaultCellState, ...current.state, ...transform.cellState };
                 return { action: 'updated', obj: current.obj };
             default:
-                    // TODO check if props need to be updated
                 return { action: 'none' };
         }
     }
 }
 
 function runTask<T>(t: T | Task<T>, ctx: RuntimeContext) {
-        if (typeof (t as any).run === 'function') return (t as Task<T>).runInContext(ctx);
+    if (typeof (t as any).runInContext === 'function') return (t as Task<T>).runInContext(ctx);
     return t as T;
 }
 

src/mol-state/transform.ts

@@ -40,8 +40,8 @@ export namespace Transform {
         return { ...t, params, version: UUID.create() };
     }
 
-    export function createRoot(ref: Ref): Transform {
-        return create(RootRef, Transformer.ROOT, {}, { ref });
+    export function createRoot(): Transform {
+        return create(RootRef, Transformer.ROOT, {}, { ref: RootRef });
     }
 
     export interface Serialized {

src/mol-state/tree.ts

@@ -4,78 +4,7 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import { Transform } from './transform';
-import { ImmutableTree, TransientTree } from './tree/base';
-import { Transformer } from './transformer';
-import { StateObject } from './object';
+import { StateTree } from './tree/immutable';
+import { TransientTree } from './tree/transient';
 
-export { StateTree, ImmutableTree, TransientTree }
-
-interface StateTree extends ImmutableTree { }
-
-namespace StateTree {
-    export interface Transient extends TransientTree { }
-    export interface Serialized extends ImmutableTree.Serialized { }
-
-    export function create() {
-        return ImmutableTree.create(Transform.createRoot(Transform.RootRef));
-    }
-
-    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();
-    }
-
-    export function toJSON(tree: StateTree) {
-        return ImmutableTree.toJSON(tree) as Serialized;
-    }
-
-    export function fromJSON(data: Serialized): StateTree {
-        return ImmutableTree.fromJSON(data);
-    }
-
-    export interface Builder {
-        getTree(): StateTree
-    }
-
-    export function build(tree: StateTree) {
-        return new Builder.Root(tree);
-    }
-
-    export namespace Builder {
-        interface State {
-            tree: StateTree.Transient
-        }
-
-        export class Root implements Builder {
-            private state: State;
-            to<A extends StateObject>(ref: Transform.Ref) { return new To<A>(this.state, ref, this); }
-            toRoot<A extends StateObject>() { return new To<A>(this.state, this.state.tree.root.ref, this); }
-            delete(ref: Transform.Ref) {
-                this.state.tree.remove(ref);
-                return this;
-            }
-            getTree(): StateTree { return this.state.tree.asImmutable(); }
-            constructor(tree: StateTree) { this.state = { tree: tree.asTransient() } }
-        }
-
-        export class To<A extends StateObject> implements Builder {
-            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);
-                return new To(this.state, t.ref, this.root);
-            }
-
-            and() { return this.root; }
-
-            getTree(): StateTree { return this.state.tree.asImmutable(); }
-
-            constructor(private state: State, private ref: Transform.Ref, private root: Root) {
-                if (!this.state.tree.nodes.has(ref)) {
-                    throw new Error(`Could not find node '${ref}'.`);
-                }
-            }
-        }
-    }
-}
\ No newline at end of file
+export { StateTree, TransientTree }
\ No newline at end of file

src/mol-state/tree/builder.ts

+/**
+ * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { StateTree } from './immutable';
+import { TransientTree } from './transient';
+import { StateObject } from '../object';
+import { Transform } from '../transform';
+import { Transformer } from '../transformer';
+
+export { StateTreeBuilder }
+
+interface StateTreeBuilder {
+    getTree(): StateTree
+}
+
+namespace StateTreeBuilder {
+    interface State {
+        tree: TransientTree
+    }
+
+    export class Root implements StateTreeBuilder {
+        private state: State;
+        to<A extends StateObject>(ref: Transform.Ref) { return new To<A>(this.state, ref, this); }
+        toRoot<A extends StateObject>() { return new To<A>(this.state, this.state.tree.root.ref, this); }
+        delete(ref: Transform.Ref) {
+            this.state.tree.remove(ref);
+            return this;
+        }
+        getTree(): StateTree { return this.state.tree.asImmutable(); }
+        constructor(tree: StateTree) { this.state = { tree: tree.asTransient() } }
+    }
+
+    export class To<A extends StateObject> implements StateTreeBuilder {
+        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);
+            return new To(this.state, t.ref, this.root);
+        }
+
+        and() { return this.root; }
+
+        getTree(): StateTree { return this.state.tree.asImmutable(); }
+
+        constructor(private state: State, private ref: Transform.Ref, private root: Root) {
+            if (!this.state.tree.nodes.has(ref)) {
+                throw new Error(`Could not find node '${ref}'.`);
+            }
+        }
+    }
+}
\ No newline at end of file

src/mol-state/tree/base.ts → src/mol-state/tree/immutable.ts

@@ -6,22 +6,25 @@
 
 import { Map as ImmutableMap, OrderedSet } from 'immutable';
 import { Transform } from '../transform';
+import { TransientTree } from './transient';
+import { StateTreeBuilder } from './builder';
+import { Transformer } from '../transformer';
 
-export { ImmutableTree, TransientTree }
+export { StateTree }
 
 /**
  * An immutable tree where each node requires a unique reference.
  * Represented as an immutable map.
  */
-interface ImmutableTree {
-    readonly version: number,
+interface StateTree {
     readonly root: Transform,
-    readonly nodes: ImmutableTree.Nodes,
-    readonly children: ImmutableTree.Children,
-    asTransient(): TransientTree
+    readonly nodes: StateTree.Nodes,
+    readonly children: StateTree.Children,
+    asTransient(): TransientTree,
+    build(): StateTreeBuilder.Root
 }
 
-namespace ImmutableTree {
+namespace StateTree {
     type Ref = Transform.Ref
 
     export interface ChildSet {
@@ -36,34 +39,45 @@ namespace ImmutableTree {
     export type Nodes = ImmutableMap<Ref, Transform>
     export type Children = ImmutableMap<Ref, ChildSet>
 
-    class Impl implements ImmutableTree {
+    class Impl implements StateTree {
         get root() { return this.nodes.get(Transform.RootRef)! }
 
         asTransient(): TransientTree {
             return new TransientTree(this);
         }
 
-        constructor(public nodes: ImmutableTree.Nodes, public children: Children, public version: number) {
+        build(): StateTreeBuilder.Root {
+            return new StateTreeBuilder.Root(this);
+        }
+
+        constructor(public nodes: StateTree.Nodes, public children: Children) {
         }
     }
 
     /**
      * Create an instance of an immutable tree.
      */
-    export function create(root: Transform): ImmutableTree {
-        return new Impl(ImmutableMap([[root.ref, root]]), ImmutableMap([[root.ref, OrderedSet()]]), 0);
+    export function createEmpty(): StateTree {
+        const root = Transform.createRoot();
+        return create(ImmutableMap([[root.ref, root]]), ImmutableMap([[root.ref, OrderedSet()]]));
+    }
+
+    export function create(nodes: Nodes, children: Children): StateTree {
+        return new Impl(nodes, children);
     }
 
-    export function construct(nodes: Nodes, children: Children, version: number): ImmutableTree {
-        return new Impl(nodes, children, version);
+    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: ImmutableTree, state: any, f: (node: Node, tree: ImmutableTree, state: any) => boolean | undefined | void };
+    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!)); }
     function _doPostOrder(ctx: VisitorCtx, root: Node) {
         const children = ctx.tree.children.get(root.ref);
-        if (checkSetRef && children.size) {
+        if (children && children.size) {
             children.forEach(_postOrderFunc, ctx);
         }
         ctx.f(root, ctx.tree, ctx.state);
@@ -72,7 +86,7 @@ namespace ImmutableTree {
     /**
      * Visit all nodes in a subtree in "post order", meaning leafs get visited first.
      */
-    export function doPostOrder<S>(tree: ImmutableTree, root: Node, state: S, f: (node: Node, tree: ImmutableTree, state: S) => boolean | undefined | void): S {
+    export function doPostOrder<S>(tree: StateTree, root: Node, state: S, f: (node: Node, tree: StateTree, state: S) => boolean | undefined | void): S {
         const ctx: VisitorCtx = { tree, state, f };
         _doPostOrder(ctx, root);
         return ctx.state;
@@ -83,7 +97,7 @@ namespace ImmutableTree {
         const ret = ctx.f(root, ctx.tree, ctx.state);
         if (typeof ret === 'boolean' && !ret) return;
         const children = ctx.tree.children.get(root.ref);
-        if (checkSetRef && children.size) {
+        if (children && children.size) {
             children.forEach(_preOrderFunc, ctx);
         }
     }
@@ -92,7 +106,7 @@ namespace ImmutableTree {
      * Visit all nodes in a subtree in "pre order", meaning leafs get visited last.
      * If the visitor function returns false, the visiting for that branch is interrupted.
      */
-    export function doPreOrder<S>(tree: ImmutableTree, root: Node, state: S, f: (node: Node, tree: ImmutableTree, state: S) => boolean | undefined | void): S {
+    export function doPreOrder<S>(tree: StateTree, root: Node, state: S, f: (node: Node, tree: StateTree, state: S) => boolean | undefined | void): S {
         const ctx: VisitorCtx = { tree, state, f };
         _doPreOrder(ctx, root);
         return ctx.state;
@@ -102,177 +116,49 @@ namespace ImmutableTree {
     /**
      * Get all nodes in a subtree, leafs come first.
      */
-    export function subtreePostOrder<T>(tree: ImmutableTree, root: Node) {
+    export function subtreePostOrder<T>(tree: StateTree, root: Node) {
         return doPostOrder<Node[]>(tree, root, [], _subtree);
     }
 
 
-    function _visitChildToJson(this: Ref[], ref: Ref) { this.push(ref); }
-    interface ToJsonCtx { tree: ImmutableTree, nodes: [Transform.Serialized, any[]][] }
-    function _visitNodeToJson(this: ToJsonCtx, node: Node) {
-        const children: Ref[] = [];
-        this.tree.children.get(node.ref).forEach(_visitChildToJson as any, children);
-        this.nodes.push([Transform.toJSON(node), children]);
+    // function _visitChildToJson(this: Ref[], ref: Ref) { this.push(ref); }
+    // interface ToJsonCtx { nodes: Transform.Serialized[] }
+    function _visitNodeToJson(node: Node, tree: StateTree, ctx: Transform.Serialized[]) {
+        // const children: Ref[] = [];
+        // tree.children.get(node.ref).forEach(_visitChildToJson as any, children);
+        ctx.push(Transform.toJSON(node));
     }
 
     export interface Serialized {
-        nodes: [any /** value */, number[] /** children indices */][]
+        /** Tree nodes serialized in pre-order */
+        nodes: Transform.Serialized[]
     }
 
-    export function toJSON<T>(tree: ImmutableTree): Serialized {
-        const ctx: ToJsonCtx = { tree, nodes: [] };
-
-        tree.nodes.forEach(_visitNodeToJson as any, ctx);
-
-        const map = new Map<string, number>();
-        let i = 0;
-        for (const n of ctx.nodes) map.set(n[0].ref, i++);
-
-        for (const n of ctx.nodes) {
-            const children = n[1];
-            for (i = 0; i < children.length; i++) {
-                children[i] = map.get(children[i]);
-            }
-        }
-        return {
-            nodes: ctx.nodes
-        };
+    export function toJSON<T>(tree: StateTree): Serialized {
+        const nodes: Transform.Serialized[] = [];
+        doPreOrder(tree, tree.root, nodes, _visitNodeToJson);
+        return { nodes };
     }
 
-    export function fromJSON<T>(data: Serialized): ImmutableTree {
+    export function fromJSON<T>(data: Serialized): StateTree {
         const nodes = ImmutableMap<Ref, Node>().asMutable();
         const children = ImmutableMap<Ref, OrderedSet<Ref>>().asMutable();
 
-        const values = data.nodes.map(n => Transform.fromJSON(n[0]));
-        let i = 0;
-        for (const value of values) {
-            const node = data.nodes[i++];
-            const ref = value.ref;
-            nodes.set(ref, value);
-            children.set(ref, OrderedSet(node[1].map(c => values[c].ref)));
-        }
-        return new Impl(nodes.asImmutable(), children.asImmutable(), 0);
-    }
-}
-
-class TransientTree implements ImmutableTree {
-    nodes = this.tree.nodes.asMutable();
-    children = this.tree.children.asMutable();
+        for (const s of data.nodes) {
+            const node = Transform.fromJSON(s);
+            nodes.set(node.ref, node);
 
-    version: number = this.tree.version + 1;
-
-    private changedValue = false;
-    private mutations: Map<Transform.Ref, OrderedSet<Transform.Ref>> = new Map();
-
-    get root() { return this.nodes.get(Transform.RootRef)! }
-
-    asTransient() {
-        return this.asImmutable().asTransient();
-    }
-
-    private addChild(parent: Transform.Ref, child: Transform.Ref) {
-        if (this.mutations.has(parent)) {
-            this.mutations.get(parent)!.add(child);
-        } else {
-            const set = (this.children.get(parent) as OrderedSet<Transform.Ref>).asMutable();
-            set.add(child);
-            this.children.set(parent, set);
-            this.mutations.set(parent, set);
-        }
+            if (!children.has(node.ref)) {
+                children.set(node.ref, OrderedSet<Ref>().asMutable());
             }
 
-    private removeChild(parent: Transform.Ref, child: Transform.Ref) {
-        if (this.mutations.has(parent)) {
-            this.mutations.get(parent)!.remove(child);
-        } else {
-            const set = (this.children.get(parent) as OrderedSet<Transform.Ref>).asMutable();
-            set.remove(child);
-            this.children.set(parent, set);
-            this.mutations.set(parent, set);
-        }
-    }
-
-    private clearRoot() {
-        const parent = Transform.RootRef;
-        if (this.children.get(parent).size === 0) return;
-        const set = OrderedSet<Transform.Ref>();
-        this.children.set(parent, set);
-        this.mutations.set(parent, set);
-    }
-
-    add(transform: Transform) {
-        const ref = transform.ref;
-
-        const children = this.children.get(transform.parent);
-        if (!children) parentNotPresent(transform.parent);
-
-        if (!children.has(transform.ref)) {
-            this.addChild(transform.parent, transform.ref);
-        }
-
-        if (!this.children.has(transform.ref)) {
-            this.children.set(transform.ref, OrderedSet());
-        }
-
-        this.nodes.set(ref, transform);
-        return this;
-    }
-
-    set(transform: Transform) {
-        ensurePresent(this.nodes, transform.ref);
-        this.changedValue = true;
-        this.nodes.set(transform.ref, transform);
-        return this;
-    }
-
-    remove(ref: Transform.Ref): Transform[] {
-        const { nodes, mutations, children } = this;
-        const node = nodes.get(ref);
-        if (!node) return [];
-
-        const st = ImmutableTree.subtreePostOrder(this, node);
-        if (ref === Transform.RootRef) {
-            st.pop();
-            this.clearRoot();
-        } else {
-            this.removeChild(node.parent, node.ref);
-        }
-
-        for (const n of st) {
-            nodes.delete(n.ref);
-            children.delete(n.ref);
-            mutations.delete(n.ref);
-        }
-
-        return st;
-    }
-
-    asImmutable() {
-        if (!this.changedValue && this.mutations.size === 0) return this.tree;
-
-        this.mutations.forEach(fixChildMutations, this.children);
-        return ImmutableTree.construct(this.nodes.asImmutable(), this.children.asImmutable(), this.version);
-    }
-
-    constructor(private tree: ImmutableTree) {
-
-    }
-}
-
-function fixChildMutations(this: ImmutableTree.Children, m: OrderedSet<Transform.Ref>, k: Transform.Ref) { this.set(k, m.asImmutable()); }
-
-function checkSetRef(oldRef: Transform.Ref, newRef: Transform.Ref) {
-    if (oldRef !== newRef) {
-        throw new Error(`Cannot setValue of node '${oldRef}' because the new value has a different ref '${newRef}'.`);
-    }
+            if (node.ref !== node.parent) children.get(node.parent).add(node.ref);
         }
 
-function parentNotPresent(ref: Transform.Ref) {
-    throw new Error(`Parent '${ref}' must be present in the tree.`);
+        for (const s of data.nodes) {
+            children.set(s.ref, children.get(s.ref).asImmutable());
         }
 
-function ensurePresent(nodes: ImmutableTree.Nodes, ref: Transform.Ref) {
-    if (!nodes.has(ref)) {
-        throw new Error(`Node '${ref}' is not present in the tree.`);
+        return create(nodes.asImmutable(), children.asImmutable());
     }
 }
\ No newline at end of file

src/mol-state/tree/transient.ts

+/**
+ * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+import { OrderedSet } from 'immutable';
+import { Transform } from '../transform';
+import { StateTree } from './immutable';
+import { StateTreeBuilder } from './builder';
+
+export { TransientTree }
+
+class TransientTree implements StateTree {
+    nodes = this.tree.nodes;
+    children = this.tree.children;
+
+    private changedNodes = false;
+    private changedChildren = false;
+    private _mutations: Map<Transform.Ref, OrderedSet<Transform.Ref>> = void 0 as any;
+
+    private get mutations() {
+        if (this._mutations) return this._mutations;
+        this._mutations = new Map();
+        return this._mutations;
+    }
+
+    get root() { return this.nodes.get(Transform.RootRef)! }
+
+    build(): StateTreeBuilder.Root {
+        return new StateTreeBuilder.Root(this);
+    }
+
+    asTransient() {
+        return this.asImmutable().asTransient();
+    }
+
+    private addChild(parent: Transform.Ref, child: Transform.Ref) {
+        if (!this.changedChildren) {
+            this.changedChildren = true;
+            this.children = this.children.asMutable();
+        }
+
+        if (this.mutations.has(parent)) {
+            this.mutations.get(parent)!.add(child);
+        } else {
+            const set = (this.children.get(parent) as OrderedSet<Transform.Ref>).asMutable();
+            set.add(child);
+            this.children.set(parent, set);
+            this.mutations.set(parent, set);
+        }
+    }
+
+    private removeChild(parent: Transform.Ref, child: Transform.Ref) {
+        if (!this.changedChildren) {
+            this.changedChildren = true;
+            this.children = this.children.asMutable();
+        }
+
+        if (this.mutations.has(parent)) {
+            this.mutations.get(parent)!.remove(child);
+        } else {
+            const set = (this.children.get(parent) as OrderedSet<Transform.Ref>).asMutable();
+            set.remove(child);
+            this.children.set(parent, set);
+            this.mutations.set(parent, set);
+        }
+    }
+
+    private clearRoot() {
+        const parent = Transform.RootRef;
+        if (this.children.get(parent).size === 0) return;
+        const set = OrderedSet<Transform.Ref>();
+        this.children.set(parent, set);
+        this.mutations.set(parent, set);
+    }
+
+    add(transform: Transform) {
+        const ref = transform.ref;
+
+        if (this.nodes.has(transform.ref)) {
+            const node = this.nodes.get(transform.ref);
+            if (node.parent !== transform.parent) alreadyPresent(transform.ref);
+        }
+
+        const children = this.children.get(transform.parent);
+        if (!children) parentNotPresent(transform.parent);
+
+        if (!children.has(transform.ref)) {
+            this.addChild(transform.parent, transform.ref);
+        }
+
+        if (!this.children.has(transform.ref)) {
+            if (!this.changedChildren) {
+                this.changedChildren = true;
+                this.children = this.children.asMutable();
+            }
+            this.children.set(transform.ref, OrderedSet());
+        }
+
+        if (!this.changedNodes) {
+            this.changedNodes = true;
+            this.nodes = this.nodes.asMutable();
+        }
+
+        this.nodes.set(ref, transform);
+        return this;
+    }
+
+    set(transform: Transform) {
+        ensurePresent(this.nodes, transform.ref);
+
+        if (!this.changedNodes) {
+            this.changedNodes = true;
+            this.nodes = this.nodes.asMutable();
+        }
+
+        this.nodes.set(transform.ref, transform);
+        return this;
+    }
+
+    remove(ref: Transform.Ref): Transform[] {
+        const node = this.nodes.get(ref);
+        if (!node) return [];
+
+        const st = StateTree.subtreePostOrder(this, node);
+        if (ref === Transform.RootRef) {
+            st.pop();
+            if (st.length === 0) return st;
+            this.clearRoot();
+        } else {
+            if (st.length === 0) return st;
+            this.removeChild(node.parent, node.ref);
+        }
+
+        if (!this.changedNodes) {
+            this.changedNodes = true;
+            this.nodes = this.nodes.asMutable();
+        }
+
+        for (const n of st) {
+            this.nodes.delete(n.ref);
+            this.children.delete(n.ref);
+            if (this._mutations) this._mutations.delete(n.ref);
+        }
+
+        return st;
+    }
+
+    asImmutable() {
+        if (!this.changedNodes && !this.changedChildren && !this._mutations) return this.tree;
+        if (this._mutations) this._mutations.forEach(fixChildMutations, this.children);
+        return StateTree.create(
+            this.changedNodes ? this.nodes.asImmutable() : this.nodes,
+            this.changedChildren ? this.children.asImmutable() : this.children);
+    }
+
+    constructor(private tree: StateTree) {
+
+    }
+}
+
+function fixChildMutations(this: StateTree.Children, m: OrderedSet<Transform.Ref>, k: Transform.Ref) { this.set(k, m.asImmutable()); }
+
+function alreadyPresent(ref: Transform.Ref) {
+    throw new Error(`Transform '${ref}' is already present in the tree.`);
+}
+
+function parentNotPresent(ref: Transform.Ref) {
+    throw new Error(`Parent '${ref}' must be present in the tree.`);
+}
+
+function ensurePresent(nodes: StateTree.Nodes, ref: Transform.Ref) {
+    if (!nodes.has(ref)) {
+        throw new Error(`Node '${ref}' is not present in the tree.`);
+    }
+}
\ No newline at end of file

src/perf-tests/state.ts

@@ -77,7 +77,7 @@ export async function testState() {
     hookEvents(state);
 
     const tree = state.tree;
-    const builder = StateTree.build(tree);
+    const builder = tree.build();
     builder.toRoot<Root>()
         .apply(CreateSquare, { a: 10 }, { ref: 'square' })
         .apply(CaclArea);