mol-state: wip immutable tree

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

@@ -15,36 +15,34 @@ export { ImmutableTree, TransientTree }
  */
 interface ImmutableTree {
     readonly version: number,
+    readonly root: Transform,
     readonly nodes: ImmutableTree.Nodes,
-    readonly root: ImmutableTree.Node,
-    get(ref: ImmutableTree.Ref): Transform | undefined,
-    //getChildren(ref: ImmutableTree.Ref): OrderedSet<ImmutableTree.Node>
+    readonly children: ImmutableTree.Children,
     asTransient(): TransientTree
 }
 
 namespace ImmutableTree {
-    export type Ref = Transform.Ref
-    export interface Node extends Readonly<TransientTree.Node> { }
-    export interface Nodes extends ImmutableMap<ImmutableTree.Ref, Node> { }
+    type Ref = Transform.Ref
 
-    class Impl implements ImmutableTree {
-        readonly version: number;
-        readonly nodes: ImmutableTree.Nodes;
+    export interface ChildSet {
+        readonly size: number,
+        readonly values: OrderedSet<Ref>['values'],
+        has(ref: Ref): boolean,
+        readonly forEach: OrderedSet<Ref>['forEach']
+    }
 
-        get root() { return this.nodes.get(Transform.RootRef)! }
+    export type Node = Transform
+    export type Nodes = ImmutableMap<Ref, Transform>
+    export type Children = ImmutableMap<Ref, ChildSet>
 
-        get(ref: Ref) {
-            const n = this.nodes.get(ref);
-            return n ? n.value : void 0;
-        }
+    class Impl implements ImmutableTree {
+        get root() { return this.nodes.get(Transform.RootRef)! }
 
         asTransient(): TransientTree {
             return new TransientTree(this);
         }
 
-        constructor(nodes: ImmutableTree.Nodes, version: number) {
-            this.nodes = nodes;
-            this.version = version;
+        constructor(public nodes: ImmutableTree.Nodes, public children: Children, public version: number) {
         }
     }
 
@@ -52,39 +50,40 @@ namespace ImmutableTree {
      * Create an instance of an immutable tree.
      */
     export function create<T>(root: Transform): ImmutableTree {
-        const r: Node = { value: root, version: 0, parent: root.ref, children: OrderedSet() };
-        return new Impl(ImmutableMap([[root.ref, r]]), 0);
+        return new Impl(ImmutableMap([[root.ref, root]]), ImmutableMap([[root.ref, OrderedSet()]]), 0);
     }
 
-    export function fromNodes<T>(nodes: Nodes, version: number): ImmutableTree {
-        return new Impl(nodes, version);
+    export function construct<T>(nodes: Nodes, children: Children, version: number): ImmutableTree {
+        return new Impl(nodes, children, version);
     }
 
-    type VisitorCtx = { nodes: Nodes, state: any, f: (node: Node, nodes: Nodes, state: any) => boolean | undefined | void };
+    type VisitorCtx = { tree: ImmutableTree, state: any, f: (node: Node, tree: ImmutableTree, state: any) => boolean | undefined | void };
 
-    function _postOrderFunc(this: VisitorCtx, c: ImmutableTree.Ref | undefined) { _doPostOrder(this, this.nodes.get(c!)!); }
+    function _postOrderFunc(this: VisitorCtx, c: Ref | undefined) { _doPostOrder(this, this.tree.nodes.get(c!)); }
     function _doPostOrder(ctx: VisitorCtx, root: Node) {
-        if (root.children.size) {
-            root.children.forEach(_postOrderFunc, ctx);
+        const children = ctx.tree.children.get(root.ref);
+        if (checkSetRef && children.size) {
+            children.forEach(_postOrderFunc, ctx);
         }
-        ctx.f(root, ctx.nodes, ctx.state);
+        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, nodes: Nodes, state: S) => boolean | undefined | void): S {
-        const ctx: VisitorCtx = { nodes: tree.nodes, state, f };
+    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: ImmutableTree.Ref | undefined) { _doPreOrder(this, this.nodes.get(c!)!); }
+    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.nodes, ctx.state);
+        const ret = ctx.f(root, ctx.tree, ctx.state);
         if (typeof ret === 'boolean' && !ret) return;
-        if (root.children.size) {
-            root.children.forEach(_preOrderFunc, ctx);
+        const children = ctx.tree.children.get(root.ref);
+        if (checkSetRef && children.size) {
+            children.forEach(_preOrderFunc, ctx);
         }
     }
 
@@ -92,13 +91,13 @@ 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, nodes: Nodes, state: S) => boolean | undefined | void): S {
-        const ctx: VisitorCtx = { nodes: tree.nodes, state, f };
+    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, nodes: Nodes, subtree: Node[]) { subtree.push(n); }
+    function _subtree(n: Node, _: any, subtree: Node[]) { subtree.push(n); }
     /**
      * Get all nodes in a subtree, leafs come first.
      */
@@ -108,19 +107,19 @@ namespace ImmutableTree {
 
 
     function _visitChildToJson(this: Ref[], ref: Ref) { this.push(ref); }
-    interface ToJsonCtx { nodes: [Transform.Serialized, any, any[]][] }
+    interface ToJsonCtx { tree: ImmutableTree, nodes: [Transform.Serialized, any[]][] }
     function _visitNodeToJson(this: ToJsonCtx, node: Node) {
         const children: Ref[] = [];
-        node.children.forEach(_visitChildToJson as any, children);
-        this.nodes.push([Transform.toJSON(node.value), node.parent, children]);
+        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 /** parent index */, number[] /** children indices */][]
+        nodes: [any /** value */, number[] /** children indices */][]
     }
 
     export function toJSON<T>(tree: ImmutableTree): Serialized {
-        const ctx: ToJsonCtx = { nodes: [] };
+        const ctx: ToJsonCtx = { tree, nodes: [] };
 
         tree.nodes.forEach(_visitNodeToJson as any, ctx);
 
@@ -129,8 +128,7 @@ namespace ImmutableTree {
         for (const n of ctx.nodes) map.set(n[0].ref, i++);
 
         for (const n of ctx.nodes) {
-            n[1] = map.get(n[1]);
-            const children = n[2];
+            const children = n[1];
             for (i = 0; i < children.length; i++) {
                 children[i] = map.get(children[i]);
             }
@@ -141,102 +139,131 @@ namespace ImmutableTree {
     }
 
     export function fromJSON<T>(data: Serialized): ImmutableTree {
-        const nodes = ImmutableMap<ImmutableTree.Ref, Node>().asMutable();
+        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,
-                version: 0,
-                parent: values[node[1]].ref,
-                children: OrderedSet(node[2].map(c => values[c].ref))
-            });
+            nodes.set(ref, value);
+            children.set(ref, OrderedSet(node[1].map(c => values[c].ref)));
         }
-        return new Impl(nodes.asImmutable(), 0);
+        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 mutations: Map<Transform.Ref, TransientTree.Node> = new Map();
 
-    get root() { return this.nodes.get(Transform.RootRef)! }
+    private mutations: Map<Transform.Ref, OrderedSet<Transform.Ref>> = new Map();
 
-    get(ref: Transform.Ref) {
-        const n = this.nodes.get(ref);
-        return n ? n.value : void 0;
-    }
+    get root() { return this.nodes.get(Transform.RootRef)! }
 
     asTransient() {
         return this.asImmutable().asTransient();
     }
 
-    mutate(ref: ImmutableTree.Ref): TransientTree.Node {
-        return mutateNode(this.nodes, this.mutations, ref);
+    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);
+        }
     }
 
-    add(parentRef: ImmutableTree.Ref, value: Transform) {
-        const ref = value.ref;
-        ensureNotPresent(this.nodes, ref);
-        const parent = this.mutate(parentRef);
-        const node: TransientTree.Node = { version: 0, value, parent: parentRef, children: OrderedSet<string>().asMutable() };
-        this.mutations.set(ref, node);
-        parent.children.add(ref);
-        this.nodes.set(ref, node);
-        return node;
+    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);
+        }
     }
 
