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 });
     }
 }

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
+    };
 }

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 }) {
@@ -111,212 +109,213 @@ namespace State {
     }
 }
 
-    type Ref = Transform.Ref
+type Ref = Transform.Ref
 
-    interface UpdateContext {
-        stateCtx: StateContext,
-        taskCtx: RuntimeContext,
-        oldTree: StateTree,
-        tree: StateTree,
-        cells: State.Cells,
-        transformCache: Map<Ref, unknown>
-    }
+interface UpdateContext {
+    stateCtx: StateContext,
+    taskCtx: RuntimeContext,
+    oldTree: StateTree,
+    tree: StateTree,
+    cells: State.Cells,
+    transformCache: Map<Ref, unknown>
+}
 
-    async function update(ctx: UpdateContext) {
-        const roots = findUpdateRoots(ctx.cells, ctx.tree);
-        const deletes = findDeletes(ctx);
-        for (const d of deletes) {
-            const obj = ctx.cells.has(d) ? ctx.cells.get(d)!.obj : void 0;
-            ctx.cells.delete(d);
-            ctx.transformCache.delete(d);
-            ctx.stateCtx.events.object.removed.next({ ref: d, obj });
-            // TODO: handle current object change
-        }
+async function update(ctx: UpdateContext) {
+    const roots = findUpdateRoots(ctx.cells, ctx.tree);
+    const deletes = findDeletes(ctx);
+    for (const d of deletes) {
+        const obj = ctx.cells.has(d) ? ctx.cells.get(d)!.obj : void 0;
+        ctx.cells.delete(d);
+        ctx.transformCache.delete(d);
+        ctx.stateCtx.events.object.removed.next({ ref: d, obj });
+        // TODO: handle current object change
+    }
 
-        initObjectState(ctx, roots);
+    initCells(ctx, roots);
+    initObjectStatus(ctx, roots);
 
-        for (const root of roots) {
-            await updateSubtree(ctx, root);
-        }
+    for (const root of roots) {
+        await updateSubtree(ctx, root);
     }
+}
 
-    function findUpdateRoots(cells: State.Cells, tree: StateTree) {
-        const findState = { roots: [] as Ref[], cells };
-        ImmutableTree.doPreOrder(tree, tree.root, findState, _findUpdateRoots);
-        return findState.roots;
+function findUpdateRoots(cells: State.Cells, tree: StateTree) {
+    const findState = { roots: [] as Ref[], cells };
+    StateTree.doPreOrder(tree, tree.root, findState, _findUpdateRoots);
+    return findState.roots;
+}
+
+function _findUpdateRoots(n: Transform, _: any, s: { roots: Ref[], cells: Map<Ref, StateObjectCell> }) {
+    const cell = s.cells.get(n.ref);
+    if (!cell || cell.version !== n.version) {
+        s.roots.push(n.ref);
+        return false;
     }
+    return true;
+}
 
-    function _findUpdateRoots(n: Transform, _: any, s: { roots: Ref[], cells: Map<Ref, StateObjectCell> }) {
-        if (!s.cells.has(n.ref)) {
-            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;
-        }
+type FindDeletesCtx = { newTree: StateTree, cells: State.Cells, deletes: Ref[] }
+function _visitCheckDelete(n: Transform, _: any, ctx: FindDeletesCtx) {
+    if (!ctx.newTree.nodes.has(n.ref) && ctx.cells.has(n.ref)) ctx.deletes.push(n.ref);
+}
+function findDeletes(ctx: UpdateContext): Ref[] {
+    const deleteCtx: FindDeletesCtx = { newTree: ctx.tree, cells: ctx.cells, deletes: [] };
+    StateTree.doPostOrder(ctx.oldTree, ctx.oldTree.root, deleteCtx, _visitCheckDelete);
+    return deleteCtx.deletes;
+}
 
-        return true;
-    }
+function setObjectStatus(ctx: UpdateContext, ref: Ref, status: StateObjectCell.Status, errorText?: string) {
+    const obj = ctx.cells.get(ref)!;
+    const changed = obj.status !== status;
+    obj.status = status;
+    obj.errorText = errorText;
+    if (changed) ctx.stateCtx.events.object.stateChanged.next({ ref });
+}
 
-    type FindDeletesCtx = { newTree: StateTree, cells: State.Cells, deletes: Ref[] }
-    function _visitCheckDelete(n: Transform, _: any, ctx: FindDeletesCtx) {
-        if (!ctx.newTree.nodes.has(n.ref) && ctx.cells.has(n.ref)) ctx.deletes.push(n.ref);
-    }
-    function findDeletes(ctx: UpdateContext): Ref[] {
-        const deleteCtx: FindDeletesCtx = { newTree: ctx.tree, cells: ctx.cells, deletes: [] };
-        ImmutableTree.doPostOrder(ctx.oldTree, ctx.oldTree.root, deleteCtx, _visitCheckDelete);
-        return deleteCtx.deletes;
-    }
+function _initObjectStatusVisitor(t: Transform, _: any, ctx: UpdateContext) {
+    setObjectStatus(ctx, t.ref, 'pending');
+}
 
-    function setObjectState(ctx: UpdateContext, ref: Ref, status: StateObjectCell.Status, errorText?: string) {
-        let changed = false;
-        if (ctx.cells.has(ref)) {
-            const obj = ctx.cells.get(ref)!;
-            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 });
+/** Return "resolve set" */
+function initObjectStatus(ctx: UpdateContext, roots: Ref[]) {
+    for (const root of roots) {
+        StateTree.doPreOrder(ctx.tree, ctx.tree.nodes.get(root), ctx, _initObjectStatusVisitor);
     }
+}
 
-    function _initVisitor(t: Transform, _: any, ctx: UpdateContext) {
-        setObjectState(ctx, t.ref, 'pending');
-    }
-    /** Return "resolve set" */
-    function initObjectState(ctx: UpdateContext, roots: Ref[]) {
-        for (const root of roots) {
-            ImmutableTree.doPreOrder(ctx.tree, ctx.tree.nodes.get(root), ctx, _initVisitor);
-        }
+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);
-        const wrap = ctx.cells.get(ref)!;
-        if (wrap.obj) {
-            ctx.stateCtx.events.object.removed.next({ ref });
-            ctx.transformCache.delete(ref);
-            wrap.obj = void 0;
-        }
+function doError(ctx: UpdateContext, ref: Ref, errorText: string) {
+    setObjectStatus(ctx, ref, 'error', errorText);
+    const wrap = ctx.cells.get(ref)!;
+    if (wrap.obj) {
+        ctx.stateCtx.events.object.removed.next({ ref });
+        ctx.transformCache.delete(ref);
+        wrap.obj = void 0;
+    }
 
-        const children = ctx.tree.children.get(ref).values();
-        while (true) {
-            const next = children.next();
-            if (next.done) return;
-            doError(ctx, next.value, 'Parent node contains error.');
-        }
+    const children = ctx.tree.children.get(ref).values();
+    while (true) {
+        const next = children.next();
+        if (next.done) return;
+        doError(ctx, next.value, 'Parent node contains error.');
     }
+}
 
-    function findAncestor(tree: StateTree, cells: State.Cells, root: Ref, types: { type: StateObject.Type }[]): StateObject {
-        let current = tree.nodes.get(root)!;
-        while (true) {
-            current = tree.nodes.get(current.parent)!;
-            if (current.ref === Transform.RootRef) {
-                return cells.get(Transform.RootRef)!.obj!;
-            }
-            const obj = cells.get(current.ref)!.obj!;
-            for (const t of types) if (obj.type === t.type) return cells.get(current.ref)!.obj!;
+function findAncestor(tree: StateTree, cells: State.Cells, root: Ref, types: { type: StateObject.Type }[]): StateObject {
+    let current = tree.nodes.get(root)!;
+    while (true) {
+        current = tree.nodes.get(current.parent)!;
+        if (current.ref === Transform.RootRef) {
+            return cells.get(Transform.RootRef)!.obj!;
         }
+        const obj = cells.get(current.ref)!.obj!;
+        for (const t of types) if (obj.type === t.type) return cells.get(current.ref)!.obj!;
     }
+}
 
-    async function updateSubtree(ctx: UpdateContext, root: Ref) {
-        setObjectState(ctx, root, 'processing');
-
-        try {
-            const update = await updateNode(ctx, root);
-            setObjectState(ctx, root, 'ok');
-            if (update.action === 'created') {
-                ctx.stateCtx.events.object.created.next({ ref: root, obj: update.obj! });
-            } else if (update.action === 'updated') {
-                ctx.stateCtx.events.object.updated.next({ ref: root, obj: update.obj });
-            } else if (update.action === 'replaced') {
-                ctx.stateCtx.events.object.replaced.next({ ref: root, oldObj: update.oldObj, newObj: update.newObj });
-            }
-        } catch (e) {
-            doError(ctx, root, '' + e);
-            return;
+async function updateSubtree(ctx: UpdateContext, root: Ref) {
+    setObjectStatus(ctx, root, 'processing');
+
+    try {
+        const update = await updateNode(ctx, root);
+        setObjectStatus(ctx, root, 'ok');
+        if (update.action === 'created') {
+            ctx.stateCtx.events.object.created.next({ ref: root, obj: update.obj! });
+        } else if (update.action === 'updated') {
+            ctx.stateCtx.events.object.updated.next({ ref: root, obj: update.obj });
+        } else if (update.action === 'replaced') {
+            ctx.stateCtx.events.object.replaced.next({ ref: root, oldObj: update.oldObj, newObj: update.newObj });
         }
+    } catch (e) {
+        doError(ctx, root, '' + e);
+        return;
+    }
 
-        const children = ctx.tree.children.get(root).values();
-        while (true) {
-            const next = children.next();
-            if (next.done) return;
-            await updateSubtree(ctx, next.value);
-        }
+    const children = ctx.tree.children.get(root).values();
+    while (true) {
+        const next = children.next();
+        if (next.done) return;
+        await updateSubtree(ctx, next.value);
     }
+}
 
-    async function updateNode(ctx: UpdateContext, currentRef: Ref) {
-        const { oldTree, tree, cells } = ctx;
-        const transform = tree.nodes.get(currentRef);
-        const parent = findAncestor(tree, cells, currentRef, transform.transformer.definition.from);
-        // console.log('parent', transform.transformer.id, transform.transformer.definition.from[0].type, parent ? parent.ref : 'undefined')
-        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 }
-            });
-            return { action: 'created', obj };
-        } else {
-            // console.log('updating...', transform.transformer.id);
-            const current = cells.get(currentRef)!;
-            const oldParams = oldTree.nodes.get(currentRef)!.params;
-
-            const updateKind = current.status === 'ok' || current.ref === Transform.RootRef
-                ? await updateObject(ctx, currentRef, transform.transformer, parent, current.obj!, oldParams, transform.params)
-                : Transformer.UpdateResult.Recreate;
-
-            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 };
-                }
-                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' };
+async function updateNode(ctx: UpdateContext, currentRef: Ref) {
+    const { oldTree, tree, cells } = ctx;
+    const transform = tree.nodes.get(currentRef);
+    const parent = findAncestor(tree, cells, currentRef, transform.transformer.definition.from);
+    // console.log('parent', transform.transformer.id, transform.transformer.definition.from[0].type, parent ? parent.ref : 'undefined')
+    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);
+        const cell = cells.get(currentRef)!;
+        cell.obj = obj;
+        cell.version = transform.version;
+
+        return { action: 'created', obj };
+    } else {
+        const current = cells.get(currentRef)!;
+        const oldParams = oldTree.nodes.get(currentRef)!.params;
+
+        const updateKind = current.status === 'ok' || current.ref === Transform.RootRef
+            ? await updateObject(ctx, currentRef, transform.transformer, parent, current.obj!, oldParams, transform.params)
+            : Transformer.UpdateResult.Recreate;
+
+        switch (updateKind) {
+            case Transformer.UpdateResult.Recreate: {
+                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;
+                return { action: 'updated', obj: current.obj };
+            default:
+                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);
-        return t as T;
-    }
+function runTask<T>(t: T | Task<T>, ctx: RuntimeContext) {
+    if (typeof (t as any).runInContext === 'function') return (t as Task<T>).runInContext(ctx);
+    return t as T;
+}
 
-    function createObject(ctx: UpdateContext, ref: Ref, transformer: Transformer, a: StateObject, params: any) {
-        const cache = { };
+function createObject(ctx: UpdateContext, ref: Ref, transformer: Transformer, a: StateObject, params: any) {
+    const cache = {};
+    ctx.transformCache.set(ref, cache);
+    return runTask(transformer.definition.apply({ a, params, cache }, ctx.stateCtx.globalContext), ctx.taskCtx);
+}
+
+async function updateObject(ctx: UpdateContext, ref: Ref, transformer: Transformer, a: StateObject, b: StateObject, oldParams: any, newParams: any) {
+    if (!transformer.definition.update) {
+        return Transformer.UpdateResult.Recreate;
+    }
+    let cache = ctx.transformCache.get(ref);
+    if (!cache) {
+        cache = {};
         ctx.transformCache.set(ref, cache);
-        return runTask(transformer.definition.apply({ a, params, cache }, ctx.stateCtx.globalContext), ctx.taskCtx);
     }
-
-    async function updateObject(ctx: UpdateContext, ref: Ref, transformer: Transformer, a: StateObject, b: StateObject, oldParams: any, newParams: any) {
-        if (!transformer.definition.update) {
-            return Transformer.UpdateResult.Recreate;
-        }
-        let cache = ctx.transformCache.get(ref);
-        if (!cache) {
-            cache = { };
-            ctx.transformCache.set(ref, cache);
-        }
-        return runTask(transformer.definition.update({ a, oldParams, b, newParams, cache }, ctx.stateCtx.globalContext), ctx.taskCtx);
-    }
+    return runTask(transformer.definition.update({ a, oldParams, b, newParams, cache }, ctx.stateCtx.globalContext), ctx.taskCtx);
+}

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}'.`);
-                }
-            }
-        }
-    }
-}
+export { StateTree, TransientTree }

