Use a special data structure for transferring feature id maps to the main thread

src/data/feature_position_map.js

@@ -0,0 +1,110 @@
+// @flow
+
+import { register } from '../util/web_worker_transfer';
+import assert from 'assert';
+
+type SerializedFeaturePositionMap = {
+    ids: Float64Array;
+    positions: Uint32Array;
+};
+
+type FeaturePosition = {
+    index: number;
+    start: number;
+    end: number;
+};
+
+// A transferable data structure that maps feature ids to their indices and buffer offsets
+export default class FeaturePositionMap {
+    ids: Array<number>;
+    positions: Array<number>;
+    indexed: boolean;
+
+    constructor() {
+        this.ids = [];
+        this.positions = [];
+        this.indexed = false;
+    }
+
+    add(id: number, index: number, start: number, end: number) {
+        this.ids.push(id);
+        this.positions.push(index, start, end);
+    }
+
+    getPositions(id: number): Array<FeaturePosition> {
+        assert(this.indexed);
+
+        // binary search for the first occurrence of id in this.ids;
+        // relies on ids/positions being sorted by id, which happens in serialization
+        let i = 0;
+        let j = this.ids.length - 1;
+        while (i < j) {
+            const m = (i + j) >> 1;
+            if (this.ids[m] >= id) {
+                j = m;
+            } else {
+                i = m + 1;
+            }
+        }
+        const positions = [];
+        while (this.ids[i] === id) {
+            const index = this.positions[3 * i];
+            const start = this.positions[3 * i + 1];
+            const end = this.positions[3 * i + 2];
+            positions.push({index, start, end});
+            i++;
+        }
+        return positions;
+    }
+
+    static serialize(map: FeaturePositionMap, transferables: Array<ArrayBuffer>): SerializedFeaturePositionMap {
+        const ids = new Float64Array(map.ids);
+        const positions = new Uint32Array(map.positions);
+
+        sort(ids, positions, 0, ids.length - 1);
+
+        transferables.push(ids.buffer, positions.buffer);
+
+        return {ids, positions};
+    }
+
+    static deserialize(obj: SerializedFeaturePositionMap): FeaturePositionMap {
+        const map = new FeaturePositionMap();
+        // after transferring, we only use these arrays statically (no pushes),
+        // so TypedArray vs Array distinction that flow points out doesn't matter
+        map.ids = (obj.ids: any);
+        map.positions = (obj.positions: any);
+        map.indexed = true;
+        return map;
+    }
+}
+
+// custom quicksort that sorts ids, indices and offsets together (by ids)
+function sort(ids, positions, left, right) {
+    if (left >= right) return;
+
+    const pivot = ids[(left + right) >> 1];
+    let i = left - 1;
+    let j = right + 1;
+
+    while (true) {
+        do i++; while (ids[i] < pivot);
+        do j--; while (ids[j] > pivot);
+        if (i >= j) break;
+        swap(ids, i, j);
+        swap(positions, 3 * i, 3 * j);
+        swap(positions, 3 * i + 1, 3 * j + 1);
+        swap(positions, 3 * i + 2, 3 * j + 2);
+    }
+
+    sort(ids, positions, left, j);
+    sort(ids, positions, j + 1, right);
+}
+
+function swap(arr, i, j) {
+    const tmp = arr[i];
+    arr[i] = arr[j];
+    arr[j] = tmp;
+}
+
+register('FeaturePositionMap', FeaturePositionMap);

src/data/program_configuration.js

@@ -7,6 +7,7 @@ import { register, serialize, deserialize } from '../util/web_worker_transfer';
 import { PossiblyEvaluatedPropertyValue } from '../style/properties';
 import { StructArrayLayout1f4, StructArrayLayout2f8, StructArrayLayout4f16 } from './array_types';
 import EvaluationParameters from '../style/evaluation_parameters';
