mol-state: simplify tree

src/mol-plugin/context.ts

@@ -161,7 +161,7 @@ export class PluginContext {
 
     _test_nextModel() {
         const models = StateSelection.select('models', this.state.data)[0].obj as SO.Molecule.Models;
-        const idx = (this.state.data.tree.getValue('structure')!.params as Transformer.Params<typeof StateTransforms.Model.CreateStructureFromModel>).modelIndex;
+        const idx = (this.state.data.tree.nodes.get('structure')!.params as Transformer.Params<typeof StateTransforms.Model.CreateStructureFromModel>).modelIndex;
         const newTree = StateTree.updateParams(this.state.data.tree, 'structure', { modelIndex: (idx + 1) % models.data.length });
         return this.state.updateData(newTree);
         // this.viewer.requestDraw(true);

src/mol-plugin/ui/controls.tsx

@@ -92,8 +92,7 @@ export class _test_CreateTransform extends React.Component<{ plugin: PluginConte
 
 export class _test_UpdateTransform extends React.Component<{ plugin: PluginContext, nodeRef: Transform.Ref }, { params: any }> {
     private getTransform() {
-        const t = this.props.plugin.state.data.tree.nodes.get(this.props.nodeRef)!;
-        return t ? t.value : t;
+        return this.props.plugin.state.data.tree.nodes.get(this.props.nodeRef)!;
     }
 
     private getParamDef() {
@@ -122,7 +121,7 @@ export class _test_UpdateTransform extends React.Component<{ plugin: PluginConte
 
     render() {
         const transform = this.getTransform();
-        if (!transform || transform.ref === this.props.plugin.state.data.tree.rootRef) {
+        if (!transform || transform.ref === Transform.RootRef) {
             return <div />;
         }
 

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

@@ -17,7 +17,7 @@ export class StateTree extends React.Component<{ plugin: PluginContext, state: S
     }
     render() {
         // const n = this.props.plugin.state.data.tree.nodes.get(this.props.plugin.state.data.tree.rootRef)!;
-        const n = this.props.state.tree.rootRef;
+        const n = this.props.state.tree.root.ref;
         return <div>
             <StateTreeNode plugin={this.props.plugin} state={this.props.state} nodeRef={n} key={n} />
             { /* n.children.map(c => <StateTreeNode plugin={this.props.plugin} nodeRef={c!} key={c} />) */}
@@ -37,7 +37,7 @@ export class StateTreeNode extends React.Component<{ plugin: PluginContext, node
 
         let label: any;
         if (cell.status !== 'ok' || !cell.obj) {
-            const name = (n.value.transformer.definition.display && n.value.transformer.definition.display.name) || n.value.transformer.definition.name;
+            const name = (n.transformer.definition.display && n.transformer.definition.display.name) || n.transformer.definition.name;
             label = <><b>{cell.status}</b> <a href='#' onClick={e => {
                 e.preventDefault();
                 PluginCommands.Data.SetCurrentObject.dispatch(this.props.plugin, { ref: this.props.nodeRef });
@@ -52,11 +52,13 @@ export class StateTreeNode extends React.Component<{ plugin: PluginContext, node
             }}>{props.label}</a> {props.description ? <small>{props.description}</small> : void 0}</>;
         }
 
+        const children = this.props.state.tree.children.get(this.props.nodeRef);
+
         return <div>
             {remove}{label}
-            {n.children.size === 0
+            {children.size === 0
                 ? void 0
-                : <div style={{ marginLeft: '7px', paddingLeft: '3px', borderLeft: '1px solid #999' }}>{n.children.map(c => <StateTreeNode plugin={this.props.plugin} state={this.props.state} nodeRef={c!} key={c} />)}</div>
+                : <div style={{ marginLeft: '7px', paddingLeft: '3px', borderLeft: '1px solid #999' }}>{children.map(c => <StateTreeNode plugin={this.props.plugin} state={this.props.state} nodeRef={c!} key={c} />)}</div>
             }
         </div>;
     }

src/mol-state/context.ts

@@ -40,9 +40,9 @@ class StateContext {
         this.ev.dispose();
     }
 
-    constructor(params: { globalContext: unknown, defaultCellState: unknown, rootRef: Transform.Ref }) {
+    constructor(params: { globalContext: unknown, defaultCellState: unknown }) {
         this.globalContext = params.globalContext;
         this.defaultCellState = params.defaultCellState;
-        this.behaviors.currentObject.next({ ref: params.rootRef });
+        this.behaviors.currentObject.next({ ref: Transform.RootRef });
     }
 }

src/mol-state/immutable-tree.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';
-
-/**
- * An immutable tree where each node requires a unique reference.
- * Represented as an immutable map.
- */
-export interface ImmutableTree<T> {
-    readonly rootRef: ImmutableTree.Ref,
-    readonly version: number,
-    readonly nodes: ImmutableTree.Nodes<T>,
-    getRef(e: T): ImmutableTree.Ref,
-    getValue(ref: ImmutableTree.Ref): T | undefined
-}
-
-export namespace ImmutableTree {
-    export type Ref = string
-    export interface MutableNode<T> { ref: ImmutableTree.Ref, value: T, version: number, parent: ImmutableTree.Ref, children: OrderedSet<ImmutableTree.Ref> }
-    export interface Node<T> extends Readonly<MutableNode<T>> { }
-    export interface Nodes<T> extends ImmutableMap<ImmutableTree.Ref, Node<T>> { }
-
-    class Impl<T> implements ImmutableTree<T> {
-        readonly rootRef: ImmutableTree.Ref;
-        readonly version: number;
-        readonly nodes: ImmutableTree.Nodes<T>;
-        readonly getRef: (e: T) => ImmutableTree.Ref;
-
-        getValue(ref: Ref) {
-            const n = this.nodes.get(ref);
-            return n ? n.value : void 0;
-        }
-
-        constructor(rootRef: ImmutableTree.Ref, nodes: ImmutableTree.Nodes<T>, getRef: (e: T) => ImmutableTree.Ref, version: number) {
-            this.rootRef = rootRef;
-            this.nodes = nodes;
-            this.getRef = getRef;
-            this.version = version;
-        }
-    }
-
-    /**
-     * Create an instance of an immutable tree.
-     */
-    export function create<T>(root: T, getRef: (t: T) => ImmutableTree.Ref): ImmutableTree<T> {
-        const ref = getRef(root);
-        const r: Node<T> = { ref, value: root, version: 0, parent: ref, children: OrderedSet() };
-        return new Impl(ref, ImmutableMap([[ref, r]]), getRef, 0);
-    }
-
-    export function asTransient<T>(tree: ImmutableTree<T>) {
-        return new Transient(tree);
-    }
-
-    type N = Node<any>
-    type Ns = Nodes<any>
-
-    type VisitorCtx = { nodes: Ns, state: any, f: (node: N, nodes: Ns, state: any) => boolean | undefined | void };
-
-    function _postOrderFunc(this: VisitorCtx, c: ImmutableTree.Ref | undefined) { _doPostOrder(this, this.nodes.get(c!)!); }
-    function _doPostOrder(ctx: VisitorCtx, root: N) {
-        if (root.children.size) {
-            root.children.forEach(_postOrderFunc, ctx);
-        }
-        ctx.f(root, ctx.nodes, ctx.state);
-    }
-
-    /**
-     * Visit all nodes in a subtree in "post order", meaning leafs get visited first.
-     */
-    export function doPostOrder<T, S>(tree: ImmutableTree<T>, root: Node<T>, state: S, f: (node: Node<T>, nodes: Nodes<T>, state: S) => boolean | undefined | void): S {
-        const ctx: VisitorCtx = { nodes: tree.nodes, state, f };
-        _doPostOrder(ctx, root);
-        return ctx.state;
-    }
-
-    function _preOrderFunc(this: VisitorCtx, c: ImmutableTree.Ref | undefined) { _doPreOrder(this, this.nodes.get(c!)!); }
-    function _doPreOrder(ctx: VisitorCtx, root: N) {
-        const ret = ctx.f(root, ctx.nodes, ctx.state);
-        if (typeof ret === 'boolean' && !ret) return;
-        if (root.children.size) {
-            root.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<T, S>(tree: ImmutableTree<T>, root: Node<T>, state: S, f: (node: Node<T>, nodes: Nodes<T>, state: S) => boolean | undefined | void): S {
-        const ctx: VisitorCtx = { nodes: tree.nodes, state, f };
-        _doPreOrder(ctx, root);
-        return ctx.state;
-    }
-
-    function _subtree(n: N, nodes: Ns, subtree: N[]) { subtree.push(n); }
-    /**
-     * Get all nodes in a subtree, leafs come first.
-     */
-    export function subtreePostOrder<T>(tree: ImmutableTree<T>, root: Node<T>) {
-        return doPostOrder<T, Node<T>[]>(tree, root, [], _subtree);
-    }
-
-
-    function _visitChildToJson(this: Ref[], ref: Ref) { this.push(ref); }
-    interface ToJsonCtx { nodes: [any, any, any[]][], refs: string[], valueToJSON: (v: any) => any }
-    function _visitNodeToJson(this: ToJsonCtx, node: Node<any>) {
-        const children: Ref[] = [];
-        node.children.forEach(_visitChildToJson as any, children);
-        this.nodes.push([this.valueToJSON(node.value), node.parent, children]);
-        this.refs.push(node.ref);
-    }
-
-    export interface Serialized {
-        root: number, // root index
-        nodes: [any /** value */, number /** parent index */, number[] /** children indices */][]
-    }
-
-    export function toJSON<T>(tree: ImmutableTree<T>, valueToJSON: (v: T) => any): Serialized {
-        const ctx: ToJsonCtx = { nodes: [], refs: [], valueToJSON };
-
-        tree.nodes.forEach(_visitNodeToJson as any, ctx);
-
-        const map = new Map<string, number>();
-        let i = 0;
-        for (const n of ctx.refs) map.set(n, i++);
-
-        for (const n of ctx.nodes) {
-            n[1] = map.get(n[1]);
-            const children = n[2];
-            for (i = 0; i < children.length; i++) {
-                children[i] = map.get(children[i]);
-            }
-        }
-        return {
-            root: map.get(tree.rootRef)!,
-            nodes: ctx.nodes
-        };
-    }
-
-    export function fromJSON<T>(data: Serialized, getRef: (v: T) => Ref, valueFromJSON: (v: any) => T): ImmutableTree<T> {
-        const nodes = ImmutableMap<ImmutableTree.Ref, Node<T>>().asMutable();
-
-        const values = data.nodes.map(n => valueFromJSON(n[0]));
-        let i = 0;
-        for (const value of values) {
-            const node = data.nodes[i++];
-            const ref = getRef(value);
-            nodes.set(ref, {
-                ref,
-                value,
-                version: 0,
-                parent: getRef(values[node[1]]),
-                children: OrderedSet(node[2].map(c => getRef(values[c])))
-            });
-        }
-        return new Impl(getRef(values[data.root]), nodes.asImmutable(), getRef, 0);
-    }
-
-    function checkSetRef(oldRef: ImmutableTree.Ref, newRef: ImmutableTree.Ref) {
-        if (oldRef !== newRef) {
-            throw new Error(`Cannot setValue of node '${oldRef}' because the new value has a different ref '${newRef}'.`);
-        }
-    }
-
-    function ensureNotPresent(nodes: Ns, ref: ImmutableTree.Ref) {
-        if (nodes.has(ref)) {
-            throw new Error(`Cannot add node '${ref}' because a different node with this ref already present in the tree.`);
-        }
-    }
-
-    function ensurePresent(nodes: Ns, ref: ImmutableTree.Ref) {
-        if (!nodes.has(ref)) {
-            throw new Error(`Node '${ref}' is not present in the tree.`);
-        }
-    }
-
-    function mutateNode(nodes: Ns, mutations: Map<ImmutableTree.Ref, N>, ref: ImmutableTree.Ref): N {
-        ensurePresent(nodes, ref);
-        if (mutations.has(ref)) {
-            return mutations.get(ref)!;
-        }
-        const node = nodes.get(ref)!;
-        const newNode: N = { ref: node.ref, value: node.value, version: node.version + 1, parent: node.parent, children: node.children.asMutable() };
-        mutations.set(ref, newNode);
-        nodes.set(ref, newNode);
-        return newNode;
-    }
-
-    export class Transient<T> implements ImmutableTree<T> {
-        nodes = this.tree.nodes.asMutable();
-        version: number = this.tree.version + 1;
-        private mutations: Map<ImmutableTree.Ref, Node<T>> = new Map();
-
-        mutate(ref: ImmutableTree.Ref): MutableNode<T> {
-            return mutateNode(this.nodes, this.mutations, ref);
-        }
-
-        get rootRef() { return this.tree.rootRef; }
-        getRef(e: T) {
-            return this.tree.getRef(e);
-        }
-
-        getValue(ref: Ref) {
-            const n = this.nodes.get(ref);
-            return n ? n.value : void 0;
-        }
-
-        add(parentRef: ImmutableTree.Ref, value: T) {
-            const ref = this.getRef(value);
-            ensureNotPresent(this.nodes, ref);
-            const parent = this.mutate(parentRef);
-            const node: Node<T> = { ref, version: 0, value, parent: parent.ref, children: OrderedSet<string>().asMutable() };
-            this.mutations.set(ref, node);
-            parent.children.add(ref);
-            this.nodes.set(ref, node);
-            return node;
-        }
-
-        setValue(ref: ImmutableTree.Ref, value: T): Node<T> {
-            checkSetRef(ref, this.getRef(value));
-            const node = this.mutate(ref);
-            node.value = value;
-            return node;
-        }
-
-        remove(ref: ImmutableTree.Ref): Node<T>[] {
-            if (ref === this.rootRef) {
-                return this.removeChildren(ref);
-            }
-
-            const { nodes, mutations } = this;
-            const node = nodes.get(ref);
-            if (!node) return [];
-            const parent = nodes.get(node.parent)!;
-            this.mutate(parent.ref).children.delete(ref);
-
-            const st = subtreePostOrder(this, node);
-            for (const n of st) {
-                nodes.delete(n.ref);
-                mutations.delete(n.ref);
-            }
-
-            return st;
-        }
-
-        removeChildren(ref: ImmutableTree.Ref): Node<T>[] {
-            const { nodes, mutations } = this;
-            let node = nodes.get(ref);
-            if (!node || !node.children.size) return [];
-            node = this.mutate(ref);
-            const st = subtreePostOrder(this, node);
-            // remove the last node which is the parent
-            st.pop();
-            node.children.clear();
-            for (const n of st) {
-                nodes.delete(n.ref);
-                mutations.delete(n.ref);
-            }
-            return st;
-        }
-
-        asImmutable() {
-            if (this.mutations.size === 0) return this.tree;
-
-            this.mutations.forEach(m => (m as MutableNode<T>).children = m.children.asImmutable());
-            return new Impl<T>(this.tree.rootRef, this.nodes.asImmutable(), this.tree.getRef, this.version);
-        }
-
-        constructor(private tree: ImmutableTree<T>) {
-
-        }
-    }
-}

src/mol-state/selection.ts

@@ -6,7 +6,7 @@
 
 import { StateObject, StateObjectCell } from './object';
 import { State } from './state';
-import { ImmutableTree } from './immutable-tree';
+import { ImmutableTree } from './tree';
 import { Transform } from './transform';
 
 namespace StateSelection {
@@ -71,7 +71,7 @@ namespace StateSelection {
         return Object.create(BuilderPrototype, { compile: { writable: false, configurable: false, value: compile } });
     }
 
-    export function root() { return build(() => (state: State) => [state.cells.get(state.tree.rootRef)!]) }
+    export function root() { return build(() => (state: State) => [state.cells.get(state.tree.root.ref)!]) }
 
 
     export function byRef(...refs: Transform.Ref[]) {
@@ -159,7 +159,7 @@ namespace StateSelection {
     export function children(b: Selector) {
         return flatMap(b, (n, s) => {
             const nodes: StateObjectCell[] = [];
-            s.tree.nodes.get(n.ref)!.children.forEach(c => nodes.push(s.cells.get(c!)!));
+            s.tree.children.get(n.ref).forEach(c => nodes.push(s.cells.get(c!)!));
             return nodes;
         });
     }
@@ -177,8 +177,8 @@ namespace StateSelection {
         let current = tree.nodes.get(root)!, len = types.length;
         while (true) {
             current = tree.nodes.get(current.parent)!;
-            if (current.ref === tree.rootRef) {
-                return cells.get(tree.rootRef);
+            if (current.ref === Transform.RootRef) {
+                return cells.get(Transform.RootRef);
             }
             const obj = cells.get(current.ref)!.obj!;
             for (let i = 0; i < len; i++) {

src/mol-state/state.ts

@@ -5,9 +5,8 @@
  */
 
 import { StateObject, StateObjectCell } from './object';
-import { StateTree } from './tree';
+import { StateTree, ImmutableTree } from './tree';
 import { Transform } from './transform';
-import { ImmutableTree } from './immutable-tree';
 import { Transformer } from './transformer';
 import { StateContext } from './context';
 import { UUID } from 'mol-util';
@@ -17,7 +16,7 @@ export { State }
 
 class State {
     private _tree: StateTree = StateTree.create();
-    private _current: Transform.Ref = this._tree.rootRef;
+    private _current: Transform.Ref = this._tree.root.ref;
     private transformCache = new Map<Transform.Ref, unknown>();
 
     get tree() { return this._tree; }
@@ -67,7 +66,7 @@ class State {
                     stateCtx: this.context,
                     taskCtx,
                     oldTree,
-                    tree: tree,
+                    tree,
                     cells: this.cells,
                     transformCache: this.transformCache
                 };
@@ -81,11 +80,11 @@ class State {
 
     constructor(rootObject: StateObject, params?: { globalContext?: unknown, defaultCellState?: unknown }) {
         const tree = this._tree;
-        const root = tree.getValue(tree.rootRef)!;
+        const root = tree.root;
         const defaultCellState = (params && params.defaultCellState) || { }
 
-        this.cells.set(tree.rootRef, {
-            ref: tree.rootRef,
+        this.cells.set(root.ref, {
+            ref: root.ref,
             obj: rootObject,
             status: 'ok',
             version: root.version,
@@ -94,8 +93,7 @@ class State {
 
         this.context = new StateContext({
             globalContext: params && params.globalContext,
-            defaultCellState,
-            rootRef: tree.rootRef
+            defaultCellState
         });
     }
 }
@@ -142,36 +140,33 @@ namespace State {
         }
     }
 
-    function findUpdateRoots(objects: State.Cells, tree: StateTree) {
-        const findState = {
-            roots: [] as Ref[],
-            objects
-        };
-
-        ImmutableTree.doPreOrder(tree, tree.nodes.get(tree.rootRef)!, findState, (n, _, s) => {
-            if (!s.objects.has(n.ref)) {
-                s.roots.push(n.ref);
-                return false;
-            }
-            const o = s.objects.get(n.ref)!;
-            if (o.version !== n.value.version) {
-                s.roots.push(n.ref);
-                return false;
-            }
+    function findUpdateRoots(cells: State.Cells, tree: StateTree) {
+        const findState = { roots: [] as Ref[], cells };
+        ImmutableTree.doPreOrder(tree, tree.root, findState, _findUpdateRoots);
+        return findState.roots;
+    }
 
-            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;
+        }
 
-        return findState.roots;
+        return true;
     }
 
     type FindDeletesCtx = { newTree: StateTree, cells: State.Cells, deletes: Ref[] }
-    function _visitCheckDelete(n: ImmutableTree.Node<any>, _: any, ctx: FindDeletesCtx) {
+    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.nodes.get(ctx.oldTree.rootRef), deleteCtx, _visitCheckDelete);
+        ImmutableTree.doPostOrder(ctx.oldTree, ctx.oldTree.root, deleteCtx, _visitCheckDelete);
         return deleteCtx.deletes;
     }
 
@@ -190,7 +185,7 @@ namespace State {
         if (changed) ctx.stateCtx.events.object.stateChanged.next({ ref });
     }
 
-    function _initVisitor(t: ImmutableTree.Node<Transform>, _: any, ctx: UpdateContext) {
+    function _initVisitor(t: Transform, _: any, ctx: UpdateContext) {
         setObjectState(ctx, t.ref, 'pending');
     }
     /** Return "resolve set" */
@@ -209,7 +204,7 @@ namespace State {
             wrap.obj = void 0;
         }
 
-        const children = ctx.tree.nodes.get(ref)!.children.values();
+        const children = ctx.tree.children.get(ref).values();
         while (true) {
             const next = children.next();
             if (next.done) return;
@@ -217,15 +212,15 @@ namespace State {
         }
     }
 
-    function findAncestor(tree: StateTree, objects: State.Cells, root: Ref, types: { type: StateObject.Type }[]): StateObject {
+    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 === tree.rootRef) {
-                return objects.get(tree.rootRef)!.obj!;
+            if (current.ref === Transform.RootRef) {
+                return cells.get(Transform.RootRef)!.obj!;
             }
-            const obj = objects.get(current.ref)!.obj!;
-            for (const t of types) if (obj.type === t.type) return objects.get(current.ref)!.obj!;
+            const obj = cells.get(current.ref)!.obj!;
+            for (const t of types) if (obj.type === t.type) return cells.get(current.ref)!.obj!;
         }
     }
 
@@ -247,7 +242,7 @@ namespace State {
             return;
         }
 
-        const children = ctx.tree.nodes.get(root)!.children.values();
+        const children = ctx.tree.children.get(root).values();
         while (true) {
             const next = children.next();
             if (next.done) return;
@@ -257,7 +252,7 @@ namespace State {
 
     async function updateNode(ctx: UpdateContext, currentRef: Ref) {
         const { oldTree, tree, cells } = ctx;
-        const transform = tree.getValue(currentRef)!;
+        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)) {
@@ -274,9 +269,9 @@ namespace State {
         } else {
             // console.log('updating...', transform.transformer.id);
             const current = cells.get(currentRef)!;
-            const oldParams = oldTree.getValue(currentRef)!.params;
+            const oldParams = oldTree.nodes.get(currentRef)!.params;
 
-            const updateKind = current.status === 'ok' || current.ref === ctx.tree.rootRef
+            const updateKind = current.status === 'ok' || current.ref === Transform.RootRef
                 ? await updateObject(ctx, currentRef, transform.transformer, parent, current.obj!, oldParams, transform.params)
                 : Transformer.UpdateResult.Recreate;
 

src/mol-state/tree.ts

@@ -5,36 +5,34 @@
  */
 
 import { Transform } from './transform';
-import { ImmutableTree } from './immutable-tree';
+import { ImmutableTree, TransientTree } from './tree/base';
 import { Transformer } from './transformer';
 import { StateObject } from './object';
 
-interface StateTree extends ImmutableTree<Transform> { }
+export { StateTree, ImmutableTree, TransientTree }
+
+interface StateTree extends ImmutableTree { }
 
 namespace StateTree {
-    export interface Transient extends ImmutableTree.Transient<Transform> { }
+    export interface Transient extends TransientTree { }
     export interface Serialized extends ImmutableTree.Serialized { }
 
-    function _getRef(t: Transform) { return t.ref; }
-
     export function create() {
-        return ImmutableTree.create<Transform>(Transform.createRoot('<:root:>'), _getRef);
+        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)!.value;
+        const t = tree.nodes.get(ref)!;
         const newTransform = Transform.updateParams(t, params);
-        const newTree = ImmutableTree.asTransient(tree);
-        newTree.setValue(ref, newTransform);
-        return newTree.asImmutable();
+        return tree.asTransient().set(newTransform).asImmutable();
     }
 
     export function toJSON(tree: StateTree) {
-        return ImmutableTree.toJSON(tree, Transform.toJSON) as Serialized;
+        return ImmutableTree.toJSON(tree) as Serialized;
     }
 
     export function fromJSON(data: Serialized): StateTree {
-        return ImmutableTree.fromJSON(data, _getRef, Transform.fromJSON);
+        return ImmutableTree.fromJSON(data);
     }
 
     export interface Builder {
@@ -53,19 +51,19 @@ namespace StateTree {
         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.rootRef as any, 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: ImmutableTree.asTransient(tree) } }
+            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(this.ref, t);
+                this.state.tree.add(t);
                 return new To(this.state, t.ref, this.root);
             }
 
@@ -80,6 +78,4 @@ namespace StateTree {
             }
         }
     }
-}
-
-export { StateTree }
+}

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

@@ -5,7 +5,7 @@
  */
 
 import { Map as ImmutableMap, OrderedSet } from 'immutable';
-import { Transform } from './transform';
+import { Transform } from '../transform';
 
 export { ImmutableTree, TransientTree }
 
@@ -28,7 +28,8 @@ namespace ImmutableTree {
         readonly size: number,
         readonly values: OrderedSet<Ref>['values'],
         has(ref: Ref): boolean,
-        readonly forEach: OrderedSet<Ref>['forEach']
+        readonly forEach: OrderedSet<Ref>['forEach'],
+        readonly map: OrderedSet<Ref>['map']
     }
 
     export type Node = Transform
@@ -49,11 +50,11 @@ namespace ImmutableTree {
     /**
      * Create an instance of an immutable tree.
      */
-    export function create<T>(root: Transform): ImmutableTree {
+    export function create(root: Transform): ImmutableTree {
         return new Impl(ImmutableMap([[root.ref, root]]), ImmutableMap([[root.ref, OrderedSet()]]), 0);
     }
 
-    export function construct<T>(nodes: Nodes, children: Children, version: number): ImmutableTree {
+    export function construct(nodes: Nodes, children: Children, version: number): ImmutableTree {
         return new Impl(nodes, children, version);
     }
 
@@ -160,6 +161,7 @@ class TransientTree implements ImmutableTree {
 
     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)! }
@@ -213,12 +215,14 @@ class TransientTree implements ImmutableTree {
         }
 
         this.nodes.set(ref, transform);
-        return 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[] {
@@ -243,24 +247,8 @@ class TransientTree implements ImmutableTree {
         return st;
     }
 
-    // removeChildren(ref: ImmutableTree.Ref): TransientTree.Node[] {
-    //     const { nodes, mutations } = this;
-    //     let node = nodes.get(ref);
-    //     if (!node || !node.children.size) return [];
-    //     node = this.mutate(ref);
-    //     const st = ImmutableTree.subtreePostOrder(this, node);
-    //     // remove the last node which is the parent
-    //     st.pop();
-    //     node.children.clear();
-    //     for (const n of st) {
-    //         nodes.delete(n.value.ref);
-    //         mutations.delete(n.value.ref);
-    //     }
-    //     return st;
-    // }
-
     asImmutable() {
-        if (this.mutations.size === 0) return this.tree;
+        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);

src/perf-tests/state.ts

@@ -121,14 +121,14 @@ export function printTTree(tree: StateTree) {
     let lines: string[] = [];
     function print(offset: string, ref: any) {
         const t = tree.nodes.get(ref)!;
-        const tr = t.value;
+        const tr = t;
 
         const name = tr.transformer.id;
-        lines.push(`${offset}|_ (${ref}) ${name} ${tr.params ? JSON.stringify(tr.params) : ''}, v${t.value.version}`);
+        lines.push(`${offset}|_ (${ref}) ${name} ${tr.params ? JSON.stringify(tr.params) : ''}, v${t.version}`);
         offset += '   ';
 
-        t.children.forEach(c => print(offset, c!));
+        tree.children.get(ref).forEach(c => print(offset, c!));
     }
-    print('', tree.rootRef);
+    print('', tree.root.ref);
     console.log(lines.join('\n'));
 }