Support MaskExtension with HeatmapLayer

modules/aggregation-layers/src/aggregation-layer.ts

@@ -53,7 +53,7 @@ export default abstract class AggregationLayer<
   updateState(opts: UpdateParameters<this>) {
     super.updateState(opts);
     const {changeFlags} = opts;
-    if (changeFlags.extensionsChanged) {
+    if (changeFlags.extensionsChanged || changeFlags.propsChanged) {
       const shaders = this.getShaders({});
       if (shaders && shaders.defines) {
         shaders.defines.NON_INSTANCED_MODEL = 1;
@@ -86,6 +86,16 @@ export default abstract class AggregationLayer<
       pickingActive: 0,
       devicePixelRatio: cssToDeviceRatio(gl)
     });
+
+    // @ts-ignore
+    const effects = this.context.deck?.effectManager?.getEffects();
+    // HACK needed to get maskChannels & maskMap
+    if (effects) {
+      for (const effect of effects) {
+        Object.assign(moduleSettings, effect.getModuleParameters?.(this));
+      }
+    }
+
     return moduleSettings;
   }
 

modules/aggregation-layers/src/heatmap-layer/heatmap-layer.ts

@@ -196,6 +196,7 @@ export default class HeatmapLayer<DataT = any, ExtraPropsT = {}> extends Aggrega
   state!: AggregationLayer['state'] & {
     supported: boolean;
     colorDomain?: number[];
+    maskRenderCount: number;
     isWeightMapDirty?: boolean;
     weightsTexture?: Texture2D;
     zoom?: number;
@@ -234,6 +235,15 @@ export default class HeatmapLayer<DataT = any, ExtraPropsT = {}> extends Aggrega
     this._updateHeatmapState(opts);
   }
 
+  postEffect() {
+    const {maskRenderCount} = this.getModuleSettings();
+    const maskChanged = maskRenderCount !== this.state.maskRenderCount;
+    if (maskChanged) {
+      this.setState({maskRenderCount});
+      this._debouncedUpdateWeightmap();
+    }
+  }
+
   _updateHeatmapState(opts: UpdateParameters<this>) {
     const {props, oldProps} = opts;
     const changeFlags = this._getChangeFlags(opts);
@@ -596,12 +606,34 @@ export default class HeatmapLayer<DataT = any, ExtraPropsT = {}> extends Aggrega
       this.state.colorDomain = colorDomain || DEFAULT_COLOR_DOMAIN;
     }
 
-    const uniforms = {
+    const uniforms: any = {
       radiusPixels,
       commonBounds,
       textureWidth: textureSize,
       weightsScale
     };
+
+    // Process extensions for transform
+    const moduleSettings = this.getModuleSettings();
+    const parameters = {
+      blend: true,
+      depthTest: false,
+      blendFunc: [GL.ONE, GL.ONE],
+      blendEquation: GL.FUNC_ADD
+    };
+    const opts = {moduleParameters: moduleSettings, uniforms, parameters, context: this.context};
+    for (const extension of this.props.extensions) {
+      extension.draw.call(this, opts, extension);
+    }
+
+    // Only instance-based masking affects the weights transform
+    // TODO: do we want to do this? On one hand it makes sense,
+    // but can result in starnge colors as the maxWeightValue is
+    // calculated on the unmasked dataset
+    if (uniforms.mask_enabled) {
+      uniforms.mask_enabled = uniforms.mask_maskByInstance;
+    }
+
     // Attribute manager sets data array count as instaceCount on model
     // we need to set that as elementCount on 'weightsTransform'
     weightsTransform.update({
@@ -611,15 +643,10 @@ export default class HeatmapLayer<DataT = any, ExtraPropsT = {}> extends Aggrega
     withParameters(this.context.gl, {clearColor: [0, 0, 0, 0]}, () => {
       weightsTransform.run({
         uniforms,
-        parameters: {
-          blend: true,
-          depthTest: false,
-          blendFunc: [GL.ONE, GL.ONE],
-          blendEquation: GL.FUNC_ADD
-        },
+        parameters,
         clearRenderTarget: true,
         attributes: this.getAttributes(),
-        moduleSettings: this.getModuleSettings()
+        moduleSettings
       });
     });
     this._updateMaxWeightValue();

modules/aggregation-layers/src/heatmap-layer/triangle-layer-vertex.glsl.ts

@@ -38,6 +38,7 @@ varying float vIntensityMax;
 
 void main(void) {
   gl_Position = project_position_to_clipspace(positions, vec3(0.0), vec3(0.0));
+  geometry.position = project_position(vec4(positions, 1.0));
   vTexCoords = texCoords;
   vec4 maxTexture = texture2D(maxTexture, vec2(0.5));
   float maxValue = aggregationMode < 0.5 ? maxTexture.r : maxTexture.g;

modules/aggregation-layers/src/heatmap-layer/triangle-layer.ts

@@ -19,7 +19,7 @@
 // THE SOFTWARE.
 
 import GL from '@luma.gl/constants';
-import {Model, Geometry, Texture2D} from '@luma.gl/core';
+import {Model, Geometry, Texture2D, UpdateParameters} from '@luma.gl/core';
 import {Layer, LayerContext, project32} from '@deck.gl/core';
 import vs from './triangle-layer-vertex.glsl';
 import fs from './triangle-layer-fragment.glsl';
@@ -39,7 +39,7 @@ export default class TriangleLayer extends Layer<_TriangleLayerProps> {
   static layerName = 'TriangleLayer';
 
   getShaders() {
-    return {vs, fs, modules: [project32]};
+    return super.getShaders({vs, fs, modules: [project32]});
   }
 
   initializeState({gl}: LayerContext): void {
@@ -53,6 +53,16 @@ export default class TriangleLayer extends Layer<_TriangleLayerProps> {
     });
   }
 
+  updateState(opts: UpdateParameters<this>) {
+    super.updateState(opts);
+    const {changeFlags} = opts;
+    if (changeFlags.extensionsChanged || changeFlags.propsChanged) {
+      this.setState({
+        model: this._getModel(opts.context.gl)
+      });
+    }
+  }
+
   _getModel(gl: WebGLRenderingContext): Model {
     const {vertexCount} = this.props;
 
@@ -72,6 +82,11 @@ export default class TriangleLayer extends Layer<_TriangleLayerProps> {
     const {texture, maxTexture, colorTexture, intensity, threshold, aggregationMode, colorDomain} =
       this.props;
 
+    if (uniforms.mask_enabled) {
+      // Instance based masking happens in weightsTransform
+      uniforms.mask_enabled = !uniforms.mask_maskByInstance;
+    }
+
     model
       .setUniforms({
         ...uniforms,

modules/aggregation-layers/src/heatmap-layer/weights-vs.glsl.ts

@@ -14,10 +14,11 @@ void main()
   float radiusTexels  = project_pixel_size(radiusPixels) * textureWidth / (commonBounds.z - commonBounds.x);
   gl_PointSize = radiusTexels * 2.;
 
-  vec3 commonPosition = project_position(positions, positions64Low);
+  geometry.position = project_position(vec4(positions, 1.0), positions64Low);
+  geometry.worldPosition = positions;
 
   // map xy from commonBounds to [-1, 1]
-  gl_Position.xy = (commonPosition.xy - commonBounds.xy) / (commonBounds.zw - commonBounds.xy) ;
+  gl_Position.xy = (geometry.position.xy - commonBounds.xy) / (commonBounds.zw - commonBounds.xy) ;
   gl_Position.xy = (gl_Position.xy * 2.) - (1.);
 }
 `;

modules/core/src/passes/layers-pass.ts

@@ -209,6 +209,12 @@ export default class LayersPass extends Pass {
       if (shouldDrawLayer && layer.props.pickable) {
         renderStatus.pickableCount++;
       }
+
+      // Hack
+      if (layer.postEffect) {
+        layer.postEffect();
+      }
+
       if (layer.isComposite) {
         renderStatus.compositeCount++;
       } else if (shouldDrawLayer) {

modules/extensions/src/mask/mask-effect.ts

@@ -41,6 +41,7 @@ export default class MaskEffect implements Effect {
   private dummyMaskMap?: Texture2D;
   private channels: (Channel | null)[] = [];
   private masks: Record<string, Mask> | null = null;
+  private maskRenderCount: number = 0;
   private maskPass?: MaskPass;
   private maskMap?: Texture2D;
   private lastViewport?: Viewport;
@@ -87,21 +88,22 @@ export default class MaskEffect implements Effect {
 
     // // Debug show FBO contents on screen
     // const color = readPixelsToArray(this.maskMap);
-    // let canvas = document.getElementById('fbo-canvas');
+    // let canvas = document.getElementById('fbo-canvas') as HTMLCanvasElement;
     // if (!canvas) {
     //   canvas = document.createElement('canvas');
     //   canvas.id = 'fbo-canvas';
     //   canvas.width = this.maskMap.width;
     //   canvas.height = this.maskMap.height;
-    //   canvas.style.zIndex = 100;
+    //   canvas.style.zIndex = '100';
     //   canvas.style.position = 'absolute';
-    //   canvas.style.right = 0;
+    //   canvas.style.right = '0';
+    //   canvas.style.top = '0';
     //   canvas.style.border = 'blue 1px solid';
     //   canvas.style.width = '256px';
     //   canvas.style.transform = 'scaleY(-1)';
     //   document.body.appendChild(canvas);
     // }
-    // const ctx = canvas.getContext('2d');
+    // const ctx = canvas.getContext('2d')!;
     // const imageData = ctx.createImageData(this.maskMap.width, this.maskMap.height);
     // for (let i = 0; i < color.length; i += 4) {
     //   imageData.data[i + 0] = color[i + 0];
@@ -177,6 +179,7 @@ export default class MaskEffect implements Effect {
             devicePixelRatio: 1
           }
         });
+        this.maskRenderCount++;
       }
     }
 
@@ -242,10 +245,12 @@ export default class MaskEffect implements Effect {
 
   getModuleParameters(): {
     maskMap: Texture2D;
+    maskRenderCount: number;
     maskChannels: Record<string, Mask> | null;
   } {
     return {
       maskMap: this.masks ? this.maskMap : this.dummyMaskMap,
+      maskRenderCount: this.maskRenderCount,
       maskChannels: this.masks
     };
   }

modules/extensions/src/mask/shader-module.ts

@@ -61,7 +61,6 @@ varying vec2 mask_texCoords;
 
     // Debug: show extent of render target
     // gl_FragColor = vec4(mask_texCoords, 0.0, 1.0);
-    gl_FragColor = texture2D(mask_texture, mask_texCoords);
 
     if (!mask) discard;
   }