-    setValue(ref: ImmutableTree.Ref, value: Transform): TransientTree.Node {
-        checkSetRef(ref, value.ref);
-        const node = this.mutate(ref);
-        node.value = value;
-        return node;
+    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);
     }
 
-    remove(ref: ImmutableTree.Ref): TransientTree.Node[] {
-        if (ref === Transform.RootRef) {
-            return this.removeChildren(ref);
+    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());
         }
 
-        const { nodes, mutations } = this;
+        this.nodes.set(ref, transform);
+        return transform;
+    }
+
+    set(transform: Transform) {
+        ensurePresent(this.nodes, transform.ref);
+        this.nodes.set(transform.ref, transform);
+    }
+
+    remove(ref: Transform.Ref): Transform[] {
+        const { nodes, mutations, children } = this;
         const node = nodes.get(ref);
         if (!node) return [];
-        this.mutate(node.parent).children.delete(ref);
 
         const st = ImmutableTree.subtreePostOrder(this, node);
-        for (const n of st) {
-            nodes.delete(n.value.ref);
-            mutations.delete(n.value.ref);
+        if (ref === Transform.RootRef) {
+            st.pop();
+            this.clearRoot();
+        } else {
+            this.removeChild(node.parent, node.ref);
         }
 
-        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);
+            nodes.delete(n.ref);
+            children.delete(n.ref);
+            mutations.delete(n.ref);
         }
