Better template atomset builder

src/mol-model/structure/_spec/atom-set.spec.ts

@@ -49,8 +49,8 @@ describe('atom set', () => {
     it('template', () => {
         const template = AtomSet.ofAtoms([p(1, 3), p(0, 1), p(0, 6), p(0, 2)], AtomSet.Empty)
         const gen = AtomSet.TemplateGenerator(template);
-        gen.add(0, AtomGroup.createNew(OrderedSet.ofSortedArray([1, 2, 6])));
-        gen.add(1, AtomGroup.createNew(OrderedSet.ofSingleton(3)));
+        gen.add(0, OrderedSet.ofSortedArray([1, 2, 6]));
+        gen.add(1, OrderedSet.ofSingleton(3));
         const set = gen.getSet();
 
         expect(AtomSet.unitGetById(set, 0)).toBe(AtomSet.unitGetById(template, 0));
@@ -61,8 +61,8 @@ describe('atom set', () => {
     it('template 1', () => {
         const template = AtomSet.ofAtoms([p(1, 3), p(0, 1), p(0, 6), p(0, 2)], AtomSet.Empty)
         const gen = AtomSet.TemplateGenerator(template);
-        gen.add(0, AtomGroup.createNew(OrderedSet.ofSortedArray([1, 2, 6])));
-        gen.add(1, AtomGroup.createNew(OrderedSet.ofSingleton(4)));
+        gen.add(0, OrderedSet.ofSortedArray([1, 2, 6]));
+        gen.add(1, OrderedSet.ofSingleton(4));
         const set = gen.getSet();
 
         expect(AtomSet.unitGetById(set, 0)).toBe(AtomSet.unitGetById(template, 0));

src/mol-model/structure/structure/atom/impl/builder.ts → src/mol-model/structure/structure/atom/impl/set-builder.ts

@@ -6,7 +6,6 @@
 
 import AtomSet from '../set'
 import Atom from '../../atom'
-import AtomGroup from '../group'
 import { OrderedSet, IntMap } from 'mol-data/int'
 import { sortArray } from 'mol-data/util/sort'
 
@@ -36,22 +35,17 @@ export class Builder {
     }
 
     getSet(): AtomSet {
-        const generator = AtomSet.Generator();
-
-        let allEqual = this.keys.length === AtomSet.unitCount(this.parent);
+        const generator = AtomSet.TemplateGenerator(this.parent);
 
         for (let i = 0, _i = this.keys.length; i < _i; i++) {
             const k = this.keys[i];
             const unit = this.units.get(k);
             const l = unit.length;
             if (!this.sorted && l > 1) sortArray(unit);
-
-            const parentGroup = AtomSet.unitGetById(this.parent, k);
-            const group = AtomGroup.createChild(parentGroup, OrderedSet.ofSortedArray(unit));
-            if (group !== parentGroup) allEqual = false;
-            generator.add(k, group);
+            generator.add(k, OrderedSet.ofSortedArray(unit));
         }
-        return allEqual ? this.parent : generator.getSet();
+
+        return generator.getSet();
     }
 
     singleton(): Atom {

src/mol-model/structure/structure/atom/impl/set.ts

@@ -215,15 +215,19 @@ export class TemplateAtomSetGenerator {
     private templateGroups: IntMap<AtomGroup>;
     private equalGroups = 0;
 
-    add(unit: number, group: AtomGroup) {
-        if (AtomGroup.size(group) === 0) return;
+    add(unit: number, set: OS) {
+        if (OS.size(set) === 0) return;
         this.keys[this.keys.length] = unit;
-        let t: AtomGroup;
-        if (this.templateGroups.has(unit) && AtomGroup.areEqual(t = this.templateGroups.get(unit), group)) {
-            this.groups.set(unit, t);
-            this.equalGroups++;
+        if (this.templateGroups.has(unit)) {
+            const t = this.templateGroups.get(unit);
+            if (OS.areEqual(t.atoms, set)) {
+                this.groups.set(unit, t);
+                this.equalGroups++;
+            } else {
+                this.groups.set(unit, AtomGroup.createNew(set));
+            }
         } else {
-            this.groups.set(unit, group);
+            this.groups.set(unit, AtomGroup.createNew(set));
         }
     }
 
@@ -348,8 +352,7 @@ function ofAtomsImpl(xs: ArrayLike<Atom>, template: AtomSetImpl) {
     const generator = TemplateGenerator(template);
     for (let i = 0, _i = keys.length; i < _i; i++) {
         const k = keys[i];
-        const group = AtomGroup.createNew(OS.ofSortedArray(normalizeArray(elements.get(k))));
-        generator.add(k, group);
+        generator.add(k, OS.ofSortedArray(normalizeArray(elements.get(k))));
     }
 
     return generator.getSet();

src/mol-model/structure/structure/atom/set.ts

@@ -4,11 +4,11 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import { SortedArray, Iterator } from 'mol-data/int'
+import { SortedArray, Iterator, OrderedSet } from 'mol-data/int'
 import Atom from '../atom'
 import AtomGroup from './group'
 import * as Impl from './impl/set'
-import * as Builders from './impl/builder'
+import * as Builders from './impl/set-builder'
 
 /** A map-like representation of grouped atom set */
 namespace AtomSet {
@@ -44,7 +44,9 @@ namespace AtomSet {
 
     export interface Generator { add(unit: number, set: AtomGroup): void, getSet(): AtomSet }
     export const Generator: () => Generator = Impl.Generator as any
-    export const TemplateGenerator: (template: AtomSet) => Generator = Impl.TemplateGenerator as any
+
+    export interface TemplateGenerator { add(unit: number, set: OrderedSet): void, getSet(): AtomSet }
+    export const TemplateGenerator: (template: AtomSet) => TemplateGenerator = Impl.TemplateGenerator as any
 
     // TODO: bounding sphere
     // TODO: distance, areWithIn?