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/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 }
-
+export { StateTree, TransientTree }
\ No newline at end of file
-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

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();
-
-    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);
-        }
+        for (const s of data.nodes) {
+            const node = Transform.fromJSON(s);
+            nodes.set(node.ref, node);
 
-        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 [];
+            if (!children.has(node.ref)) {
+                children.set(node.ref, OrderedSet<Ref>().asMutable());
+            }
 
-        const st = ImmutableTree.subtreePostOrder(this, node);
-        if (ref === Transform.RootRef) {
-            st.pop();
-            this.clearRoot();
-        } else {
-            this.removeChild(node.parent, node.ref);
+            if (node.ref !== node.parent) children.get(node.parent).add(node.ref);
         }
 
-        for (const n of st) {
-            nodes.delete(n.ref);
-            children.delete(n.ref);
-            mutations.delete(n.ref);
+        for (const s of data.nodes) {
+            children.set(s.ref, children.get(s.ref).asImmutable());
         }
 
-        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.`);
+        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);