src/mol-state/tree/base.ts

@@ -1,278 +0,0 @@
-/**
- * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
- *
- * @author David Sehnal <david.sehnal@gmail.com>
- */
-
-import { Map as ImmutableMap, OrderedSet } from 'immutable';
-import { Transform } from '../transform';
-
-export { ImmutableTree, TransientTree }
-
-/**
- * An immutable tree where each node requires a unique reference.
- * Represented as an immutable map.
- */
-interface ImmutableTree {
-    readonly version: number,
-    readonly root: Transform,
-    readonly nodes: ImmutableTree.Nodes,
-    readonly children: ImmutableTree.Children,
-    asTransient(): TransientTree
-}
-
-namespace ImmutableTree {
-    type Ref = Transform.Ref
-
-    export interface ChildSet {
-        readonly size: number,
-        readonly values: OrderedSet<Ref>['values'],
-        has(ref: Ref): boolean,
-        readonly forEach: OrderedSet<Ref>['forEach'],
-        readonly map: OrderedSet<Ref>['map']
-    }
-
-    export type Node = Transform
-    export type Nodes = ImmutableMap<Ref, Transform>
-    export type Children = ImmutableMap<Ref, ChildSet>
-
-    class Impl implements ImmutableTree {
-        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) {
-        }
-    }
-
-    /**
-     * 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 construct(nodes: Nodes, children: Children, version: number): ImmutableTree {
-        return new Impl(nodes, children, version);
-    }
-
-    type VisitorCtx = { tree: ImmutableTree, state: any, f: (node: Node, tree: ImmutableTree, 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) {
-            children.forEach(_postOrderFunc, ctx);
-        }
-        ctx.f(root, ctx.tree, ctx.state);
-    }
-
-    /**
-     * 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 {
-        const ctx: VisitorCtx = { tree, state, f };
-        _doPostOrder(ctx, root);
-        return ctx.state;
-    }
-
-    function _preOrderFunc(this: VisitorCtx, c: Ref | undefined) { _doPreOrder(this, this.tree.nodes.get(c!)); }
-    function _doPreOrder(ctx: VisitorCtx, root: Node) {
-        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) {
-            children.forEach(_preOrderFunc, ctx);
-        }
-    }
-
-    /**
-     * 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 {
-        const ctx: VisitorCtx = { tree, state, f };
-        _doPreOrder(ctx, root);
-        return ctx.state;
-    }
-
-    function _subtree(n: Node, _: any, subtree: Node[]) { subtree.push(n); }
-    /**
-     * Get all nodes in a subtree, leafs come first.
-     */
-    export function subtreePostOrder<T>(tree: ImmutableTree, 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]);
-    }
-
-    export interface Serialized {
-        nodes: [any /** value */, number[] /** children indices */][]
-    }
-
-    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 fromJSON<T>(data: Serialized): ImmutableTree {
-        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();
-
-    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);
-        }
-    }
-
-    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}'.`);
-    }
-}
-
-function parentNotPresent(ref: Transform.Ref) {
-    throw new Error(`Parent '${ref}' must be present in the tree.`);
-}
-
-function ensurePresent(nodes: ImmutableTree.Nodes, ref: Transform.Ref) {
-    if (!nodes.has(ref)) {
-        throw new Error(`Node '${ref}' is not present in the tree.`);
-    }
-}