+import FeaturePositionMap from './feature_position_map';
 import {
     Uniform,
     Uniform1f,
@@ -29,14 +30,6 @@ import type {
 import type {PossiblyEvaluated} from '../style/properties';
 import type {FeatureStates} from '../source/source_state';
 
-type FeaturePaintBufferMap = {
-    [feature_id: string]: Array<{
-        index: number,
-        start: number,
-        end: number
-    }>
-};
-
 function packColor(color: Color): [number, number] {
     return [
         packUint8ToFloat(255 * color.r, 255 * color.g),
@@ -366,16 +359,15 @@ export default class ProgramConfiguration {
     layoutAttributes: Array<StructArrayMember>;
 
     _buffers: Array<VertexBuffer>;
-
-    _idMap: FeaturePaintBufferMap;
+    _featureMap: FeaturePositionMap;
     _bufferOffset: number;
 
     constructor() {
         this.binders = {};
         this.cacheKey = '';
 
         this._buffers = [];
-        this._idMap = {};
+        this._featureMap = new FeaturePositionMap();
         this._bufferOffset = 0;
     }
 
@@ -414,27 +406,18 @@ export default class ProgramConfiguration {
         for (const property in this.binders) {
             this.binders[property].populatePaintArray(newLength, feature);
         }
-        if (feature.id) {
-            const featureId = String(feature.id);
-            this._idMap[featureId] = this._idMap[featureId] || [];
-            this._idMap[featureId].push({
-                index: index,
-                start: this._bufferOffset,
-                end: newLength
-            });
+        if (feature.id !== undefined) {
+            this._featureMap.add(+feature.id, index, this._bufferOffset, newLength);
         }
-
         this._bufferOffset = newLength;
     }
 
     updatePaintArrays(featureStates: FeatureStates, vtLayer: VectorTileLayer, layer: TypedStyleLayer): boolean {
         let dirty: boolean = false;
         for (const id in featureStates) {
-            const posArray = this._idMap[id];
-            if (!posArray) continue;
+            const positions = this._featureMap.getPositions(+id);
 
-            const featureState = featureStates[id];
-            for (const pos of posArray) {
+            for (const pos of positions) {
                 const feature = vtLayer.feature(pos.index);
 
                 for (const property in this.binders) {
@@ -444,7 +427,7 @@ export default class ProgramConfiguration {
                         //AHM: Remove after https://github.com/mapbox/mapbox-gl-js/issues/6255
                         const value = layer.paint.get(property);
                         (binder: any).expression = value.value;
-                        binder.updatePaintArray(pos.start, pos.end, feature, featureState);
+                        binder.updatePaintArray(pos.start, pos.end, feature, featureStates[id]);
                         dirty = true;
                     }
                 }

test/unit/data/feature_position_map.test.js

@@ -0,0 +1,42 @@
+import { test } from 'mapbox-gl-js-test';
+
+import FeatureMap from '../../../src/data/feature_position_map';
+import { serialize, deserialize } from '../../../src/util/web_worker_transfer';
+
+test('FeaturePositionMap', (t) => {
+
+    test('Can be queried after serialization/deserialization', (t) => {
+        const featureMap = new FeatureMap();
+        featureMap.add(7, 1, 0, 1);
+        featureMap.add(3, 2, 1, 2);
+        featureMap.add(7, 3, 2, 3);
+        featureMap.add(4, 4, 3, 4);
+        featureMap.add(2, 5, 4, 5);
+        featureMap.add(7, 6, 5, 7);
+
+        const featureMap2 = deserialize(serialize(featureMap, []));
+
+        const compareIndex = (a, b) => a.index - b.index;
+
+        t.same(featureMap2.getPositions(7).sort(compareIndex), [
+            {index: 1, start: 0, end: 1},
+            {index: 3, start: 2, end: 3},
+            {index: 6, start: 5, end: 7}
+        ].sort(compareIndex));
+
+        t.end();
+    });
+
+    test('Can not be queried before serialization/deserialization', (t) => {
+        const featureMap = new FeatureMap();
+        featureMap.add(0, 1, 2, 3);
+
+        t.throws(() => {
+            featureMap.getPositions(0);
+        });
+
+        t.end();
+    });
+
+    t.end();
+});