mol-state: StateTransformer.Definition.isDecorator

- moved from StateTransform.isDecorator
- when using StateBuilder.apply:
  * decorators are always "inserted"
  * if a node has a decorator, the transformer is applied to the decorator instead (recursive)

src/mol-plugin-state/builder/structure.ts

@@ -113,7 +113,7 @@ export class StructureBuilder {
     async insertModelProperties(model: StateObjectRef<SO.Molecule.Model>, params?: StateTransformer.Params<StateTransforms['Model']['CustomModelProperties']>) {
         const state = this.dataState;
         const props = state.build().to(model)
-            .insert(StateTransforms.Model.CustomModelProperties, params, { tags: StructureBuilderTags.ModelProperties, isDecorator: true });
+            .apply(StateTransforms.Model.CustomModelProperties, params, { tags: StructureBuilderTags.ModelProperties });
         await this.plugin.updateDataState(props, { revertOnError: true });
         return props.selector;
     }
@@ -152,7 +152,7 @@ export class StructureBuilder {
     async insertStructureProperties(structure: StateObjectRef<SO.Molecule.Structure>, params?: StateTransformer.Params<StateTransforms['Model']['CustomStructureProperties']>) {
         const state = this.dataState;
         const props = state.build().to(structure)
-            .insert(StateTransforms.Model.CustomStructureProperties, params, { tags: StructureBuilderTags.StructureProperties, isDecorator: true });
+            .apply(StateTransforms.Model.CustomStructureProperties, params, { tags: StructureBuilderTags.StructureProperties });
         await this.plugin.updateDataState(props, { revertOnError: true });
         return props.selector;
     }

src/mol-plugin-state/builder/structure/hierarchy-preset.ts

@@ -52,12 +52,10 @@ const firstModel = TrajectoryHierarchyPresetProvider({
         const representation =  await plugin.builders.structure.representation.applyPreset(structureProperties, params.representationPreset || 'auto');
 
         return {
-            model: modelProperties,
-            modelRoot: model,
+            model,
             modelProperties,
             unitcell,
-            structure: structureProperties,
-            structureRoot: structure,
+            structure,
             structureProperties,
             representation
         };

src/mol-plugin-state/manager/structure/component.ts

@@ -109,7 +109,7 @@ class StructureComponentManager extends PluginComponent<StructureComponentManage
         return this.plugin.dataTransaction(async () => {
             await this.clearComponents(structures);
             for (const s of structures) {
-                await this.plugin.builders.structure.representation.applyPreset(s.childRoot, provider, params);
+                await this.plugin.builders.structure.representation.applyPreset(s.cell, provider, params);
             }
         }, { canUndo: 'Preset' });
     }
@@ -273,7 +273,7 @@ class StructureComponentManager extends PluginComponent<StructureComponentManage
 
             const componentKey = UUID.create22();
             for (const s of xs) {
-                const component = await this.plugin.builders.structure.tryCreateComponentFromSelection(s.childRoot, params.selection, componentKey, {
+                const component = await this.plugin.builders.structure.tryCreateComponentFromSelection(s.cell, params.selection, componentKey, {
                     label: params.label || (params.selection === StructureSelectionQueries.current ? 'Custom Selection' : ''),
                 });
                 if (params.representation === 'none' || !component) continue;

src/mol-plugin-state/manager/structure/hierarchy-state.ts

@@ -56,13 +56,11 @@ export interface ModelRef extends RefBase<'model', SO.Molecule.Model> {
     properties?: ModelPropertiesRef,
     structures: StructureRef[],
     genericRepresentations?: GenericRepresentationRef[],
-    unitcell?: ModelUnitcellRef,
-    /** to support decorators */
-    childRoot: StateObjectCell<SO.Molecule.Model>
+    unitcell?: ModelUnitcellRef
 }
 
 function ModelRef(cell: StateObjectCell<SO.Molecule.Model>, trajectory?: TrajectoryRef): ModelRef {
-    return { kind: 'model', cell, version: cell.transform.version, trajectory, structures: [], childRoot: cell };
+    return { kind: 'model', cell, version: cell.transform.version, trajectory, structures: [] };
 }
 
 export interface ModelPropertiesRef extends RefBase<'model-properties', SO.Molecule.Model, StateTransforms['Model']['CustomModelProperties']> {
@@ -90,13 +88,11 @@ export interface StructureRef extends RefBase<'structure', SO.Molecule.Structure
         surroundings?: StructureComponentRef,
     },
     genericRepresentations?: GenericRepresentationRef[],
-    volumeStreaming?: StructureVolumeStreamingRef,
-    /** to support decorators */
-    childRoot: StateObjectCell<SO.Molecule.Structure>
+    volumeStreaming?: StructureVolumeStreamingRef
 }
 
 function StructureRef(cell: StateObjectCell<SO.Molecule.Structure>, model?: ModelRef): StructureRef {
-    return { kind: 'structure', cell, version: cell.transform.version, model, components: [], childRoot: cell };
+    return { kind: 'structure', cell, version: cell.transform.version, model, components: [] };
 }
 
 export interface StructurePropertiesRef extends RefBase<'structure-properties', SO.Molecule.Structure, StateTransforms['Model']['CustomStructureProperties']> {
@@ -285,14 +281,6 @@ function _doPreOrder(ctx: VisitorCtx, root: StateTransform) {
         }
     }
 
-    if (cell.transform.isDecorator) {
-        if (state.currentModel && SO.Molecule.Model.is(cell.obj)) {
-            state.currentModel.childRoot = cell;
-        } else if (state.currentStructure && !state.currentComponent && SO.Molecule.Structure.is(cell.obj)) {
-            state.currentStructure.childRoot = cell;
-        }
-    }
-
     if (!onLeave && !cell.state.isGhost && state.currentComponent && SO.Molecule.Structure.Representation3D.is(cell.obj)) {
         createOrUpdateRefList(state, cell, state.currentComponent.representations, StructureRepresentationRef, cell, state.currentComponent);
     } else if (!cell.state.isGhost && SO.isRepresentation3D(cell.obj)) {

src/mol-plugin-state/transforms/model.ts

@@ -685,6 +685,7 @@ type CustomModelProperties = typeof CustomModelProperties
 const CustomModelProperties = PluginStateTransform.BuiltIn({
     name: 'custom-model-properties',
     display: { name: 'Custom Model Properties' },
+    isDecorator: true,
     from: SO.Molecule.Model,
     to: SO.Molecule.Model,
     params: (a, ctx: PluginContext) => {
@@ -734,6 +735,7 @@ type CustomStructureProperties = typeof CustomStructureProperties
 const CustomStructureProperties = PluginStateTransform.BuiltIn({
     name: 'custom-structure-properties',
     display: { name: 'Custom Structure Properties' },
+    isDecorator: true,
     from: SO.Molecule.Structure,
     to: SO.Molecule.Structure,
     params: (a, ctx: PluginContext) => {

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

@@ -93,7 +93,7 @@ class StateTreeNode extends PluginUIComponent<{ cell: StateObjectCell, depth: nu
     hasDecorator(children: _StateTree.ChildSet) {
         if (children.size !== 1) return false;
         const ref = children.values().next().value;
-        return !!this.props.cell.parent.tree.transforms.get(ref).isDecorator;
+        return !!this.props.cell.parent.tree.transforms.get(ref).transformer.definition.isDecorator;
     }
 
     render() {

src/mol-plugin-ui/structure/volume.tsx

@@ -46,7 +46,7 @@ export class VolumeStreamingControls extends CollapsableControls<{}, VolumeStrea
 
     renderEnable() {
         const pivot = this.pivot;
-        return <ApplyActionControl state={pivot.cell.parent} action={InitVolumeStreaming} initiallyCollapsed={true} nodeRef={pivot.childRoot.transform.ref} simpleApply={{ header: 'Enable', icon: 'check' }} />;
+        return <ApplyActionControl state={pivot.cell.parent} action={InitVolumeStreaming} initiallyCollapsed={true} nodeRef={pivot.cell.transform.ref} simpleApply={{ header: 'Enable', icon: 'check' }} />;
     }
 
     renderParams() {

src/mol-plugin/behavior/dynamic/volume-streaming/behavior.ts

@@ -276,9 +276,7 @@ export namespace VolumeStreaming {
             const parent = this.plugin.helpers.substructureParent.get(loci.structure);
             if (!parent) return Box3D.empty();
             const root = this.getStructureRoot();
-            if (!root || !root.obj) return Box3D.empty();
-            const rootParent = this.plugin.helpers.substructureParent.get(root.obj.data);
-            if (parent !== rootParent) return Box3D.empty();
+            if (!root || root.obj?.data !== parent.obj?.data) return Box3D.empty();
 
             const extendedLoci = StructureElement.Loci.extendToWholeResidues(loci)
             const box = StructureElement.Loci.getBoundary(extendedLoci).box

src/mol-plugin/util/substructure-parent-helper.ts

@@ -5,15 +5,14 @@
  */
 
 import { Structure } from '../../mol-model/structure';
-import { State, StateObject, StateSelection, StateObjectCell, StateTransform } from '../../mol-state';
-import { PluginContext } from '../context';
 import { PluginStateObject } from '../../mol-plugin-state/objects';
-import { arraySetAdd, arraySetRemove } from '../../mol-util/array';
+import { State, StateObject, StateObjectCell, StateSelection } from '../../mol-state';
+import { PluginContext } from '../context';
 
 export { SubstructureParentHelper };
 
 class SubstructureParentHelper {
-    private decorators = new Map<string, string[]>();
+    // private decorators = new Map<string, string[]>();
     private root = new Map<Structure, { ref: string, count: number }>();
     private tracked = new Map<string, Structure>();
 
@@ -21,46 +20,7 @@ class SubstructureParentHelper {
     get(s: Structure, ignoreDecorators = false): StateObjectCell<PluginStateObject.Molecule.Structure> | undefined {
         const r = this.root.get(s);
         if (!r) return;
-        const decorators = this.decorators.get(r.ref);
-        if (ignoreDecorators || !decorators) return this.plugin.state.data.cells.get(r.ref);
-        return this.plugin.state.data.cells.get(this.findDeepestDecorator(r.ref, decorators));
-    }
-
-    private findDeepestDecorator(ref: string, decorators: string[]) {
-        if (decorators.length === 0) return ref;
-        if (decorators.length === 1) return decorators[0];
-
-        const cells = this.plugin.state.data.cells;
-        let depth = 0, ret = ref;
-        for (const dr of decorators) {
-            let c = cells.get(dr);
-            let d = 0;
-            while (c && c.transform.ref !== StateTransform.RootRef) {
-                d++;
-                c = cells.get(c.transform.parent);
-            }
-            if (d > depth) {
-                ret = dr;
-                depth = d;
-            }
-        }
-        return ret;
-    }
-
-    private addDecorator(root: string, ref: string) {
-        if (this.decorators.has(root)) {
-            arraySetAdd(this.decorators.get(root)!, ref);
-        } else {
-            this.decorators.set(root, [ref]);
-        }
-    }
-
-    private tryRemoveDecorator(root: string, ref: string) {
-        if (this.decorators.has(root)) {
-            const xs = this.decorators.get(root)!;
-            arraySetRemove(xs, ref);
-            if (xs.length === 0) this.decorators.delete(root);
-        }
+        return this.plugin.state.data.cells.get(r.ref);
     }
 
     private addMapping(state: State, ref: string, obj: StateObject) {
@@ -68,11 +28,9 @@ class SubstructureParentHelper {
 
         this.tracked.set(ref, obj.data);
 
-        let parentRef;
         // if the structure is already present in the tree, do not rewrite the root.
         if (this.root.has(obj.data)) {
             const e = this.root.get(obj.data)!;
-            parentRef = e.ref;
             e.count++;
         } else {
             const parent = state.select(StateSelection.Generators.byRef(ref).rootOfType([PluginStateObject.Molecule.Structure]))[0];
@@ -80,17 +38,9 @@ class SubstructureParentHelper {
             if (!parent) {
                 this.root.set(obj.data, { ref, count: 1 });
             } else {
-                parentRef = parent.transform.ref;
-                this.root.set(obj.data, { ref: parentRef, count: 1 });
+                this.root.set(obj.data, { ref: parent.transform.ref, count: 1 });
             }
         }
-
-        if (!parentRef) return;
-
-        const cell = state.cells.get(ref);
-        if (cell?.transform.isDecorator) {
-            this.addDecorator(parentRef, ref);
-        }
     }
 
     private removeMapping(ref: string) {
@@ -101,8 +51,6 @@ class SubstructureParentHelper {
 
         const root = this.root.get(s)!;
 
-        this.tryRemoveDecorator(root.ref, ref);
-
         if (root.count > 1) {
             root.count--;
         } else {

src/mol-state/state/builder.ts

@@ -111,11 +111,26 @@ namespace StateBuilder {
 
         readonly ref: StateTransform.Ref;
 
+        private getApplyRoot(ref?: StateTransform.Ref): StateTransform.Ref {
+            const children = this.state.tree.children.get(ref || this.ref);
+            if (children.size !== 1) return ref || this.ref;
+            const child = this.state.tree.transforms.get(children.first());
+            if (child.transformer.definition.isDecorator) return this.getApplyRoot(child.ref);
+            return ref || this.ref;
+        }
+
         /**
          * Apply the transformed to the parent node
          * If no params are specified (params <- undefined), default params are lazily resolved.
          */
-        apply<T extends StateTransformer<A, any, any>>(tr: T, params?: Partial<StateTransformer.Params<T>>, options?: Partial<StateTransform.Options>): To<StateTransformer.To<T>, T> {const t = tr.apply(this.ref, params, options);
+        apply<T extends StateTransformer<A, any, any>>(tr: T, params?: Partial<StateTransformer.Params<T>>, options?: Partial<StateTransform.Options>): To<StateTransformer.To<T>, T> {
+            if (tr.definition.isDecorator) {
+                return this.insert(tr, params, options);
+            }
+
+            const applyRoot = this.getApplyRoot();
+            const t = tr.apply(applyRoot, params, options);
+
             this.state.tree.add(t);
             this.editInfo.count++;
             this.editInfo.lastUpdate = t.ref;
@@ -141,9 +156,16 @@ namespace StateBuilder {
         /**
          * Apply the transformed to the parent node
          * If no params are specified (params <- undefined), default params are lazily resolved.
+         * The transformer cannot be a decorator to be able to use this.
          */
         applyOrUpdateTagged<T extends StateTransformer<A, any, any>>(tags: string | string[], tr: T, params?: Partial<StateTransformer.Params<T>>, options?: Partial<StateTransform.Options>): To<StateTransformer.To<T>, T> {
-            const children = this.state.tree.children.get(this.ref).values();
+            if (tr.definition.isDecorator) {
+                throw new Error(`Can't use applyOrUpdateTagged on decorator transformers.`);
+            }
+
+            const applyRoot = this.getApplyRoot();
+
+            const children = this.state.tree.children.get(applyRoot).values();
             while (true) {
                 const child = children.next();
                 if (child.done) break;
@@ -155,7 +177,7 @@ namespace StateBuilder {
                 }
             }
 
-            const t = tr.apply(this.ref, params, { ...options, tags: tagsUnion(tags, options && options.tags) });
+            const t = tr.apply(applyRoot, params, { ...options, tags: tagsUnion(tags, options && options.tags) });
             this.state.tree.add(t);
             this.editInfo.count++;
             this.editInfo.lastUpdate = t.ref;

src/mol-state/transform.ts

@@ -14,7 +14,6 @@ interface Transform<T extends StateTransformer = StateTransformer> {
     readonly transformer: T,
     readonly state: Transform.State,
     readonly tags?: string[],
-    readonly isDecorator?: boolean,
     readonly ref: Transform.Ref,
     /**
      * Sibling-like dependency
@@ -91,7 +90,6 @@ namespace Transform {
     export interface Options {
         ref?: string,
         tags?: string | string[],
-        isDecorator?: boolean,
         state?: State,
         dependsOn?: Ref[]
     }
@@ -109,7 +107,6 @@ namespace Transform {
             transformer,
             state: options?.state || { },
             tags,
-            isDecorator: options?.isDecorator,
             ref,
             dependsOn: options && options.dependsOn,
             params,
@@ -188,7 +185,6 @@ namespace Transform {
             params: t.params ? pToJson(t.params) : void 0,
             state,
             tags: t.tags,
-            isDecorator: t.isDecorator || void 0,
             ref: t.ref,
             dependsOn: t.dependsOn,
             version: t.version
@@ -206,7 +202,6 @@ namespace Transform {
             params: t.params ? pFromJson(t.params) : void 0,
             state: t.state || { },
             tags: t.tags,
-            isDecorator: t.isDecorator,
             ref: t.ref as Ref,
             dependsOn: t.dependsOn,
             version: t.version

src/mol-state/transformer.ts

@@ -106,7 +106,16 @@ namespace Transformer {
         readonly from: StateObject.Ctor[],
         readonly to: StateObject.Ctor[],
         readonly display: { readonly name: string, readonly description?: string },
-        params?(a: A | undefined, globalCtx: unknown): { [K in keyof P]: PD.Any }
+        params?(a: A | undefined, globalCtx: unknown): { [K in keyof P]: PD.Any },
+
+        /**
+         * Decorators are special Transformers mapping the object to the same type.
+         *
+         * Special rules apply:
+         * - applying decorator always "inserts" it instead
+         * - applying to a decorated Transform is applied to the decorator instead (transitive)
+         */
+        readonly isDecorator?: boolean
     }
 
     const registry = new Map<Id, Transformer<any, any>>();
@@ -175,7 +184,8 @@ namespace Transformer {
             to: B | B[],
             /** The source StateObject can be undefined: used for generating docs. */
             params?: PD.For<P> | ((a: StateObject.From<A> | undefined, globalCtx: any) => PD.For<P>),
-            display?: string | { name: string, description?: string }
+            display?: string | { name: string, description?: string },
+            isDecorator?: boolean
         }
 
         export interface Root {
@@ -201,6 +211,7 @@ namespace Transformer {
                     : !!info.params
                         ? info.params as any
                         : void 0,
+                isDecorator: info.isDecorator,
                 ...def
             });
         }

src/mol-state/tree/spine.ts

@@ -56,7 +56,7 @@ namespace StateTreeSpine {
         const cells = state.cells;
         let current = cells.get(currentRef)!;
         const ret: StateObjectCell[] = [current];
-        while (current?.transform.isDecorator) {
+        while (current?.transform.transformer.definition.isDecorator) {
             current = cells.get(current.transform.parent)!;
             ret.push(current);
         }