Refactor Structure to use IntMap for units

src/mol-data/int/map.ts

@@ -28,16 +28,6 @@ namespace IntMap {
         return iterableToArray(map.keys());
     }
 
-    export function ofObj<T>(obj: { [key: number]: T }): IntMap<T> {
-        const keys = Object.keys(obj);
-        const ret = new Map<number, T>();
-        for (let i = 0, _i = keys.length; i < _i; i++) {
-            const k = keys[i];
-            ret.set(+k, obj[k as any]);
-        }
-        return ret as IntMap<T>;
-    }
-
     export function Mutable<T>(): Mutable<T> {
         return new Map<number, T>() as Mutable<T>;
     }
@@ -45,6 +35,27 @@ namespace IntMap {
     export function asImmutable<T>(map: IntMap<T>): IntMap<T> {
         return map;
     }
+
+    export function copy<T>(map: IntMap<T>): Mutable<T> {
+        const ret = Mutable<T>();
+        const it = map.keys();
+        while (true) {
+            const { done, value } = it.next();
+            if (done) break;
+            ret.set(value, map.get(value));
+        }
+        return ret;
+    }
+
+    export function addFrom<T>(map: Mutable<T>, src: IntMap<T>) {
+        const it = src.keys();
+        while (true) {
+            const { done, value } = it.next();
+            if (done) break;
+            map.set(value, src.get(value));
+        }
+        return map;
+    }
 }
 
 export default IntMap;

src/mol-model/structure/query/generators.ts