+
         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;
 
-        this.mutations.forEach(m => (m as TransientTree.Node).children = m.children.asImmutable());
-        return ImmutableTree.fromNodes(this.nodes.asMutable(), this.version);
+        this.mutations.forEach(fixChildMutations, this.children);
+        return ImmutableTree.construct(this.nodes.asImmutable(), this.children.asImmutable(), this.version);
     }
 
     constructor(private tree: ImmutableTree) {
@@ -244,37 +271,20 @@ class TransientTree implements ImmutableTree {
     }
 }
 
-namespace TransientTree {
-    export interface Node { value: Transform, version: number, parent: ImmutableTree.Ref, children: OrderedSet<ImmutableTree.Ref> }
-}
-
+function fixChildMutations(this: ImmutableTree.Children, m: OrderedSet<Transform.Ref>, k: Transform.Ref) { this.set(k, m.asImmutable()); }
 
-function checkSetRef(oldRef: ImmutableTree.Ref, newRef: ImmutableTree.Ref) {
+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 ensureNotPresent(nodes: ImmutableTree.Nodes, 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 parentNotPresent(ref: Transform.Ref) {
+    throw new Error(`Parent '${ref}' must be present in the tree.`);
 }
 
-function ensurePresent(nodes: ImmutableTree.Nodes, ref: ImmutableTree.Ref) {
+function ensurePresent(nodes: ImmutableTree.Nodes, ref: Transform.Ref) {
     if (!nodes.has(ref)) {
         throw new Error(`Node '${ref}' is not present in the tree.`);
     }
-}
-
-function mutateNode(nodes: ImmutableTree.Nodes, mutations: Map<ImmutableTree.Ref, TransientTree.Node>, ref: ImmutableTree.Ref): TransientTree.Node {
-    ensurePresent(nodes, ref);
-    if (mutations.has(ref)) {
-        return mutations.get(ref)!;
-    }
-    const node = nodes.get(ref)!;
-    const newNode: TransientTree.Node = { value: node.value, version: node.version + 1, parent: node.parent, children: node.children.asMutable() };
-    mutations.set(ref, newNode);
-    nodes.set(ref, newNode);
-    return newNode;
 }

src/mol-state/state.ts

@@ -268,7 +268,7 @@ namespace State {
                 obj,
                 status: 'ok',
                 version: transform.version,
-                state: { ...ctx.stateCtx.defaultCellState, ...transform.defaultProps }
+                state: { ...ctx.stateCtx.defaultCellState, ...transform.cellState }
             });
             return { action: 'created', obj };
         } else {
@@ -288,13 +288,13 @@ namespace State {
                         obj,
                         status: 'ok',
                         version: transform.version,
-                        state: { ...ctx.stateCtx.defaultCellState, ...current.state, ...transform.defaultProps }
+                        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.defaultProps };
+                    current.state = { ...ctx.stateCtx.defaultCellState, ...current.state, ...transform.cellState };
                     return { action: 'updated', obj: current.obj };
                 default:
                     // TODO check if props need to be updated

src/mol-state/transform.ts

@@ -4,16 +4,17 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import { StateObject } from './object';
+import { StateObject, StateObjectCell } from './object';
 import { Transformer } from './transformer';
 import { UUID } from 'mol-util';
 
 export interface Transform<A extends StateObject = StateObject, B extends StateObject = StateObject, P = unknown> {
+    readonly parent: Transform.Ref,
     readonly transformer: Transformer<A, B, P>,
     readonly params: P,
     readonly ref: Transform.Ref,
     readonly version: string,
-    readonly defaultProps?: unknown
+    readonly cellState?: Partial<StateObjectCell.State>
 }
 
 export namespace Transform {
@@ -21,16 +22,17 @@ export namespace Transform {
 
     export const RootRef = '-=root=-' as Ref;
 
-    export interface Options { ref?: Ref, defaultProps?: unknown }
+    export interface Options { ref?: Ref, cellState?: Partial<StateObjectCell.State> }
 
-    export function create<A extends StateObject, B extends StateObject, P>(transformer: Transformer<A, B, P>, params?: P, options?: Options): Transform<A, B, P> {
+    export function create<A extends StateObject, B extends StateObject, P>(parent: Ref, transformer: Transformer<A, B, P>, params?: P, options?: Options): Transform<A, B, P> {
         const ref = options && options.ref ? options.ref : UUID.create() as string as Ref;
         return {
+            parent,
             transformer,
             params: params || {} as any,
             ref,
             version: UUID.create(),
-            defaultProps: options && options.defaultProps
+            cellState: options && options.cellState
         }
     }
 
@@ -39,15 +41,16 @@ export namespace Transform {
     }
 
     export function createRoot(ref: Ref): Transform {
-        return create(Transformer.ROOT, {}, { ref });
+        return create(RootRef, Transformer.ROOT, {}, { ref });
     }
 
     export interface Serialized {
+        parent: string,
         transformer: string,
         params: any,
         ref: string,
         version: string,
-        defaultProps?: unknown
+        cellState?: Partial<StateObjectCell.State>
     }
 
     function _id(x: any) { return x; }
@@ -56,11 +59,12 @@ export namespace Transform {
             ? t.transformer.definition.customSerialization.toJSON
             : _id;
         return {
+            parent: t.parent,
             transformer: t.transformer.id,
             params: pToJson(t.params),
             ref: t.ref,
             version: t.version,
-            defaultProps: t.defaultProps
+            cellState: t.cellState
         };
     }
 
@@ -70,11 +74,12 @@ export namespace Transform {
             ? transformer.definition.customSerialization.toJSON
             : _id;
         return {
+            parent: t.parent as Ref,
             transformer,
             params: pFromJson(t.params),
             ref: t.ref as Ref,
             version: t.version,
-            defaultProps: t.defaultProps
+            cellState: t.cellState
         };
     }
 }

src/mol-state/transformer.ts

@@ -10,7 +10,7 @@ import { Transform } from './transform';
 import { ParamDefinition as PD } from 'mol-util/param-definition';
 
 export interface Transformer<A extends StateObject = StateObject, B extends StateObject = StateObject, P = unknown> {
-    apply(params?: P, props?: Partial<Transform.Options>): Transform<A, B, P>,
+    apply(parent: Transform.Ref, params?: P, props?: Partial<Transform.Options>): Transform<A, B, P>,
     readonly namespace: string,
     readonly id: Transformer.Id,
     readonly definition: Transformer.Definition<A, B, P>
@@ -114,7 +114,7 @@ export namespace Transformer {
         }
 
         const t: Transformer<A, B, P> = {
-            apply(params, props) { return Transform.create<A, B, P>(t as any, params, props); },
+            apply(parent, params, props) { return Transform.create<A, B, P>(parent, t as any, params, props); },
             namespace,
             id,
             definition

src/mol-state/tree.ts

@@ -64,7 +64,7 @@ namespace StateTree {
 
         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(params, props);
+                const t = tr.apply(this.ref, params, props);
                 this.state.tree.add(this.ref, t);
                 return new To(this.state, t.ref, this.root);
             }