src/mol-state/tree/builder.ts

@@ -0,0 +1,53 @@
+/**
+ * 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}'.`);
+            }
+        }
+    }
+}

src/mol-state/tree/immutable.ts

@@ -0,0 +1,164 @@
+/**
+ * Copyright (c) 2018 mol* contributors, licensed under MIT, See LICENSE file for more info.
+ *
+ * @author David Sehnal <david.sehnal@gmail.com>
+ */
+
+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 }
+
+/**
+ * An immutable tree where each node requires a unique reference.
+ * Represented as an immutable map.
+ */
+interface StateTree {
+    readonly root: Transform,
+    readonly nodes: StateTree.Nodes,
+    readonly children: StateTree.Children,
+    asTransient(): TransientTree,
+    build(): StateTreeBuilder.Root
+}
+
+namespace StateTree {
+    type Ref = Transform.Ref
+
+    export interface ChildSet {
+        readonly size: number,
+        readonly values: OrderedSet<Ref>['values'],
+        has(ref: Ref): boolean,
+        readonly forEach: OrderedSet<Ref>['forEach'],
+        readonly map: OrderedSet<Ref>['map']
+    }
+
+    export type Node = Transform
+    export type Nodes = ImmutableMap<Ref, Transform>
+    export type Children = ImmutableMap<Ref, ChildSet>
+
+    class Impl implements StateTree {
+        get root() { return this.nodes.get(Transform.RootRef)! }
+
+        asTransient(): TransientTree {
+            return new TransientTree(this);
+        }
+
+        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 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 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!)); }
+    function _doPostOrder(ctx: VisitorCtx, root: Node) {
+        const children = ctx.tree.children.get(root.ref);
+        if (children && children.size) {
+            children.forEach(_postOrderFunc, ctx);
+        }
+        ctx.f(root, ctx.tree, ctx.state);
+    }
+
+    /**
+     * Visit all nodes in a subtree in "post order", meaning leafs get visited first.
+     */
+    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;
+    }
+
+    function _preOrderFunc(this: VisitorCtx, c: Ref | undefined) { _doPreOrder(this, this.tree.nodes.get(c!)); }
+    function _doPreOrder(ctx: VisitorCtx, root: Node) {
+        const ret = ctx.f(root, ctx.tree, ctx.state);
+        if (typeof ret === 'boolean' && !ret) return;
+        const children = ctx.tree.children.get(root.ref);
+        if (children && children.size) {
+            children.forEach(_preOrderFunc, ctx);
+        }
+    }
+
+    /**
+     * 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: 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;
+    }
+
+    function _subtree(n: Node, _: any, subtree: Node[]) { subtree.push(n); }
+    /**
+     * Get all nodes in a subtree, leafs come first.
+     */
+    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 { 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 {
+        /** Tree nodes serialized in pre-order */
+        nodes: Transform.Serialized[]
+    }
+
+    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): StateTree {
+        const nodes = ImmutableMap<Ref, Node>().asMutable();
+        const children = ImmutableMap<Ref, OrderedSet<Ref>>().asMutable();
+
+        for (const s of data.nodes) {
+            const node = Transform.fromJSON(s);
+            nodes.set(node.ref, node);
+
+            if (!children.has(node.ref)) {
+                children.set(node.ref, OrderedSet<Ref>().asMutable());
+            }
+
+            if (node.ref !== node.parent) children.get(node.parent).add(node.ref);
+        }
+
+        for (const s of data.nodes) {
+            children.set(s.ref, children.get(s.ref).asImmutable());
+        }
+
+        return create(nodes.asImmutable(), children.asImmutable());
+    }
+}

src/mol-state/tree/transient.ts

@@ -0,0 +1,177 @@
+/**
+ * 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.`);
+    }
+}

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);