@@ -45,7 +45,7 @@ function atomGroupsLinear(atomTest: Atom.Predicate): Query {
 
         for (let i = 0, _i = unitIds.length; i < _i; i++) {
             const unitId = unitIds[i];
-            l.unit = units[unitId];
+            l.unit = units.get(unitId);
             const set = AtomSet.unitGetByIndex(atoms, i);
 
             builder.beginUnit();
@@ -69,7 +69,7 @@ function atomGroupsSegmented({ entityTest, chainTest, residueTest, atomTest }: A
 
         for (let i = 0, _i = unitIds.length; i < _i; i++) {
             const unitId = unitIds[i];
-            const unit = units[unitId];
+            const unit = units.get(unitId);
             l.unit = unit;
             const set = AtomSet.unitGetByIndex(atoms, i);
 
@@ -160,7 +160,7 @@ function atomGroupsGrouped({ entityTest, chainTest, residueTest, atomTest, group
 
         for (let i = 0, _i = unitIds.length; i < _i; i++) {
             const unitId = unitIds[i];
-            const unit = units[unitId];
+            const unit = units.get(unitId);
             l.unit = unit;
             const set = AtomSet.unitGetByIndex(atoms, i);
 

src/mol-model/structure/query/query.ts

@@ -7,5 +7,7 @@
 import { Structure } from '../structure'
 import Selection from './selection'
 
+// TODO: Query { (s: Structure): Computation<Selection> }
+
 interface Query { (s: Structure): Selection }
 export default Query

src/mol-model/structure/query/selection.ts

@@ -6,6 +6,7 @@
 
 import Iterator from 'mol-data/iterator'
 import { HashSet } from 'mol-data/util'
+import { IntMap } from 'mol-data/int'
 import { Structure, Atom, AtomSet } from '../structure'
 
 type Selection =
@@ -105,20 +106,18 @@ export default Selection
 
 function unionUnits(xs: Structure[]): Structure['units'] {
     let prev = xs[0].units;
-    let sameModel = true;
+    let sameUnits = true;
     for (let i = 1, _i = xs.length; i < _i; i++) {
-        if (xs[i].units !== prev) sameModel = false;
+        if (xs[i].units !== prev) sameUnits = false;
     }
-    if (sameModel) return prev;
+    if (sameUnits) return prev;
 
-    let ret: any = { ...prev };
+    const ret = IntMap.copy(prev);
     for (let i = 1, _i = xs.length; i < _i; i++) {
         const units = xs[i].units;
-        if (units !== prev) {
-            const keys = Object.keys(units);
-            for (let j = 0; j < keys.length; j++) ret[keys[j]] = (units as any)[keys[j]];
-        }
-        prev = xs[i];
+        if (units !== prev) IntMap.addFrom(ret, units);
+        prev = units;
     }
+
     return ret;
 }

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

@@ -29,7 +29,7 @@ namespace Atom {
     export interface Predicate extends Property<boolean> { }
 
     export function updateLocation(structure: Structure, l: Location, atom: Atom) {
-        l.unit = structure.units[unit(atom)];
+        l.unit = structure.units.get(unit(atom));
         l.atom = index(atom);
         return l;
     }

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

@@ -4,7 +4,7 @@
  * @author David Sehnal <david.sehnal@gmail.com>
  */
 
-import { OrderedSet, Iterator } from 'mol-data/int'
+import { OrderedSet, Iterator, IntMap } from 'mol-data/int'
 import { UniqueArray } from 'mol-data/util'
 import SymmetryOperator from 'mol-math/geometry/symmetry-operator'
 import { Model, Format } from '../model'
@@ -14,14 +14,14 @@ import Atom from './atom'
 
 
 interface Structure extends Readonly<{
-    units: { readonly [id: number]: Unit },
+    units: IntMap<Unit>,
     atoms: AtomSet
 }> { }
 
 namespace Structure {
-    export const Empty = { units: {}, atoms: AtomSet.Empty };
+    export const Empty: Structure = { units: IntMap.Empty, atoms: AtomSet.Empty };
 
-    export function create(units: Structure['units'], atoms: AtomSet): Structure {
+    export function create(units: IntMap<Unit>, atoms: AtomSet): Structure {
         return { units, atoms };
     }
 
@@ -36,8 +36,7 @@ namespace Structure {
 
         for (let c = 0; c < chains.count; c++) {
             const unit = Unit.create(model, SymmetryOperator.Default);
-            builder.addUnit(unit);
-            builder.setAtoms(unit.id, OrderedSet.ofBounds(chains.segments[c], chains.segments[c + 1]));
+            builder.add(unit, OrderedSet.ofBounds(chains.segments[c], chains.segments[c + 1]));
         }
 
         return builder.getStructure();
@@ -52,12 +51,12 @@ namespace Structure {
     }
 
     class BuilderImpl implements Builder {
-        private units = Object.create(null);
+        private units = IntMap.Mutable<Unit>();
         private atoms = AtomSet.Generator();
         atomCount = 0;
 
         add(unit: Unit, atoms: OrderedSet) { this.addUnit(unit); this.setAtoms(unit.id, atoms); }
-        addUnit(unit: Unit) { this.units[unit.id] = unit; }
+        addUnit(unit: Unit) { this.units.set(unit.id, unit); }
         setAtoms(unitId: number, atoms: OrderedSet) { this.atoms.add(unitId, atoms); this.atomCount += OrderedSet.size(atoms); }
         getStructure(): Structure { return this.atomCount > 0 ? Structure.create(this.units, this.atoms.getSet()) : Empty; }
     }
@@ -72,9 +71,11 @@ namespace Structure {
     }
 
     export function getModels(s: Structure) {
+        const { units, atoms } = s;
         const arr = UniqueArray.create<Model['id'], Model>();
-        for (const k of Object.keys(s.units)) {
-            const u = s.units[+k];
+        const ids = AtomSet.unitIds(atoms);
+        for (let i = 0; i < ids.length; i++) {
+            const u = units.get(ids[i]);
             UniqueArray.add(arr, u.model.id, u.model);
         }
         return arr.array;

src/mol-model/structure/structure/symmetry.ts

@@ -34,7 +34,7 @@ function buildAssemblyImpl(structure: Structure, name: string) {
 
         for (const oper of g.operators) {
             for (let uI = 0, _uI = unitIds.length; uI < _uI; uI++) {
-                const unit = units[unitIds[uI]];
+                const unit = units.get(unitIds[uI]);
                 assembler.add(Unit.create(unit.model, oper), AtomSet.unitGetByIndex(atoms, uI));
             }
         }

src/perf-tests/structure.ts

@@ -74,7 +74,7 @@ export namespace PropertyAccess {
         let s = 0;
 
         for (let i = 0, _i = unitIds.length; i < _i; i++) {
-            l.unit = units[unitIds[i]];
+            l.unit = units.get(unitIds[i]);
             const set = AtomSet.unitGetByIndex(atoms, i);
 
 
@@ -96,7 +96,7 @@ export namespace PropertyAccess {
 
         let vA = 0, cC = 0, rC = 0;
         for (let i = 0, _i = unitIds.length; i < _i; i++) {
-            const unit = units[unitIds[i]];
+            const unit = units.get(unitIds[i]);
             l.unit = unit;
             const set = AtomSet.unitGetByIndex(atoms, i);
 
@@ -134,7 +134,7 @@ export namespace PropertyAccess {
         let s = 0;
 
         for (let i = 0, _i = unitIds.length; i < _i; i++) {
-            const unit = units[unitIds[i]];
+            const unit = units.get(unitIds[i]);
             l.unit = unit;
             const set = AtomSet.unitGetByIndex(atoms, i);
             const residuesIt = Segmentation.transientSegments(unit.hierarchy.residueSegments, set);
@@ -155,7 +155,7 @@ export namespace PropertyAccess {
         const l = Atom.Location();
         while (atomsIt.hasNext) {
             const a = atomsIt.move();
-            l.unit = units[Atom.unit(a)];
+            l.unit = units.get(Atom.unit(a));
             l.atom = Atom.index(a);
             s += p(l);
         }
@@ -200,7 +200,7 @@ export namespace PropertyAccess {
 
         for (let i = 0, _i = unitIds.length; i < _i; i++) {
             const unitId = unitIds[i];
-            const unit = units[unitId];
+            const unit = units.get(unitId);
             const set = AtomSet.unitGetByIndex(atoms, i);
             //const { residueIndex, chainIndex } = unit;
             const p = unit.conformation.atomId.value;